gilesthomas-com

本記事は長期連載の集大成。著者はSebastian Raschkaの書籍を元に、ノートのみを頼りにスクラッチからLLMを構築・訓練。PyTorch版のコードを一切参照せず、JAXを使用。当初は入力と同じ系列を出力するだけの「A-to-Aモデル」から出発し、LayerNormやTransformerブロックなどを段階的に追加。最終的にGPT-2 Small相当のモデルをRTX 3090で37時間15分かけて訓練し、損失3.418を達成。これは同等のPyTorchモデル（3.538）や元のGPT-2 small（3.499）を上回る結果となった。

Sebastian Raschkaの書籍を参考にした学習の集大成として、ゼロからLLMを構築・訓練するプロジェクト。PyTorch版の最高モデル（3.2Bトークン学習、GPT-2小規模版に迫る品質）と同等の成果を、別フレームワークJAXを用いて再現する。本記事前半では、最小限の「A-to-Aモデル」を訓練するハーネスをFlax NNXとOptaxで構築し、データセット読み込みの最適化方法などを詳述する。

Charles Ye氏、Jasmine Cui氏、Dylan Hadfield-Menell氏の論文「Prompt Injection as Role Confusion」を基に、LLMが<system>や<user>などの「役割タグ」をほぼ無視し、代わりにテキストの「トーン（文体）」から役割を推測する傾向があることを解説。この現象が多くの脱獄（ジェイルブレイク）を説明できる可能性を示す。また、著者は自身の数当てゲーム実験や、埋め込みベクトルへの直接的な役割タグ付与による対策の可能性についても考察している。

JAX/Flax NNXの訓練ループで損失が減少しない問題をデバッグするため、パラメータ配列のハッシュ値を比較する手法を紹介。77Mものパラメータを逐一確認する代わりに、numpyのtobytes()とPythonのhash()を使って微小な変化も検出。原因は@jax.jitではなく@nnx.jitを使うべきところを間違えていたことで、FlaxのNNX APIにおける非関数型的なインプレース更新とJITコンパイルの互換性について学べる内容。

自宅LANを10Gb/sにアップグレードした際、10GBASE-T SFP+モジュールが過熱し、約95℃で自動シャットダウンを繰り返す「フラッピング」が発生。Marvellチップ搭載のMikroTik S+RJ10から、Broadcom BCM84891搭載の10Gtek ASF-10G-T80-INTに交換したところ、スイッチのCPU温度が約5℃低下し、リンクの安定性も改善された。ただし新しいモジュールはIntel製光ファイバーモジュールとして認識され、温度監視はできなくなった。

JAXでCPU上に作成した配列がコミットされていない場合、要素を取得するたびにGPUにデータが移動し、1回のルックアップに1秒以上かかる問題が発生する。jax.device_putで明示的にデバイスにコミットすることで、このオーバーヘッドを解消でき、ルックアップ時間を0.0002秒未満に抑えられる。

PyTorchからJAXへのLLMコード移植中、19GiBの大規模データセットをロードしようとした際にCUDAのメモリ不足エラーが発生。調査の結果、JAXはデフォルトでGPU（最速バックエンド）にデータを配置することが判明。本記事では、jax.devices()によるバックエンド・デバイスの確認方法や、jax.default_deviceコンテキストマネージャを使ってCPUに一時的にデフォルトデバイスを切り替える方法を解説する。

PyTorchのLLMコードをJAX/Flaxに移植する際、モデルのチェックポイント保存にSafetensorsを使いたい場合の注意点と解決方法。SafetensorsのFlax/JAX APIは単純なフラットな辞書構造を期待するが、Flaxの`nnx.State.to_pure_dict`で得られる辞書はネスト構造のためそのまま渡すとエラーになる。代わりに`nnx.to_flat_state`を使ってフラット化し、ドット区切りのキーに変換してからSafetensorsに渡すことで正しく保存・読み込みが可能になる。

長年PyTorchでLLMのトレーニングやRNNの再実装などに取り組んできた筆者が、最近注目したJAXフレームワークについての所感をまとめた記事。PyTorchとJAXの違いを「エンジニアリング vs 数学」「手続き型 vs 関数型」「最大主義 vs 最小主義」など4つのテーゼで大胆に整理し、jax.gradによる自動微分の仕組みやGradTracerによるトレーシングの仕組みを具体例とともに解説。最後にJAXの美しさと将来への懸念にも言及している。

自宅の10Gb/sスイッチで使っているMikroTik 10GBASE-T SFP+モジュールの高温対策として、Raspberry Pi用のミニヒートシンクを装着したところ、約3.5°Cの温度低下を確認。効果は劇的ではないものの、一定の改善が見られた。また、コミュニティの指摘から、SFP+モジュールにはMarvellチップ搭載の旧型（30m対応）とBroadcomチップ搭載の新型（100m対応）の2世代が存在することも紹介している。

筆者は自宅のネットワークを2.5Gb/sから10Gb/sにアップグレードした経緯を詳述。既存の構内配線の種類が不明だったため、書斎内の配線から段階的に導入し、各部屋やパッチパネル、ISP側のルーターなどを順次10Gb/s対応機器に交換。SFP+モジュールやUSBドングルの発熱が課題となるも、最終的には3台のスイッチとルーターを経由した実測で8〜9Gb/sの速度を達成した。

