HN 日本語サマリー

← 一覧へ戻る
AI・機械学習

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

Training a 3.8B LLM to 0.384 CORE for $998 – Hugo Vergnes (hugovergnes.github.io)

96 pointsby Anon8414 コメント

要約

Hugo Vergnes氏は、個人が数千ドルの予算で意味のある大規模言語モデル(LLM)をトレーニングできることを実証するため、3.8Bパラメータのモデルを998ドルでトレーニングしました。このプロジェクトでは、nanoGPTやnanochatに触発され、設定駆動型のフレームワークと最適化手法(FP8、ボキャブラリパディング、学習率スケジュールの調整など)を用いて、650億トークンを43時間で学習させ、COREスコア0.384を達成しました。これは、研究室や大企業以外でも手の届く範囲で最先端のAI開発が可能であることを示しています。

全文翻訳

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

更新 v1.1 — 2026-10-03 (🕒🦾 を忘れないで、わかる人にはわかる)

  1. 話題にしていただいています!共有してくださった皆様、そして興味深いフィードバックをくださった皆様に感謝します:hackernews, Vision AI。
  2. 明確化のためにいくつかの段落を書き直しました。
  3. さらに多くの実験とアブレーションを追加しました!

「nanoGPTのおもちゃ」と「研究室が必要」の間には、個人が数千ドルで意味のあるモデルをトレーニングできる、広大で十分に説明されていない領域が存在します。この記事はその領域への一歩です。

私は、自分自身でランダム性から言語と理解がどのように現れるのかを見て学びたかったのです。そして、ゼロから始めることでしか学べない部分を理解したかったのです。このプロジェクトは夜間にコーディングされ、5090でデバッグされ、レンタルしたB200で完成しました。Andrej Karpathy氏のnanochatに強く影響を受けました。彼の素晴らしい作品の数々もぜひチェックしてみてください!

結果は、65Bトークンを43時間で、998ドルでトレーニングされた、COREで0.384を記録した3.8Bパラメータのモデルです。

以下に、その旅で学んだことを記します。

モデルパラメータ数トークン数ハードウェア時間コストCORE
GPT-2 (OpenAI)1.5B————0.2565
nanochat d26~561M11.2B8× H100~3h—~0.258
nanochat d32~1B—8× H100~33h~$10000.310
little-lm 3.8B (1024 ctx)3.848B57.3B8× B20035.9h$8200.338
little-lm 3.8B (2048 ctx)3.848B65.3B8× B20043h$9980.384

私のモデルはnanochatのバリアントよりも大きいことに注意してください。同様の時間しかかかりませんでしたが、ハードウェアは異なりました。nanochatはH100を使用し、私はB200を使用しました。これは当時の実験では、単位あたりの作業量でより価値がありました。nanochatの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個のvocab × kv_dimのテーブルが、2層ごとに配置されています。


結果

初期実験

共有する価値のあるものができる前に、いくつかの期待外れの実験がありました。

探索として、FineWeb-Eduで16.4Bトークンを約5.8日間、単一のA100でトレーニングした858M Llamaモデルをトレーニングしました。AdamWを2.5e-4、コサイン減衰をゼロまで、5%のウォームアップ、勾配累積によるバッチサイズ256、コンテキスト長2048。すべて非常に標準的な設定です。

その結果、PIQA(CORE内で評価するためのより小さなタスク)で約60.45%という、かなり貧弱なモデルになりました。GPT-2 124Mは約63%を記録します。明らかに、そのコンピューティングリソースはうまく投資されていませんでした。そして実際、生成結果は典型的なGPT-2のように、繰り返しが多く、ほとんど意味をなさない文章でした。

