Google ColabでフルスクラッチLLMを学習してみた話
はじめに
小規模Transformerをスケールさせると何が起きるか？ということで、
LLM開発の雰囲気を少しでも体感したく、Google Colab上でフルスクラッチLLMの学習を試してみました。
本記事では、
・Hugging Faceなどの既成モデルではなく
・フルスクラッチLLM（kazmot (x:@kmotohas) さん提供の実装）をスケールした場合に
・データ量やモデルサイズを増やすと何が起きるのか
を、実験結果ベースで共有します。
なお、本記事は精度追求やSOTA更新が目的ではなく、
「LLM学習を実際に回すと、どこがボトルネックになるのか」を掴むことを目的としています。
実験環境と前提
・実行環境：Google Colab Pro+
・フルスクラッチLLMコード：kazmot (x:@kmotohas) さん提供の実装を利用
・GPUメモリ制約を考慮し、小規模モデルから段階的に拡張
フルスクラッチLLMの概要
モデルアーキテクチャ
・Decoder-only Transformer（小規模LLM）
・総パラメータ数：約1,300万（初期設定）
詳細設定（初期）
・埋め込み次元数：128
・Transformerブロック数：4
・各ブロックのMulti-Head Attentionヘッド数：8
・層数：4
・コンテクスト長：128
👉 Colab環境で無理なく学習を回せる最小構成になっています。
事前学習用コーパス
・fineweb-2-edu-japanese
このデータセットを以下の形で利用しています。
・平均文字数：約654文字
・初期実験では 100件のみ抽出
👉 まずは挙動確認のため小規模から始めます。
トークナイザー
・GPT-2 BPE
・語彙数：50,257
学習設定
・ドロップアウト率：0.1
・バッチサイズ：2
・Optimizer：AdamW
・学習率：0.0005
学習結果
学習結果①：素のままで動かす
まずはパラメータ変更なしで学習を回してみます。
学習結果
・初期Loss：11.0228
・最終Loss：3.7960
・改善率：65.6%
※改善率は以下で算出しています。
(初期Loss - 最終Loss) / 初期Loss
所感
・Lossはしっかり下がる
・ただし生成結果は、意味の通らないトークン列が中心
・「学習はしているが、言語モデルとしては未熟」な状態
学習結果②：データ数とパラメータを増やす
次に、データ量とモデル容量を同時に拡張します。
変更点
データセット
・100件 → 10,000件
モデルパラメータ
・ヘッド数：8 → 16
・層数：4 → 8
・コンテクスト長：128 → 256
・総パラメータ数：約1,300万 → 約3,200万
👉 GPUメモリにまだ余裕があったため、素直にスケールさせました。
学習結果
・初期Loss：10.9922
・最終Loss：2.3745
・改善率：78.4%
学習前の生成例：
わたしが Carney Pac towed infall shuffaren amongst Participation
percentagesroman back abilities coolerVC uncle Orient internal Kanyeidine
Several bloc coerд 65lifeadalvor>< )); Unit Cousmart
学習後の生成例：
わたしがちょっと楽しくなっているのです。
それは、それはそれを楽しむことができます。
それは、それはそれを
評価
良くなった点：
・日本語として破綻していない文章が生成され始めた
・Lossが大きく改善
問題点：
・文のループが多く、長文生成は不安定
・Lossの振れ幅が大きく、学習が安定しない
学習結果③：バッチサイズ・層数・コンテクスト長をさらに拡張
Lossの振れが激しかったため、
ミニバッチが小さすぎることによる勾配のばらつきが原因と考え、設定を見直します。
また、GPUメモリ使用率にまだ余裕があったのでパラメーターもスケールさせます。
変更点
・バッチサイズ：2 → 32
・層数：8 → 10
・コンテクスト長：256 → 2048
・総パラメータ数：約3,200万（ほぼ据え置き）
学習結果
・初期Loss：11.0036
・最終Loss：1.9126
・改善率：82.6%
学習後の生成例：
わたしが、その言葉を聞くことができる。
その言葉は、その言葉を言葉にしている。
その言
評価
・Lossは右肩下がりで安定
・学習曲線の振れも大きく改善
・改善率も今回が最良
👉 一方で、生成文はまだ短く、意味の展開は弱いままでした。
なお、モデルサイズをさらにスケールさせることを試みましたが、GPUメモリ不足エラーとなりました。
結果として、Colab Pro+環境ではこの構成が実行限界でした。
まとめ
今回の実験から、以下のことが分かりました。
・データ数・パラメータ数を増やすことで、生成品質は確実に向上する
・しかし、バッチサイズやコンテクスト長を適切に取らないと学習は不安定になる
・Colab環境では、GPUメモリと学習時間が明確なボトルネックになる
今後試したいこと
・勾配累積による実効バッチサイズの拡張
・トークナイザーの変更
・Attention構造や正規化方法の見直し
「LLMを学習させるとはどういうことか」を体感する、という目的は十分に達成できたと感じています。
Discussion