ISPが10Gb/sプランを提供開始したのを機に、著者は自宅の有線ネットワークを10GbEへアップグレードする過程で得た知見を綴る。20年以上前から家庭や小規模オフィスで使われてきたギガビットイーサネット（1Gb/s）は安定しているが、10Gb/sへの移行は熱管理やケーブル規格（CAT-5Eでは非対応）など新たな課題を突きつける。本記事では、10GBASE-Tの熱問題、SFP+やDAC（ダイレクトアタッチ銅線ケーブル）、アクティブ光ケーブル（AOC）といった低消費電力な代替手段について、過去のネットワーク技術の歴史を振り返りながら解説する。

『Build a Large Language Model (from Scratch)』の付録を読み終えた。PyTorchの解説、DataLoader、DDP、勾配クリッピング、学習率スケジューリング、LoRAなど、多くの有用な内容があった。付録を先に読めば時間を節約できただろうが、自分で調べて実装することでより深く理解できたと感じている。

著者は、自前のマシンで44時間かけてトレーニングしたモデルがGPT-2 smallにほぼ匹敵する性能に達し、「スクラッチからのLLM構築」プロジェクトの介入シリーズを完了した。学習率調整、ドロップアウト除去、重み減衰などの様々な介入を試し、最終的にローカルトレーニングで良好な結果を得た。次は書籍の付録を確認し、最終目標としてJAXフレームワークを使用した完全独自実装に挑戦する予定。

Sebastian Raschkaの書籍に基づくGPT-2-smallスタイルのLLM構築プロジェクトの一環で、命令ファインチューニングの評価方法を改善し、複数のモデル間で比較可能な結果を得るための新たなテストを実施。テスト損失と命令追従スコアの相関や、データセットの特性（FineWeb-Edu）がモデルの性能に与える影響について考察している。

GPT-2-smallスタイルのLLMを32億トークンで訓練し、57のチェックポイントで生成テキストの進化を観察。初期は無意味なトークンサラダだったが、訓練の約1/3（約10億トークン）で一貫したテキストを生成できるようになり、最終的にはより洗練された出力へと進化した。LLMは比較的早期に一貫性を獲得するが、正確性の向上にはさらなる訓練が必要である。

Sebastian Raschkaの書籍に基づいてGPT-2-smallスタイルのLLMを開発中。クラウドでのトレーニングで効果的な介入を特定した後、ローカルマシンで同等の品質を達成するため、バッチサイズを一致させる必要があった。RTX 3090では6シーケンスしか処理できないが、勾配蓄積を活用することで96の実効バッチサイズを実現し、クラウドトレーニングと同等のモデル品質をローカルで達成する方法を探る。

著者は、163MパラメータのGPT-2スタイルモデルに対して、勾配クリッピング、重み減衰の調整、ドロップアウトの削除、学習率のスケジューリング、QKVバイアスの追加など、複数の介入を組み合わせてクラウド上で訓練を実施。テストセットの損失が3.577761とこれまでで最良の結果を得たが、元のGPT-2の重みの性能にはまだ及ばなかった。

著者は、ゼロからトレーニングした163MパラメータGPT-2スタイルモデルの性能向上を目指し、勾配クリッピング、ドロップアウト削除、アテンション重みバイアス追加、学習率スケジューリング、重み減衰調整、重み共有、float32フル精度トレーニングなど、さまざまな介入を試した。結果、学習率スケジューリングが最も効果的だったが、改善幅は小さく、ランダムシードの違いによる初期重みの変動が結果に大きな影響を与える可能性があることが分かった。

この記事では、AMP（自動混合精度）と低精度行列乗算の最適化を無効にし、完全なfloat32精度でGPT-2小型ベースモデルを訓練した実験について報告しています。結果、テスト損失はわずかに改善したものの、訓練時間とコストが大幅に増加し、実用的なメリットは限定的であると結論づけています。

Lambda Labsで8x A100インスタンスを確保するため、利用可能になるまでAPIを監視し、起動時にTelegram通知を送信する自動化ツール「lambda-manager」を開発しました。現在の混雑状況では6時間経過してもまだ成功していませんが、トレーニング実行の機会を待っています。

Sebastian Raschkaの著書では、weight tying（重み共有）はパラメータ数を減らすが、モデルの性能を悪化させると指摘されている。しかし、著者は163Mパラメータの小さなモデルで実際に試し、この古典的な手法が有効かどうかを検証している。コード的には、埋め込み層の重みを出力層の重みとして転置して共有するだけで実装できる。

この記事では、GPT-2 smallベースモデルの訓練において、過学習を防ぐ正則化手法の一つである「weight decay」について探求しています。weight decayはモデルの重みの大きさ（L2ノルムの二乗）を損失関数に加えることでモデルの複雑さを抑制し、最適化の過程で自動的に適切なモデルサイズを見つけることを目指します。著者はAdamWオプティマイザにおけるweight decayの役割と、その値（0.1）が最適かどうかを検証しています。

この記事では、スクラッチから構築したGPT-2小型モデルの学習損失を改善する過程で、学習率の設定について探求しています。固定学習率の限界、学習率スケジューリング（ステップ減衰、指数減衰、コサイン減衰）、ウォームアップ戦略を解説し、Chinchilla論文で言及されたコサインサイクルについても考察します。適切な学習率の初期値を見つける難しさにも触れています。

GPT-2 smallモデルのテスト損失を改善するため、アテンション重み行列にバイアスを追加する実験を行った。QKVバイアスを有効にしたモデルは、ベースラインよりもテスト損失が0.023改善し、トレーニング中の損失スパイクも小さくなるなど、予想外の効果を示した。