事後分析は非常に明確でした:

  1. コサイン減衰は攻撃的すぎる。 カーブは、ステップの約70%を過ぎると完全に平坦になりました。計算予算の最後の30%は、学習率が低すぎる可能性があるため、実質的に何も生み出しませんでした。線形クールダウンは、より遅い段階でも有用なレートを維持します。
  2. ピークLRが保守的すぎる。 858Mパラメータに対して2.5e-4は低すぎます。それらの小規模モデルに対しては、かなり積極的に設定できます。
  3. すべてにAdamW。 このスケールでは、Muonは行列パラメータに対してトークンあたりで有意に優れているはずです。実際、これはアブレーション実行でかなり早く実証されました。
  4. データ。 FineWeb-Eduはまともです。しかし、ClimbMix(この時期にNVIDIAからリリースされたもの)の方が優れていることがわかりました。

5つの変更が加えられ、はるかに優れたモデルができました。

台形LRスケジュール。 5%のウォームアップ。平坦に保持し、最後の50%からピークの5%にかけて線形クールダウンで終了します。ポイントは、モデルが終盤を惰性で進むのではなく、最後まで学習し続けることです。3.8Bの実行では、評価損失は最終ステップでもまだ下降中でした。これは、858Mの実行が生成できなかったまさにその振る舞いです。

行列パラメータにはMuon、それ以外にはAdamWを使用します。Muonはステップあたりの速度は遅いですが(Newton-Schulz直交化は無料ではなく、浅い累積ベンチマークで約25%かかります)、そのコストはオプティマイザステップごとに一度だけ支払われます。勾配累積ステップが7回の場合、それは約4%に希釈されます。総実行時間に対して測定すると、全体的な収束ははるかに速くなります。

注:Muon » Adam。これまでのところ、私が試したすべてのプロジェクトで、MuonをAdamの代わりに(Vision、NLPなど)使用することで収束が大きく向上しました。

FineWeb-Eduの代わりにClimbMix。 これは収束速度において驚異的なジャンプでした。Karpathyが見出したことと全く同じです。

FP8 + ボキャブラリパディング。 torch._scaled_mm を使用したFP8トレーニングで、すべての3つのGEMMで動的なテンソルワイズスケーリングを行います。スループットが+33%向上しました。

ボキャブラリパディングは、ボキャブラリを50,257から50,304(64の倍数)にパディングしたため、テンソルコアが効率的に動作しました。スループットはごくわずかに増加しました。

2048ではなく1024コンテキスト。 コンテキストを半分にすると、メモリが固定された場合のバッチサイズが約2倍になります。スループットは約10%増加します(正確な設定によります)。MLPが依然として支配的であり、ハードウェアを効果的に使用している良い兆候です。

これらのパッチをすべて適用すると、実行ははるかに健全に見えました:

ステップトークン数評価損失CORE
2,5005.7B2.32780.2389
5,00011.5B2.20720.2752
7,50017.2B2.15710.2934
10,00022.9B2.12690.3104
12,50028.7B2.10750.3147
15,00034.4B2.07100.3224
17,50040.1B2.03950.3294
20,00045.9B2.01600.3267
22,50051.6B1.99630.3345
25,00057.3B1.98680.3384

定常状態(コンパイル、キューの充填など後)で約480,000トークン/秒。定期的なCORE評価を含め、総壁時計時間は35.9時間でした。

同じレシピを2048トークンコンテキストで再実行したところ、0.3840を記録しました。このギャップのほとんどすべては、非常にコンテキスト依存性の高い一部のタスクによるものでした。

安定した状態(コンパイル、キューの充填後など)では、GPUは92%のSMアクティビティ、40%のSM占有率でした。

これは、ストリーミングするのではなく、シャードをローカルにダウンロードしたことの成果です。実際、最初はHugging Faceから直接データをストリーミングしていました。しかし、ネットワークのハングに数回遭遇した後、トレーニング前にデータをローカルにプルすることにしました。

全体として、これらの数値にはかなり満足しています。トレーニングはコンピューティングバウンドで十分にデータが供給されており、ハードウェアをうまく使用しているという良い兆候であり、費やすすべてのドルをより良いモデルに変換できます。

