3.8B LLMを998ドルでCORE 0.384までトレーニングする

「nanoGPTのおもちゃ」と「研究室が必要」の間には、個人が数千ドルで意味のあるモデルをトレーニングできる、広大で説明不足な領域があります。私は、ランダムな重みから言語と理解がどのように生まれるのかを自分で確かめたく、ゼロから始めることでしか学べない部分を学びたかったのです。このプロジェクトは夜間に作成され、5090でデバッグされ、レンタルしたB200で完成しました。Andrej Karpathy氏のnanochatに強く影響を受けています。結果として、COREで0.384を記録した3.8Bパラメータのモデルが、650億トークンを43時間で998ドルでトレーニングされました。以下に、うまくいったこと、いかなかったこと、そしてまだ分からないことを記します。 モデル パラメータ トークン ハードウェア 時間 コスト CORE GPT-2 (OpenAI) 1.5B — — — — 0.2565 nanochat d26 ~561M 11.2B 8× H100 ~3h — ~0.258 nanochat d32 ~1B — 8× H100 ~33h ~$1000 0.310 little-lm 3.8B (1024 ctx) 3.848B 57.3B 8× B200 35.9h $820 0.338 little-lm 3.8B (2048 ctx) 3.848B 65.3B 8× B200 43h $998 0.384 私のモデルはnanochat d32よりも大きく、同様の壁時計時間で完了しました。B200はH100よりも単位あたりの作業効率で優れていました。しかし、nanochatの1,000ドルの構成とほぼ同じ費用で、これほど大きく上回る結果が得られました。研究所や数百万ドルのコンピューティング予算を持つ大企業以外で達成可能なことについての、励みになるデータポイントです。フロンティアが進むにつれて、1,000ドルでより遠くまで到達できるようになります。 セットアップ 私はlittle-lmを、小さなデコーダーオンリーLLMをトレーニングするための設定駆動型フレームワークとして構築しました。すべての実行はYAMLファイルによって完全に指定されます：モデル、データセット、オプティマイザ、スケジュール、コールバック。コンポーネントはグローバルレジストリに自己登録され、名前で解決されるため、オプティマイザやデータセットの切り替えは1行の設定変更で済みます。優れたインフラストラクチャは、ほぼ即座に元が取れます。通常のソフトウェアエンジニアリング規律（関心の分離、クリーンなインターフェース、交換可能なコンポーネントなど）は、AIの仕事において非常に重要です。初期段階で少し、その後も悪い契約や最適化不足を修正するために何度かコストがかかりましたが、この時間投資は、最初に出会う収束問題で元が取れます。優れたインフラストラクチャとは、コードを手動で編集する必要がほとんどないインフラストラクチャだとわかりました。設定を読んで何が起こるかを正確に理解でき、隠れたメカニズムがなければ、それは良い仕事をしたことを意味します。以下のレポートは、実験をブランチではなく3行のYAML差分として表現できた結果です。 最終モデルはLlamaスタイルです：RMSNorm、RoPE、GQA（24クエリヘッド、8 KVヘッド）、relu² MLP、QK-norm、logit softcap、レイヤーごとの学習可能な残差スカラー、およびResFormerスタイルのバリュー埋め込み。 コンポーネント パラメータ トークン埋め込み 154.5M LMヘッド（非結合） 154.5M 28デコーダーレイヤー 2,818.7M バリュー埋め込み（14テーブル） 721.2M 合計 3.848B バリュー埋め込みがパラメータ数の19%を占めることは注目に値します。14テーブル × ボキャブラリ × kv_dim、各レイヤーおきに1つです。 結果 初期実験 良い実行の前に、多くの悪い実行がありました。私はFineWeb-Eduで858M Llamaを16.4Bトークン、単一のA100で5.8日間トレーニングしました。AdamWを2.5e-4、コサイン減衰をゼロまで、ウォームアップ5%、バッチ256（勾配累積）、コンテキスト2048。結果：PIQA 60.45%。GPT-2 124Mは約63%を記録します。私は6日間のコンピューティング時間を費やして、2019年の7倍小さいモデルよりも悪いものを作ってしまいました。生成は繰り返しが多く、ほとんど意味をなさなかったです。損失曲線がその物語を語っていました。コサイン減衰をゼロまで。曲線は約70%のステップで完全に平坦になりました。最後の30%のコンピューティング予算は、学習率が低すぎる可能性があるため、実質的に何も生み出しませんでした。線形クールダウンは、より長い間有用なレートを維持します。ピークLRが保守的すぎました。858Mパラメータに対して2.5e-4は低すぎます。それらの小さなモデルに対してはかなり積極的に行えます。すべてにAdamWを使用。このスケールでは、Muonは行列パラメータに対してトークンあたりで大幅に優れているはずです。実際、これはアブレーション実行でかなり早く実証されました。 データ。 FineWeb-Eduはまともです。利用可能な中で最高ではありません。そのポストモーテムから5つの変更点が出てきました。これらを組み合わせることで、上記の実行と、GPT-2を大幅に上回るモデルとの違いが生まれます。 台形LRスケジュール。 5%のウォームアップ。最後の50%から5%のピークまで線形クールダウンで平坦に保持し、終了します。ポイントは、モデルが終盤に惰性で進むのではなく、最後まで学習を続けることです。3.8B実行では、評価損失は最終ステップでもまだ下降しており、これはまさに858M実行が達成できなかった動作です。 Muonを行列パラメータに、AdamWをそれ以外すべてに使用。 Muonはステップあたりの速度は遅いですが（Newton-Schulz直交化は無料ではなく、浅い累積ベンチマークで約25%）、そのコストはオプティマイザステップごとに1回支払われます。勾配累積7ステップでは、約4%に希釈されます。総実行時間に対して測定すると、全体的な収束ははるかに速くなります。 ClimbMixをFineWeb-Eduの代わりに。 これは収束速度において劇的な飛躍をもたらしました。Karpathy氏が発見したとおりです。 FP8 + ボキャブラリパディング。 torch._scaled_mmを使用したFP8トレーニングで、すべての3つのGEMMで動的なテンソルワイズスケーリングを行い、ボキャブラリを50,257から50,304（64の倍数）にパディングして、テンソルコアが快適に動作するようにしました。これらを合わせて、スループットが+33%向上しました。これは主にFP8によるものです。 コンテキスト1024（2048の代わりに）。 コンテキストを半分にすると、固定メモリでバッチサイズが約2倍になります。トークンあたりのスループットはほとんど変わりません。MLPが支配的であることは、ハードウェアを効果的に使用している良い兆候です。以下でコンテキスト長のモデルへの影響について議論します。 以下に、完全な実行を示します。 ステップ トークン 評価損失 CORE 2,500 5.7B 2.3278 0.2389 5,000 11.5B 2.2072 0.2752 7,500 17.2B 2.1571 0.2934 10,000 22.9B 2.1269 0.3104 12,500 28.7B 2.1075 0.3147 15,000 34.4B 2.0710 0.3224 17,500 40.1B 2.0395 0.3294 20,000 45.9B 2.0160 0.3267 22,500 51.6B 1.9963 0.3345 25,000 57.3B 1.9868 0.3384 定常状態での約480,000トークン/秒は、57.3Bトークンを33時間で処理することに相当します。壁時計時間は35.9時間でした。差分はCORE評価で、それぞれ約15分（実行中に10回）かかり、合計の7%を消費しました。 この全く同じレシピを2048トークンのコンテキストで再実行したところ、0.3840を記録しました。このギャップのほとんどすべては、コンテキストに大きく依存するタスクによるものでした。 GPU自体について：SMアクティビティ92%、SMオキュパンシー40%。高いアクティビティは、SMがほとんどアイドル状態にならなかったことを意味します。データローダーの枯渇やネットワーク待ちはありません。これは、ハギングフェイスの小さなネットワークハングに脆弱になるストリーミングではなく、シャードをローカルにダウンロードしたことの成果です。低いオキュパンシーは、バックツーバックの大きなGEMMがどのように見えるかです：行列乗算カーネルは、レジスタタイルのサイズのために意図的にオキュパンシーをトレードオフします。コンピューティングバウンドで十分に供給されていることは、ハードウェアをうまく使用している良い信号であり、費やすすべてのドルをより良いモデルに変換できます。これはB200あたり約1,047 TFLOP/sの持続スループット、またはBlackwellの密なFP8ピークに対して約25%のMFUです。（bf16ピークに対しては50%と表示されますが、これは線形レイヤーのすべてがFP8で実行されるわけではないため、より重要な数値です。） 分散戦略は、単純なDistributedDataParallelです。単一ノードで3.8Bの場合、勾配通信が制約になることはなく、シャードオプティマイザの仕組みは不要であることが判明しました。 スループットの向上 GPUのレンタルは安くはありません。職場では、モデルのコストよりも品質について考えることが多いです。自分の金が燃えているときは、突然スループットがはるかに重要になります。これは、ノードをレンタルする前に、単一のRTX 5090で実際の作業を必要としました。 ベースライン858Mモデル、bf16、コンパイル済み：26,144 tok/s。 最終：37,621 tok/s。 FP8（+25%）。 すべての3つのGEMM（フォワード1つとバックワード2つ）をFP8で、動的なテンソルワイズスケーリングを使用。SM90+が必要ですが、これは非常に良いスループットのジャンプです。 ボキャブラリパディング（累積+33%）。 50,257から50,304へのパディングは、47個の未使用の埋め込み行をコストとし、高速なテンソルコアパスをアンロックします。ほぼ無料です。 融合線形クロスエントロピー（累積+44%）。 LigerのFusedLinearCrossEntropyLossは、lm_headの行列乗算を損失に融合し、内部でチャンク化するため、完全な（B*T、ボキャブラリ）ロジットテンソルが決してマテリアライズされません。測定