TT Lab
はじめる
学ぶ 学習パス コース

MiniMind — 小さな言語モデルを最初から最後まで自分で学習する

MiniMind の LoRA で口調だけを変え、基盤の重みがそのままか確かめる

TT Labで続きを見る

目標

MiniMindのmodel_lora.pyのapply_lora・save_lora・load_loraで、SFTの基準モデルにrank 8のLoRAを載せます。基盤の重みを凍結してLoRAだけを学習し、答えを「…だニャン。」の口調(韓国語の猫の語尾にあたる口調です)に変え、基盤の重みが1ビットも変わっていないかと、統合したモデルが同じ出力を出すかを確認します。

なぜ重要なのか

LoRAの約束は2つです。少ないパラメーターだけを学習すること、基盤モデルはそのままにしておくことです。前者はメモリと保存を減らし、後者は、1つの基盤の上にLoRAを複数、差し替えられるようにします。ところが、凍結を忘れると、2つの約束がどちらも静かに壊れます。学習はうまくいき、口調も変わりますが、基盤の重みまで動いて、ほかのLoRAがすべて狂います。 そのため、このラボは、結果(口調)だけを見るのではなく、学習後の基盤の重みを、元のものとテンソルごとに比べます。MiniMindの実装には、アルファの倍率がなく、次元が同じLinearにだけ付くことも、コードで確認します。

ステップ

  1. /opt/mm/ref/sft.pthを読み込んでapply_lora(model, rank=8)を行い、LoRAが付いたモジュールの名前を、rank・targetsとして書いてください(出力先: /root/mm/lora/targets.json)。
  2. 載せる前と直後のロジットの最大の差を、max_abs_diffとして書いてください(出力先: /root/mm/lora/zero.json)。
  3. 名前にloraが入っていないパラメーターをすべて凍結して、学習するパラメーター数・全体の数・割合を書いてください(出力先: /root/mm/lora/params.json)。
  4. lora_nyang.jsonlでLoRAだけを150学習ステップ学習して、save_loraで保存してください(出力先: /root/mm/lora/lora.pth)。
  5. 学習が終わったモデルから、LoRAを除いた重みを保存してください(出力先: /root/mm/lora/base_after.pth)。元のSFTの重みと同じである必要があります。
  6. 取り分けておいた質問(lora_nyang_val.jsonl)で、韓国語の猫の語尾にあたる「ニャン。」で終わる答えの割合を、LoRAなしとありで測って、書いてください(出力先: /root/mm/lora/effect.json)。
  7. W + B·Aで統合した重みを保存してください(出力先: /root/mm/lora/merged.pth)。
  8. ## 어디에 붙나 ## 무엇이 바뀌나 ## 합치기の3つのセクションを書き、学習したパラメーター数と、LoRAを載せたあとの語尾「ニャン」の割合を入れてください(出力先: /root/mm/lora/report.md)。韓国語の見出しは、順に「付く場所」「何が変わるのか」「統合」という意味です。

参考

LoRAがどこに付くかを調べる

mmkit.load_model("/opt/mm/ref/sft.pth")で読み込んでapply_lora(model, rank=8)を行ったあと、lora属性ができたモジュールの名前を順に、{"rank": 8, "targets": [...]}の形で書いてください(出力先: /root/mm/lora/targets.json)。

model.named_modules()を回して、hasattr(m, 'lora')のものを集めればよいです。apply_loraは、in_features == out_featuresのLinearにだけ付きます。私たちのモデルで、どの射影がその条件を満たすのか、先に予想してみてください。

載せた直後は何も変わらない

/opt/mm/ref/val.npyの先頭64トークンに対するロジットを、apply_loraの前と直後に求めて、最大の絶対差を、max_abs_diffとして書いてください(出力先: /root/mm/lora/zero.json)。

MiniMindのLoRAは、Aを正規分布で、Bを0で初期化します。B·Aが0なので、脇の枝の出力が0で、学習は「基盤モデルからどれだけ離れるか」を0から学びます。