これは、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とボキャブラリパディングへの切り替えについてはすでに説明しました。これらを組み合わせることで、スループットが+33%向上しました。

それに加えて、最適化されたカーネルを使用できます:LigerのFusedLinearCrossEntropyLossは、lm_headの行列積を損失に融合し、内部でチャンク化するため、完全な(B*T, vocab)のロジットテンソルが決してマテリアライズされません。同じバッチサイズで直接比較すると、6%遅くなります。

設定スループットVRAM
ベースラインCE、バッチ634,724 tok/s27,852 MiB
融合CE、バッチ632,952 tok/s19,630 MiB
融合CE、バッチ835,979 tok/s24,028 MiB
融合CE、バッチ1037,621 tok/s28,872 MiB

ステップあたりの速度は遅くなりますが、VRAMをかなりの量(私の5090で8GB)節約できるため、マイクロバッチサイズの増加が6%の損失を補って余りあります。Claudeは6%低下したため、すぐに却下しましたが、全体としてはスループットをわずかに向上させる良い方法でした。未使用のVRAMを補うためにバッチサイズを増やしたところ、累積で+44%になりました。

非ゲートMLP。 小規模モデルでゲート射影(SwiGLU → relu²、3つの行列積から2つに削減)をドロップします:183,035 → 214,173 tok/s、VRAMは6GB削減。アブレーションからの注意点:SwiGLUの中間比率2.75はrelu²には転移しません。モデルは顕著に悪く学習します。非ゲートの場合は4倍を使用してください。

bf16マスターウェイト。 オプティマイザのマスターウェイトをfp32ではなくbf16で保持すると、VRAMが27%削減され、1.5B設定でのスループットが640Kから1.4M tok/sに増加しました。これは2.2倍という大幅な高速化でした。品質のコストは現実的ですが小さいです:4,000ステップでCORE 0.22対0.23。能力あたりのコストを最適化する場合、慎重なdtype処理は、利用可能な最もレバレッジが高く、あまり議論されていないノブの1つです。

ハードウェア。 同じコード、150Mモデル、FP8:RTX 5090で184,662 tok/s、B200で477,440 tok/s。ハードウェアのみで2.59倍、さらにVRAMが増えたことでバッチサイズをさらに押し上げられることを考慮すると、それ以上です。

うまくいかなかったこと

ドキュメント境界マスキングとフレックスアテンション。 ドキュメントを1つのシーケンスにパッキングすると、トークンが境界を越えてアテンドできるようになるため、これを適切に修正しました:トークンごとにドキュメントIDを付与し、アテンションをマスクして、各トークンが現在のドキュメントにのみアテンドできるようにしました。これは正しいことですが、私の実験スケールではモデルの改善にはつながりませんでした。Andrej Karpathyも、BOSアラインドパッキング下ではドキュメント間の漏洩はそれほど悪化しないことを見出しました。ベストフィットパッキングが約10行でそれを置き換え、アテンションは非条件付きF.scaled_dot_product_attention(..., is_causal=True)に戻りました。多くの人が同様の発見を報告しています。当然、これはトレーニング分布の影響を受けますが、コードの簡潔さのために、効果のないこのメカニズムは破棄しました。

LigerのRMSNormとRoPE。 RoPEはマイクロベンチマークで2.2倍高速でしたが、エンドツーエンドのスループットには測定可能な変化をもたらしませんでした。RoPEはこのスケールでのクリティカルな計算ボトルネックの一部ではないため、ある程度予想されていました。RMSNormはPyTorch 2.9の組み込みF.rms_norm(0.41ms対0.25ms)よりも明らかに遅かったです。どちらも元に戻され、複雑さに見合うものではありませんでした。

nanochatスタイルの初期化。 埋め込みはN(0, 0.8)、線形重みは一様、出力射影はゼロ初期化され、残差ストリームは純粋なアイデンティティとして開始され、LMヘッドはN(0, 0.001)です。理論的にはGPT-2のどこでもN(0, 0.02)よりもはるかに優れています。損失カーブはわずかに低く開始し、2つのカーブは約1,500ステップで重なります。私の実験では測定可能な品質の違いはありませんでした。より原則的な初期化であるため、保持しました。

