ノーパソでつよつよLLM計画!

2024年11月30日

概要

みなさん、Llamaが貧弱だと感じたことはありますか？ありますよね？(圧)

だったら、自分で強くすればよいのです!以上!

方法

…これで終わるわけにも行かないので、方法を書き散らしておきます。

まずは、Llamaを教師モデルとして自分のモデルを訓練するために、そのプログラムをPerplexityに書いてもらいます。そのプログラムがこちらです。(動作確認していないので、ご注意)

<h1 id="%E5%88%9D%E5%9B%9E%E7%94%A8">初回用</h1>
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import MllamaForConditionalGeneration, AutoProcessor
from logging import getLogger, Formatter, StreamHandler, DEBUG

<h1 id="%E3%83%AD%E3%82%AC%E3%83%BC%E3%81%AE%E8%A8%AD%E5%AE%9A">ロガーの設定</h1>
logger = getLogger(__name__)
logger.setLevel(DEBUG)
handler = StreamHandler()
handler.setLevel(DEBUG)
formatter = Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

<h1 id="WikipediaDataset%E3%81%AE%E5%AE%9A%E7%BE%A9">WikipediaDatasetの定義</h1>
class WikipediaDataset(Dataset):
    def __init__(self, file_path, max_length=512):
        self.file_path = file_path
        self.max_length = max_length
        self.data = self.load_annotations()

    def load_annotations(self):
        with open(self.file_path, 'r', encoding='utf-8') as f:
            return [line.strip() for line in f if line.strip()]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx][:self.max_length]

<h1 id="SmallModel%E3%81%AE%E5%AE%9A%E7%BE%A9">SmallModelの定義</h1>
class SmallModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SmallModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

<h1 id="%E3%83%AD%E3%83%BC%E3%82%AB%E3%83%AB%E3%81%AELlama%203.2%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E3%83%AD%E3%83%BC%E3%83%89">ローカルのLlama 3.2モデルのロード</h1>
model_path = "path/to/your/local/llama3.2/model"
teacher_model = MllamaForConditionalGeneration.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
processor = AutoProcessor.from_pretrained(model_path)

<h1 id="SmallModel%E3%81%AE%E5%88%9D%E6%9C%9F%E5%8C%96">SmallModelの初期化</h1>
input_size = 768  # 入力サイズ（実際のタスクに合わせて調整）
hidden_size = 256
output_size = teacher_model.config.vocab_size
student_model = SmallModel(input_size, hidden_size, output_size).to("cuda")

<h1 id="%E3%83%87%E3%83%BC%E3%82%BF%E3%82%BB%E3%83%83%E3%83%88%E3%81%A8DataLoader%E3%81%AE%E6%BA%96%E5%82%99">データセットとDataLoaderの準備</h1>
dataset = WikipediaDataset("path/to/wiki.txt")
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

<h1 id="%E6%90%8D%E5%A4%B1%E9%96%A2%E6%95%B0%E3%81%A8%E3%82%AA%E3%83%97%E3%83%86%E3%82%A3%E3%83%9E%E3%82%A4%E3%82%B6%E3%81%AE%E8%A8%AD%E5%AE%9A">損失関数とオプティマイザの設定</h1>
criterion = nn.KLDivLoss(reduction='batchmean')
optimizer = optim.Adam(student_model.parameters())

<h1 id="%E5%AD%A6%E7%BF%92%E3%83%AB%E3%83%BC%E3%83%97">学習ループ</h1>
num_epochs = 10
temperature = 2.0

for epoch in range(num_epochs):
    for batch in dataloader:
        # 入力の処理
        inputs = processor(batch, return_tensors="pt", padding=True, truncation=True).to("cuda")
        
        # 教師モデルの出力を取得
        with torch.no_grad():
            teacher_outputs = teacher_model(**inputs).logits
        
        # 生徒モデルの出力を取得
        student_outputs = student_model(inputs.input_ids)
        
        # 知識蒸留損失の計算
        loss = criterion(
            nn.functional.log_softmax(student_outputs / temperature, dim=-1),
            nn.functional.softmax(teacher_outputs / temperature, dim=-1)
        )
        
        # 逆伝播と最適化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    logger.debug(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}")