LoRAだけを学習するように凍結する

学習スクリプト(/root/mm/lora/train_lora.py)を書きはじめてください。LoRAを載せたモデルで、名前にloraが入っていないパラメーターはrequires_grad=False、入っているものはTrueにして、学習するパラメーター数(trainable)・全体の数(total)・割合(ratio)を書いてください(出力先: /root/mm/lora/params.json)。

MiniMindのtrain_lora.pyが使っている規則そのままです。1か所あたりA(8×128) + B(128×8)で、付いた場所の数を掛ければよいです。全体には、LoRAのパラメーターも入ります。

口調をLoRAで学ぶ

/opt/mm/data/lora_nyang.jsonlをmmkit.sft_encodeで変換して、LoRAのパラメーターだけをAdamW(lr 5e-3)に渡して150学習ステップ(バッチ16・シード42)学習し、save_lora(model, "/root/mm/lora/lora.pth")で保存してください(出力先: /root/mm/lora/lora.pth)。

save_loraは、LoRAが付いたモジュールごとに、…lora.A.weightと…lora.B.weightだけをfp16で保存します。ファイルのサイズを、基盤モデル(4MB強)と比べてみてください。採点ツールは、キーがLoRAだけか、Bが0から動いたかを見ます。

基盤の重みがそのままかを確かめる

学習が終わったモデルのstate_dictから、名前に.lora.が入ったものを除いた残りを保存してください(出力先: /root/mm/lora/base_after.pth)。採点ツールが、元の/opt/mm/ref/sft.pthとテンソルごとに比べます。

凍結が正しくできていれば、1ビットも違いません。requires_gradをオフにせず、オプティマイザーにすべてのパラメーターを渡したなら、基盤の重みも勾配を受けて動き、ここで引っかかります。PyTorchのオプティマイザーは、勾配のない(None)パラメーターを飛ばすので、凍結さえ正しくできていれば、オプティマイザーに何を渡しても、基盤はそのままです。

口調がどれだけ変わったかを測る

lora_nyang_val.jsonlの質問60個に、グリーディ生成で答えさせて、猫の語尾「ニャン。」にあたる韓国語の語で終わる答えの割合を、基盤モデル(base_rate)とLoRAを載せたモデル(lora_rate)で測って、書いてください(出力先: /root/mm/lora/effect.json)。

LoRAを載せたモデルは、基盤モデルを読み込んで、apply_lora(rank=8)のあとload_lora(model, 경로)で作ります(プレースホルダーはパスです)。取り分けておいた質問は、LoRAの学習のときに見たことのない組み合わせなので、口調が「質問の形」ではなく、「答え方」として学ばれたかを見ます。

統合して脇の枝をなくす

lora.pthのA・Bで、該当するモジュールの重みにB @ Aを足して(アルファの倍率なしで)統合したstate_dictを、保存してください(出力先: /root/mm/lora/merged.pth)。採点ツールは、この重みをLoRAのないモデルに入れて、LoRAを載せたモデルと同じロジットを出すかを見ます。

Aの形は(rank, 入力)、Bは(出力, rank)なので、B @ Aは(出力, 入力)で、元のweightと同じ形です。MiniMindのmerge_loraも同じことをして、fp16で保存します。ここでは、比較のために、fp32で保存します。

LoRAの約束を確認した記録を残す

## 어디에 붙나 ## 무엇이 바뀌나 ## 합치기の3つのセクションを書き、ステップ3のtrainableとステップ6のlora_rateを数字で入れてください(出力先: /root/mm/lora/report.md)。韓国語の見出しは、順に「付く場所」「何が変わるのか」「統合」という意味です。

LoRAのファイルが基盤モデルの何分の1なのか、アルファの倍率がない実装なので、rankを変えるとき何を見直す必要があるのかも、1行書いてみてください。