見出し画像

LLama-3.2-3BをローカルPCのブラウザ上で試す

はまち

WebLLMとWebGPUという技術を利用してローカルPCのブラウザ上でLlama3.2を実行できるというXポストを見つけました。

詳しい仕組みは全く理解できていないのですが、とにかくさっそく手元のゲーミングノートPC（RTX3060 laptop）で試してみました。


ロードするモデルを選択してダウンロード

Llama-3.2-3Bは日本語は得意じゃなさそうなので英語で試してみます。

プロンプト：トランスフォーマーアーキテクチャーを使った基本的なLLMのモデルをpytorch lightningで実装してみて
ちゃんとローカルのGPUが回ってるっぽい
今回の推論速度

今までで一番手軽にローカル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 parser

HParams.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_steps

main.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)


この記事が気に入ったらサポートをしてみませんか？