<h1 id="%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E4%BF%9D%E5%AD%98">モデルの保存</h1>
torch.save(student_model.state_dict(), "small_model.pth")

logger.debug("モデルが保存されました。")                                                                                                                                           )

<h1 id="2%E5%9B%9E%E7%9B%AE%E4%BB%A5%E9%99%8D%E7%94%A8">2回目以降用</h1>
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import MllamaForConditionalGeneration, AutoProcessor
from logging import getLogger, Formatter, StreamHandler, DEBUG

<h1 id="%E3%83%AD%E3%82%AC%E3%83%BC%E3%81%AE%E8%A8%AD%E5%AE%9A">ロガーの設定</h1>
logger = getLogger(__name__)
logger.setLevel(DEBUG)
handler = StreamHandler()
handler.setLevel(DEBUG)
formatter = Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

<h1 id="WikipediaDataset%E3%81%AE%E5%AE%9A%E7%BE%A9">WikipediaDatasetの定義</h1>
class WikipediaDataset(Dataset):
    def __init__(self, file_path, max_length=512):
        self.file_path = file_path
        self.max_length = max_length
        self.data = self.load_annotations()

    def load_annotations(self):
        with open(self.file_path, 'r', encoding='utf-8') as f:
            return [line.strip() for line in f if line.strip()]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx][:self.max_length]

<h1 id="SmallModel%E3%81%AE%E5%AE%9A%E7%BE%A9">SmallModelの定義</h1>
class SmallModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SmallModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

<h1 id="%E3%83%AD%E3%83%BC%E3%82%AB%E3%83%AB%E3%81%AELlama%203.2%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E3%83%AD%E3%83%BC%E3%83%89%EF%BC%88%E6%95%99%E5%B8%AB%E3%83%A2%E3%83%87%E3%83%AB%EF%BC%89">ローカルのLlama 3.2モデルのロード（教師モデル）</h1>
model_path = "path/to/your/local/llama3.2/model"
teacher_model = MllamaForConditionalGeneration.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
processor = AutoProcessor.from_pretrained(model_path)

<h1 id="%E6%97%A2%E5%AD%98%E3%81%AE%E7%94%9F%E5%BE%92%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E8%AA%AD%E3%81%BF%E8%BE%BC%E3%81%BF">既存の生徒モデルの読み込み</h1>
student_model = SmallModel(input_size=768, hidden_size=256, output_size=teacher_model.config.vocab_size)
student_model.load_state_dict(torch.load("small_model.pth"))
student_model.eval()  # 教師モデルとして評価モードに設定

<h1 id="%E6%96%B0%E3%81%97%E3%81%84%E7%94%9F%E5%BE%92%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E5%88%9D%E6%9C%9F%E5%8C%96">新しい生徒モデルの初期化</h1>
new_student_model = SmallModel(input_size=768, hidden_size=256, output_size=teacher_model.config.vocab_size).to("cuda")

<h1 id="%E3%83%87%E3%83%BC%E3%82%BF%E3%82%BB%E3%83%83%E3%83%88%E3%81%A8DataLoader%E3%81%AE%E6%BA%96%E5%82%99">データセットとDataLoaderの準備</h1>
dataset = WikipediaDataset("path/to/wiki.txt")
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

<h1 id="%E6%90%8D%E5%A4%B1%E9%96%A2%E6%95%B0%E3%81%A8%E3%82%AA%E3%83%97%E3%83%86%E3%82%A3%E3%83%9E%E3%82%A4%E3%82%B6%E3%81%AE%E8%A8%AD%E5%AE%9A">損失関数とオプティマイザの設定</h1>
criterion_soft = nn.KLDivLoss(reduction='batchmean')
criterion_hard = nn.CrossEntropyLoss()
optimizer = optim.Adam(new_student_model.parameters())

