ゼロからLLMを書く 第34a回 — LLMトレーニング実行のためのJAXトレーニングループ構築
Sebastian Raschkaの書籍を参考にした学習の集大成として、ゼロからLLMを構築・訓練するプロジェクト。PyTorch版の最高モデル（3.2Bトークン学習、GPT-2小規模版に迫る品質）と同等の成果を、別フレームワークJAXを用いて再現する。本記事前半では、最小限の「A-to-Aモデル」を訓練するハーネスをFlax NNXとOptaxで構築し、データセット読み込みの最適化方法などを詳述する。
背景メモ
・著者のGiles Thomasは、Python製ワークフローエンジン「Prefect」の創業者で、個人ブログでLLMの仕組みをスクラッチから実装する連載を長年続けている。<br>・本エントリはその連載の一環。これまではSebastian Raschkaの書籍『Build a Large Language Model (from Scratch)』を教科書にPyTorchで進めてきたが、「教科書を見ずに自分の記憶とノートだけで完全にゼロからLLMを実装し、訓練する」という最終課題として、あえてフレームワークをPyTorchからJAXに切り替えている。<br>・JAXとはGoogleが開発した数値計算フレームワークで、関数型プログラミングのスタイルを採用し、自動微分とGPU/TPU上での高速実行が特徴。PyTorchに比べコミュニティは小さいが、深層学習研究の最前線ではよく使われる。<br>・Flax NNXとOptaxはJAXエコシステムの中核ライブラリで、Flax NNXがニューラルネットワークの構成要素（線形層など）を提供し、Optaxが最適化アルゴリズム（AdamWなど）を提供する。<br>・まずは訓練ループが正しく動くかを確認するため、最も単純なモデル（トークンを埋め込みに変換し、すぐに元の語彙空間に射影し直すだけ）を使い、入力と同じトークンを出力させる「A-to-Aモデル」でテストしている。これは次回以降で本格的なTransformerアーキテクチャに拡張する前段階。