ストリーミングデータセット。 開始してさまざまなデータセットを試すのに最適です。しかし、実際のお金を使う際には単純にリスクがあります。ネットワークが健全に見えても、ローカルシャードの方がスループットが2〜3%高くなりました。時折発生するネットワークの低下は、それ以上のコストがかかることがあり、トレーニングが完全に停止することさえあります。数時間以上の実行では、最初に一度ダウンロードする(数分程度)方が良いと、かなり早い段階で確信しました。

バリュー埋め込みのアブレーション

バリュー埋め込みは、3.8Bモデルに対して721Mパラメータでした。私は同じモデルを同じ設定でvalue_embeddings: falseでトレーニングし、すでに費用を支払った元の実行と比較しました。ステップ12,500(29Bトークン)まで実行しました。

パラメータ数12.5K時点の損失12.5K時点のCOREスループット
バリュー埋め込みあり3.848B2.10750.3147479,445 tok/s
バリュー埋め込みなし3.128B2.11710.3047477,908 tok/s

損失が0.46%改善し、COREが3.2%改善しました。パラメータ数は19%増加しました。スループットは同一です。なぜなら、バリュー埋め込みはルックアップであり、メモリとオプティマイザの状態を消費しますが、実質的にFLOPは消費しないからです。

2つの興味深い発見があります:

  • バリュー埋め込みは、約1,200ステップ分のトレーニングに相当する効果をもたらしました。 これを価格設定する方法:ステップ10,000から12,500の間で、私のベースライン損失は0.0194低下しました。つまり、2,500ステップでその程度の改善が得られます。バリュー埋め込みの利点は0.0096で、その約半分です。25,000ステップのうち約1,200ステップと見なせます。したがって、パラメータ数が19%増加すると、トレーニングが約5%向上する価値があります。
  • COREは損失よりも約7倍多く変動しました(0.46%対3.2%)。そして、そのギャップはトレーニング中に着実に縮小しました。COREを使用して意思決定を行う場合は、これは知っておくと良いでしょう。COREは精度メトリックであるため、決定境界付近の項目はわずかなロジット変化で反転し、ランダムベースラインに対して中央揃えされているため、スコアがまだ低い相対的な差が拡大します。

バリュー埋め込みは小規模モデルにとって有用であり、スループットコストはほとんどかかりません。このための少量のVRAMを費やすことで、モデルはCOREに役立つ可能性のある特定の概念に対するバイアスを得ます。

GQAのアブレーション

このアブレーションスタディのために、1.5Bの中間モデルサイズを使用しました。これは、チンチラ最適性まで、またはそれを超えてトレーニングするのにずっと安価です。

同じ1.5Bレシピを3つの方法で実行しました:

  • 完全なマルチヘッドアテンション(MHA)
  • 2:1(6 KVヘッド)
  • 3:1(4 KVヘッド)。
KVヘッド数20K時点の損失20K時点のCOREスループットトレーニング時間
MHA122.14330.28571.235M tok/s9.4h
GQA 2:162.16520.26181.414M tok/s8.2h
GQA 3:142.17620.28121.380M tok/s8.4h

GQAは、評価損失のわずかなコストで、より高速にトレーニングされます。MHAのマイクロバッチサイズを減らしてOOMエラーを回避する必要がありました(これがスループット低下の一部を説明しています)。約12〜15%高速で、評価損失は1.0〜1.5%のコストです。予想通り、KVヘッドを共有するヘッドが多いほど、損失ペナルティは大きくなります。