<h1 id="%E5%AD%A6%E7%BF%92%E3%83%AB%E3%83%BC%E3%83%97">学習ループ</h1>
num_epochs = 10
temperature = 2.0
alpha = 0.5  # ソフトターゲットとハードターゲットのバランス

for epoch in range(num_epochs):
    for batch in dataloader:
        # 入力の処理
        inputs = processor(batch, return_tensors="pt", padding=True, truncation=True).to("cuda")
        
        # 教師モデル（元の生徒モデル）の出力を取得
        with torch.no_grad():
            teacher_outputs = student_model(inputs.input_ids)

        # 新しい生徒モデルの出力を取得
        new_student_outputs = new_student_model(inputs.input_ids)
        
        # 知識蒸留損失の計算
        loss_soft = criterion_soft(
            nn.functional.log_softmax(new_student_outputs / temperature, dim=-1),
            nn.functional.softmax(teacher_outputs / temperature, dim=-1)
        )
        
        loss_hard = criterion_hard(new_student_outputs.view(-1, teacher_model.config.vocab_size), inputs.input_ids.view(-1))
        
        # 総合損失の計算
        loss = alpha * loss_soft + (1 - alpha) * loss_hard
        
        # 逆伝播と最適化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    logger.debug(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}")

<h1 id="%E6%96%B0%E3%81%97%E3%81%84%E7%94%9F%E5%BE%92%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E4%BF%9D%E5%AD%98">新しい生徒モデルの保存</h1>
torch.save(new_student_model.state_dict(), "new_small_model.pth")

logger.debug("新しい生徒モデルが保存されました。")

(これを作らせた会話はこちら) 余談ですが、Perplexityに書いてもらった理由は、ChatGPTだと教師モデル云々かんぬんがどれだけ知っているかわからなかったためです。

ちなみに自分は、AIをこんな感じで使い分けています。

とりあえずブレインストーミング系とかアイデア出しをしたいとき → Gemini
ファクトチェックができたらいいとき → Genspark
最新技術を使ったコードがほしいとき → Perplexity

Geminiを使っている理由としては、個人情報をGoogleだけにしまっておきたいからです。(ChatGPTはGoogleアカウントだけだと登録できない)

話が脱線しましたが、生成してもらったコードを見た感じだと、

文か何かを教師モデルに見せる
教師モデルが回答
生徒モデルも同じような回答になるようにパラメータを調整という感じでやっているようです。

ですが、手元には文か何かにあたるものがありません。

というわけで、何を用意するかというと、Wikipediaのダンプです。

憎きWikiextractor

Wikipediaのダンプを用意するとなったら、ダウンロード！

ですが、データはXMLになっていて、なおかつそもそも圧縮されています。これをやってくれるのがWikiextractor君なのですが、すでに2回悩まされています。

じゃあ成功した時のデータをとっておけという声が聞こえてきますが、なんか見つからなかったので仕方がありません。

で、3回目の今回も悩まされているのですが、永遠にできなくて、やる気をなくして現実逃避でこの記事を書いていたりします。

とりあえず、解決できたら色々追記すると思うので、待っていてください。

2025/01/09追記

ファイル整理してたら2回目に成功したときのデータが出てきました。やったね！

Wikipediaだけじゃ…

Wikipediaでは、とりあえず基本的な常識とかは身に付きますが、Markdownという結構重要なことが身に付きません。

では、どうすればいいのでしょうか。StackOverFlowのダンプです！

StackOverFlowは皆さんご存じのプログラミングの質問投稿サイトですが、Markdownを使用していたり、コードがあったりと、結構役に立ちます。

StackOverFlowのダンプは、archive.orgの中で7z形式で圧縮されたうえで公開されています。

ちなみに、ライセンスはどうかなと見たところ、WikipediaもStackOverFlowもCC BY-SA 4.0で公開されていたので、大丈夫でした。

これで、CC BY-SAとCC BY-NC-SAの組み合わせだったら、死んでいたところでした。

今後

とりあえず開発途中に書いた記事なので、今後もどんどん更新を入れていきます。

« 前の記事静カニ VS vim 次の記事 »友達の無茶