LLama-3.2-3BをローカルPCのブラウザ上で試す
Run Llama 3.2 in the browser on WebGPU 🔥
— Caleb (@calebfahlgren) September 26, 2024
• Llama 3.2 (1B + 3B)🤏 🦙
• Running 100% locally in the browser at 62 tok/sec 🏎️
• Powered by MLC WebLLM + WebGPU ⚡ pic.twitter.com/HaxRQTSrTc
WebLLMとWebGPUという技術を利用してローカルPCのブラウザ上でLlama3.2を実行できるというXポストを見つけました。
詳しい仕組みは全く理解できていないのですが、とにかくさっそく手元のゲーミングノートPC（RTX3060 laptop）で試してみました。
Llama-3.2-3Bは日本語は得意じゃなさそうなので英語で試してみます。
今までで一番手軽にローカルPC上でLLMを試せた気がします！すばらしい。
特別な環境を準備しなくてもブラウザ上で完結するので、小規模なモデルをローカルでサクッと試してもらうには良いかのも。
参考：今回出力されたコード（Llama-3.2-3B-Instruct-q4f32_1-MLC）
動作は試していないですが、それっぽい感じのコードが出力されてますね。最近のモデルは3Bとか小さくても立派に動きますね。
LLMModel.py
import torch
from torch import nn
from torch.nn import functional as F
from transformers import get_linear_schedule_with_warmup
from pytorch_lightning import LightningModule
from pytorch_lightning.callbacks import EarlyStopping
class LLMModel(LightningModule):
def __init__(self, hparams):
super(LLMModel, self).__init__()
self.hparams = hparams
self.encoder = nn.Sequential(
nn.Embedding(self.hparams.vocab_size, self.hparams hidden_size),
nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size),
nn.ReLU(),
nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size),
nn.ReLU(),
nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size),
nn.ReLU(),
nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size),
nn.ReLU(),
nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size),
nn.ReLU()
)
self.fc_out = nn.Linear(self.hparams.hidden_size, self.hparams.vocab_size)
def forward(self, x, target):
x = self.encoder(x)
x = x.permute(1, 0, 2)
x = x.contiguous().view(-1, self.hparams.hidden_size)
x = self.decoder(x)
x = x.permute(1, 0, 2)
x = x.contiguous().view(-1, self.hparams.hidden_size)
x = self.fc_out(x)
return x
def training_step(self, batch, batch_idx):
x, target = batch
x = x.to(self.device)
target = target.to(self.device)
output = self.forward(x, target)
loss = F.cross_entropy(output, target)
self.log("train_loss", loss)
return {"loss": loss}
def validation_step(self, batch, batch_idx):
x, target = batch
x = x.to(self.device)
target = target.to(self.device)
output = self.forward(x, target)
loss = F.cross_entropy(output, target)
self.log("val_loss", loss)
return {"loss": loss}
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.hparams.num_training_steps)
return {
"optimizer": optimizer,
"lr_scheduler": scheduler
}
def train(self, train_dataloader, val_dataloader, early_stopping_callback):
trainer = self.trainer
trainer.fit(self, train_dataloader, val_dataloader, early_stopping_callback)
def test(self, test_dataloader):
self.test_dataloader = test_dataloader
results = self.trainer.test(self)
return results
@staticmethod
def add_model_specific_args(parent_parser):
parser = parent_parser.add_argument_group("LLMModel")
parser.add_argument("--hidden_size", type=int, default=256)
parser.add_argument("--vocab_size", type=int, default=10000)
parser.add_argument("--learning_rate", type=float, default=3e-4)
parser.add_argument("--num_training_steps", type=int, default=100000)
parser.add_argument("--warmup_steps", type=int, default=1000)
return parserHParams.py
import torch
class HParams:
def __init__(self, hidden_size, vocab_size, learning_rate, num_training_steps, warmup_steps):
self.hidden_size = hidden_size
self.vocab_size = vocab_size
self.learning_rate = learning_rate
self.num_training_steps = num_training_steps
self.warmup_steps = warmup_stepsmain.py
import torch
from torch.utils.data import Dataset, DataLoader
from LLMModel import LLMModel
from HParams import HParams
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import EarlyStopping
# Define a dataset
class MyDataset(Dataset):
def __init__(self, data, tokenizer):
self.data = data
self.tokenizer = tokenizer
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
text = self.data[idx]
inputs = self.tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=512,
padding='max_length',
truncation=True,
return_attention_mask=True,
return_tensors='pt'
)
return {
'input_ids': inputs['input_ids'].flatten(),
'attention_mask': inputs['attention_mask'].flatten(),
'labels': torch.tensor(0) # Add labels here
}
# Set up the model and data
hparams = HParams(256, 10000, 3e-4, 100000, 1000)
tokenizer = 'bert-base-uncased'
model = LLMModel(hparams)
# Set up the data loaders
train_dataset = MyDataset(train_data, tokenizer)
val_dataset = MyDataset(val_data, tokenizer)
test_dataset = MyDataset(test_data, tokenizer)
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# Set up the trainer
trainer = Trainer(
max_epochs=10,
check_val_every_n_steps=1,
callbacks=[EarlyStopping()]
)
# Train the model
trainer.fit(model, train_dataloader, val_dataloader)