いくつかの注意点:

  • このスケールでは、GQAはトレーニングでは元が取れません。 MHAは、トレーニング時間約8.2時間に対して、約7.9時間でGQA 2:1の最終損失に達します。3:1の場合は約7.4時間に短縮されます。これはMHAの損失がクールダウンの途中で取得されるため、GQAに有利な評価です。GQAの真の利点は推論時にあり、KVキャッシュが2〜3倍小さくなります。
  • COREはこれらの実行をランク付けするにはノイズが多すぎます。 3:1はCOREでは2:1を上回りますが、損失では劣ります。バリュー埋め込みのアブレーションと同じ教訓:COREはチェックポイント間で約0.01変動するため、これほど小さな差については評価損失を信頼します。
  • 3:1は2:1よりもわずかに遅かったですが、KVヘッドが少ないことを考えると驚きでした。私の最良の推測は、実際の効果ではなく、ノード間のばらつきです。

考察

誤解を招くマイクロベンチマーク

1024トークンのコンテキストで十分なCOREスコアが得られると信じてしまう誘惑に駆られるかもしれません。タスクごとのログを遡ると、その結論は非常にコンテキストに依存する一部のタスクでは間違っていることがわかります。

22のCOREタスクのうち3つは、プロンプトが実質的に1024トークンに収まらないものがあります:

タスクプロンプトが切り取られた割合ステップ 2.5Kステップ 25K
squad10570 / 10570 (100%)0.14780.0000
boolq3265 / 3270 (99.8%)0.57980.5131
bigbench_language_id9965 / 10000 (99.7%)0.24540.2538

SQuADは顕著です。停滞するのではなく、単調にゼロまで減少します:0.1478 → 0.0617 → 0.0099 → 0.0007 → 0.0000。モデルはトレーニングが進むにつれて、通常モデルが行わないような、このタスクで着実に悪化していきます。

2つの詳細がそれを説明しています。SQuADはDCLMバンドル内の10-shotタスクであるため、各プロンプトは10個の実行例とその後に実際の質問が続きます。私の評価データの中央値は1,998トークンです。1024トークンに収まるものはありません。そして、プロンプトが長すぎる場合、私のハーネスは最後のmax_seq_lenトークンを保持します。

テストの文章は最後に位置するため、常に残ります。テスト例は約169トークンです。切り取られたのは10個のデモンストレーションでした。モデルは文章と質問を読んでいましたが、期待される出力形式を教える例を見ることはほとんどありませんでした。SQuADは正解との正確なトークン一致で採点されるため、流暢な文章は毎回ゼロになります。

それが減少を説明しています。初期の、高エントロピーのモデルは、偶然一致する短い汎用的なものを時折出力します。シャープになるにつれて、それはよくできた継続にコミットし、偶然の一致は消えます。面白​​いことに、言語が上手になるほど、偶然正しく推測するのが下手になりました。

boolqは、同じ形状のより穏やかなバージョンを示しています。ステップ10,000(0.6294)でピークに達し、0.5131に低下します。言語識別は、偶然のレベルから全く動きません。

要するに、0.338は、モデルの品質とは全く関係なく、単にコンテキスト長のサイズに起因する理由で、22タスク中3タスクが偶然以下でスコアリングされた状態で測定されました。

より大きなコンテキストの効果

見てきたように、可能な限り高いCOREスコアを得たいのであれば、より大きなコンテキストが必要です。しかし、これにはトレーニングスループットへの影響があります。

コンテキスト長を倍増させ、VRAMを一定に保つためにマイクロバッチを半分にしたため、オプティマイザステップあたりのトークン数は同じままでした。レンタル時間の最後の数時間を節約するために、約28,000ステップで停止したため、学習率のウォームダウンは完全に完了せず、以下の数値は下限値です。

COREは0.3384から0.3840に向上しました。

[画像: 1024および2048コンテキスト実行の評価損失とCORE]

ステップ20,000では、2つの実行は評価損失で3桁まで同じ(2.016対2.016)であり、COREでは0.034の差がありました。ClimbMixデータセットで、COREと評価損失の間の低い相関関係が見られたのは驚きでした。

タスク10242048切り取り済み
squad0.00000.3114100% → 47%
boolq0.51310.709599.8% → 3.2%
bigbench_language_id0.25380.258599.7% → 14%
その他の19タスク+0.008 合計

squadとboolqだけで、ゲインの83%を占めます。 boolqが最も貢献しています。なぜなら、そのランダムベースラインが0.62(多数派クラス:約62%の回答が「はい」)であり、COREがそれを中央揃えするためです。生の+0.196は、中央揃えされた+0.517になります。それら2つを剥ぎ取ると、残りの20は合計で+0.008しか移動しません。これは、トークンが14%多いこと自体がもたらすものとほぼ同じです。

言語識別は、99.7%切り取りから14%切り取りになり、+0.005移動しました。これは、現在のモデルにとって、CORE評価ベンチマークで最も難しいタスクです。

いくつかのタスクは悪化しました:commonsense_qaは0.072低下し、cs_algorithmsは0.031低下しました。22タスク全体で、両方向へのいくらかの動きは予想されます。

2048は、品質ではなく、測定上の決定として支払う価値がありました。スループットが9%(480K→437K tok/s)低下しましたが、1024ではスコアリングできなかったタスクを除くと、ほとんど何も得られませんでした。1024はトレーニングには十分であり、モデルを良好なCOREスコアに到達させるための「安価な」方法です。2048は、コンテキストに大きく依存する一部のタスクをアンロックします。

後続の実験に関するもう一つの注意点: 1024→2048の効果は、より大きなモデルでより顕著になるようです。1.5Bの中間設定では、コンテキスト長を増やすとCOREスコアは約0.01しか向上しませんでした。しかし、3.8Bモデルでは、同じコンテキスト長の増加がCOREスコアを約0.05押し上げました。より長いコンテキストは、1024で切り取られていたタスクをアンロックします。そこから得られるCOREの増加は、そのタスクにおけるモデルの能力に依存します。

2048からのCOREゲインsquadboolqその他の20タスク合計
1.5B+0.015+0.003−0.003+0.014
3.8B+0.014+0.024+0.008+0.046

今後の作業

制限事項

私がアブレーションを行わなかった3つの点。ピークLRは、nanochatのsqrt(768/d_model)から。学習率のスイープにお金をかけたくありませんでした。858Mの事後分析のため、コサインから台形に移行しましたが、より効率的なスケジュールが存在する可能性があります。QK-normはデフォルトでオンになっており、一度もオフにしませんでした。

これらのほとんどは、ここでテストされたものではなく、nanochatから引き継がれています。これは、小規模な予算を費やすための正当な方法です—誰かがすでに実験費用を負担しています—しかし、それは私がKarpathyの結果が私のモデル、データ、スケールに転移することを信頼していることを意味します。

未解決の質問

もし時間とリソースがあれば追求したい興味深い作業がたくさんあります:

  • バリュー埋め込み対再割り当て。 上記の比較はVE対何もありませんでした。重要なのは、VEを他の何かに費やすこととの比較です。
  • 同壁時計時間での1024対2048。 再実行ではコンテキストを変更し、さらに長く実行したため、測定の問題は解決しますが、品質の問題は解決しません。
  • commonsense_qaがなぜ後退したのか、そのタスクは長いプロンプトを一切含まないにもかかわらず、より長いコンテキストで0.072も低下したのか。
  • オプティマイザのシャーディング、nanochatが行っている方法。 私は単一GPUのMuonでプレーンなDDPを使用しました。これは、各ランクがオプティマイザ状態の完全なコピーを保持し、同じNewton-Schulz更新を冗長に再計算することを意味します。nanochatはDDPラッパーを完全にドロップし、ZeRO-2シャーディングをオプティマイザ内で実行し、reduce-scatter、compute、all-gatherをオーバーラップさせます。メモリの利点は確実であり、解放されたメモリはバッチサイズに変換されます。これは同じコストでトークンを意味します。冗長な直交化がなくなるかどうかは、シャーディングの方法によります:Muonは完全な勾配行列を必要とするため、行列をランク間で分割しても役立ちませんが、各ランクに独自の行列全体を割り当てる場合は役立ちます。私はそれを全く探求していませんが、追加の機械装置のコストで総トレーニングスループットをさらに向上させるための素晴らしい方法だと思います。
  • 追加のデータ探索。 COREベンチマークまたはClimbMixデータセットのいずれについても、データ分析に多くの時間を費やすことができませんでした。同じコンピューティング予算でさらに高いパフォーマンスを引き出すのに役立つと確信しています。

締めくくり

GPT-2は2019年のフロンティア結果であり、資金力のあるラボが大規模なチームで制作したもので、その1.5BモデルはCOREで0.2565を記録しました。7年後、私は夜間に、998ドルで、時間単位でレンタルしたハードウェアで、それを大幅に上回りました。

フロンティアは移動し、すべてがそれに伴いました。ラボが必要だった作業は、今や一人のエンジニアが夜間に行えるようになりました。7年後、私たちはどのようなクレイジーなマシンを構築できるのだろうか!


付録:設定

YAMLインクルードからフラット化された実行全体を1つのブロックにまとめました。

model:
  hidden_size: 3072
  intermediate_size: 12288      # 4x, non-gated
  num_hidden_layers: 28
  num_attention_heads: 24
  num_key_value_heads: 8        # 3:1 GQA
  head_dim: 128
  hidden_act: relu2
  gated_mlp: false
  qk_norm: true
  logit_softcap: 15.0
  layer_scale: true
  value_embeddings: true        # 14 tables, alternating layers
  tie_word_embeddings: false
  rope_theta: 10000.0
  rms_norm_eps: 1.0e-6
  vocab_pad_to: 64              # 50257 -> 50304
  max_position_embeddings: 2048
  dtype: bf16

engine:
  compile: true
  fp8: true
  precision: bf16
  total_batch_size: 2293760     # 20 x 2048 x 7 grad_accum x 8 GPUs
  loss: LigerFusedLinearCrossEntropyLoss(softcap=15.0)

optimizer:                      # composite, one group per parameter class
  matrix:        Muon   lr=0.02      momentum=0.95  wd=0.0
  embeddings:    AdamW  lr=0.1414    betas=(0.8, 0.995)  eps=1e-10  wd=0.001
  lm_head:       AdamW  lr=0.002828  betas=(0.8, 0.96)   eps=1e-10  wd=0.01
  value_embeds:  AdamW  lr=0.0707    betas=(0.8, 0.995)  eps=1e-10  wd=0.01
  scalars:       AdamW  lr=0.005     betas=(0.8, 0.95)   eps=1e-10  wd=0.05

scheduler:
  trapezoidal:
    warmup_ratio: 0.05
    warmdown_ratio: 0.50
    final_lr_frac: 0.05

data:
  dataset: nvidia/Nemotron-ClimbMix  (karpathy/climbmix-400b-shuffle shards)
  tokenizer: gpt2 (tiktoken)
  block_size: 2048
  packing: best-fit, BOS-aligned
  batch_size: 20 per rank
  num_workers: 11

trainer:
  max_steps: 32000              # stopped at ~28,000 -> 65.3B tokens
  eval_every: 4000              # must divide max_steps or the final CORE is skipped

AdamWの学習率はnanochatのsqrt(768/d_model)スケーリングルールに従っています。MuonのLR 0.02もそこから引き継がれています。

付録:テキスト生成例

The capital of France is Paris. It is the largest city in France and the second largest city in Europe
The french revolution happened in 1789 and 1799, and was a time of great change in france
At the center of the milky way there is a supermassive black hole. It is called Sagittarius A* (pronounced
Electrons orbit around the nucleus of an atom in a series of energy levels. The energy levels are numbered
Newton discovered the laws of motion and gravity. He also discovered the law of universal gravitation. Newton's