ローカル勾配の乗算
フォワードパス
ANDREA-120M のフォワードパスは、入力を一連の操作のシーケンスに通します:
x = embed(token_ids) # トークン埋め込み
for layer in 12_layers:
x = x + attn(LN(x)) # アテンションサブレイヤー
x = x + mlp(LN(x)) # MLPサブレイヤー
logits = LN(x) @ embed.T # 共有出力投影
loss = cross_entropy(logits, targets)
```
各操作は入力テンソルを読んで出力テンソルを生成します。前向きパスは、このバッチのクロスエントロピー損失という単一のスカラーで終了します。
後ろ向きパス
トレーニングは損失を減少させる方向に重みを更新します。更新方向を得るために、エンジンは以下を必要とします:
モデルのすべての学習可能な W に対する dL/dW
連鎖律がこれを与えます。連鎖 loss = f(g(h(x))) の場合:
dL/dx = (dL/df) * (df/dg) * (dg/dh) * (dh/dx)
各要因はローカル勾配です:1つの操作の出力が、その入力がわずかに変化したときにどのように変化するかを示します。ローカル勾配をグラフの後方へ乗算することで、損失信号をすべての重みに伝播します。
逆モード微分
バックプロパゲーションは逆順で勾配を計算します:dL/dlogits = 1 から始め、クロスエントロピー、出力量子化、レイヤーノーマライズ、12個のトランスフォーマーブロック、エンベディングへと後方へ進みます。各ステップで、入ってくる勾配をローカルヤコビアンで乗算します。
逆モードは、出力が単一のスカラー(損失)で入力が多数(重み)である場合に効率的です。1回の後方パスでモデルのすべての重みに対する勾配を生成します。前方モードでは重みごとに1回の通過が必要で、~120M重みのANDREA-120Mでは前方モードは非現実的です。
なぜ逆モードか
すべての順方向演算に後方双子が付与される
ペアリング規律
microgpt_cuda.cu はすべての操作に対して2つの CUDA カーネルを提供します:順方向出力を計算するものと、出力勾配から入力勾配を計算するもの。このペアリングは1対1です:
| フォワードカーネル | バックワードカーネル | 操作 |
|---|---|---|
k_embed_fwd | k_embed_bwd | トークン埋め込み検索 |
k_layernorm_fwd | k_layernorm_bwd | レイヤー正規化 |
k_attn_qkv_fwd | k_attn_qkv_bwd | Q, K, V 投影 |
k_attn_fwd | k_attn_bwd | スケールドドット積アテンション |
k_attn_out_fwd | k_attn_out_bwd | 出力投影 W_O |
k_mlp_fwd | k_mlp_bwd | MLP (GELU付き) |
k_residual_add | k_residual_add_bwd | 残差接続 |
k_loss_fwd | k_loss_bwd | クロスエントロピー損失 |
8つの演算ペアがトランスフォーマー全体をカバーします。加えて、いくつかのユーティリティカーネル:k_grad_norm_partial、k_grad_norm_final、k_grad_scale が勾配クリッピング用(アクティビティ75を参照)。
バックワードカーネルの仕事
後続のレイヤーから流入する勾配(grad_output)を与えられ、バックワードカーネルは以下を計算します:
1. grad_input: 操作の入力テンソルに関する勾配。これがさらに後方へ渡されます。
2. grad_weight: 操作内の学習可能パラメータに関する勾配。これがオプティマイザの状態に送られます。
両方は単一のカーネル起動で計算されます。CUDAスレッドが勾配テンソルのタイルで並列に協力します。
保存されたテンソル
後方計算ではしばしば前方パスからの値が必要です。例えば、k_layernorm_bwd は前方で計算された平均と分散を必要とし、k_mlp_bwd はGELU事前活性化を必要とします。トレーニングエンジンは前方でこれらを専用バッファに保存し、後方で読み込みます。
メモリコスト: 各保存テンソルに対してフォワード出力とほぼ同じ形状。ANDREA-120M で batch=8, seq=1024, d_model=768 の場合、1つの保存テンソルは 8 × 1024 × 768 × 4 bytes = 25 MB。12層にわたり各層で複数の保存テンソルがあると、トレーニング中のアクティベーションがVRAMを支配(24 GBカードで ~5-10 GB)。
1つのバックワードステップのトレース
メモリ内の勾配の保存場所
各重みテンソルごとに1つの勾配テンソル
ANDREA-120Mのすべての学習可能な重みテンソルには、同一形状の対応する勾配テンソルがあります。各ブロックごとに:
W_Q [768, 768] ↔ grad_W_Q [768, 768]
W_K [768, 768] ↔ grad_W_K [768, 768]
W_V [768, 768] ↔ grad_W_V [768, 768]
W_O [768, 768] ↔ grad_W_O [768, 768]
W_1 [768, 3072] ↔ grad_W_1 [768, 3072]
W_2 [3072, 768] ↔ grad_W_2 [3072, 768]
LN1.gamma [768] ↔ grad_LN1.gamma [768]
LN1.beta [768] ↔ grad_LN1.beta [768]
LN2.gamma [768] ↔ grad_LN2.gamma [768]
LN2.beta [768] ↔ grad_LN2.beta [768]
プラストークン埋め込み、位置埋め込み、および最終層正規化。グラディエントバッファの総メモリは重みのメモリと一致:~120M floats、FP32で~480 MB、FP16で~240 MB。
マイクロバッチにわたる蓄積
ANDREAのbatch_size = 8はFP16でVRAMに収まります。より大きな有効バッチサイズにはグラディエント蓄積が必要です:小さなバッチで複数のforward+backwardパスを実行し、グラディエントを同じバッファに合計し、その後1回のオプティマイザーステップを実行します。
for microbatch in range(n_microbatches):
forward(マイクロバッチ)
backward() # グラディエントバッファに追加します、上書きしません
scale_grads(1.0 / n_microbatches) # マイクロバッチ間で平均化
optimizer_step()
zero_grads() # 次のトレーニングステップのためにリセット
Backward カーネルは += セマンティクスを使用します、= ではありません。各呼び出しは既存のバッファにグラディエントの寄与を追加します。バッファは zero_grads() がクリアするまで累積和を保持します。
オプティマイザの状態
AdamW(アクティビティ 73)は、各重みごとにさらに2つのバッファを保持します:一次モーメント m と二次モーメント v。トレーニング時の総メモリ:
重み: 1× 重み数
勾配: 1× 重み数
Adam m: 1× 重み数
Adam v: 1× 重み数
保存された acts: 層とバッチによる ~2-4×
──────────────────────────────────────────
合計: ~6-8× 重み数
ANDREA-120M at FP16: ~240 MB × 4 buffers (weight, grad, m, v) + ~5-10 GB activations = ~10-12 GB total. RTX 4090の24 GB上限を余裕を持って下回る。ANDREA-12Mは1.4 GBで訓練; 10×のパラメータスケーリングで~10×のメモリ。
グラディエントバッファのサイズ調整
メモリと精度の完全制御
汎用フレームワークのコスト
PyTorch & JAX は autograd を便利にします: Pythonコードを書けば、勾配が自動的に得られます。コスト: あなたのコードとCUDAの間に汎用ディスパッチ層があります。すべての操作がPythonインタープリタのオーバーヘッド、フレームワークの帳簿管理、動的カーネル選択を通ります。1つのGPUで小さな言語モデルをトレーニングする場合、そのオーバーヘッドが重要です。
ANDREA が回避する具体的なコスト:
1. Python インタープリタのレイテンシ。 PyTorch の各オペレーションは Python/C++ の境界を越えます。トレーニングステップあたり約 100 回の カーネル起動で、約 9 ステップ/分 として、1分あたり約 900 回の境界越えが発生します。C レベルでのディスパッチによりこれを排除します。
2. フレームワークのアロケータの予測不可能性。 PyTorch のキャッシングアロケータは平均スループットは良好ですが、ピークメモリが予測不能です。ANDREA のトレーニングエンジンは起動時にすべてのバッファを事前割り当てします。トレーニング中は再割り当てなし、断片化なし、ステップ 100K での予期せぬ OOM なし。
3. 汎用カーネル選択。 PyTorch はランタイムでヒューリスティックによりカーネルを選択します。ANDREA はコンパイル時に RTX 4090 テンソルコアのタイルサイズに調整されたカーネルを選択します。
4. 混合精度の実装。 ANDREA-120MのFP16 cuBLASパスおよびANDREAのFP8 E4M3テンソルコア実験では、どのテンソルがどの精度で存在するかを精密に制御する必要があります。汎用フレームワークはこの制御を階層化されたAPIで公開しますが、カスタムCUDAはこれを直接記述します。
トレードオフ
カスタムCUDAのコスト:書くコードが増え、バグ探しが増え、コミュニティエコシステムがありません。ANDREAのmicrogpt_cuda.cuは約6000行の手書きCUDAで、デバッグに数ヶ月かかりました。各新しい操作にはフォワードカーネル、バックワードカーネル、テストを書く必要があります。
ANDREAが得るもの:
- 完全な再現性。 トレーニングパイプラインは1つのCバイナリと1つのPythonプロキシのみ。PyTorchリリース間のバージョンずれなし、フレームワークホイールとのCUDAバージョン不一致なし。
- ビット単位で正確な再開。 SIGTERMがGPUが見るすべてのテンソルを正確にキャプチャするチェックポイント書き込みをトリガー。再開は実行中の同じ損失軌道を継続。
- 予測可能なメモリ。 ANDREA-120Mが200Kステップ訓練されOOMなし。メモリはエンジン起動時にすべて計上。
- 直接ハードウェアアクセス。 テンソルコアのタイルサイズ、FP8 E4M3設定、非同期メモリコピー:すべてCUDAで直接アクセス可能、一般フレームワークでは不透明。
再現性をミッションとして
ANDREAホワイトペーパーのセクション9に完全な再現性スタックが記載されています:
``
``
トレーニングエンジン: microgpt/microgpt_cuda.cu
``
``
トレーニングプロキシ: microgpt/training_proxy.py
``
``
実験設定: experiments/ANDREA-*-TRAIN.json
``
``
データパイプライン: scripts/pull-hermes3.py, scripts/prep-megachat.py
``
``
ダッシュボード: scripts/live-loss-dashboard.html
``
``
バンディット仕様: docs/FIREHOSE-BANDIT.md
``
``
モデルドキュメント: docs/ANDREA.md ```
ハードウェア要件: 8 GB VRAM以上の1つのNVIDIA GPU (RTX 3060以上)。誰でもこれらのアーティファクトからANDREA-12Mを再現できます。カスタムCUDAパスがその理由の一部です: フレームワークのバージョン凍結なし、5年後の依存関係のサプライズなし。
### シグナルとチェックポイント
CUDAトレーニングループは2つのPOSIXシグナルに応答します:
- **SIGTERM**: 即時のチェックポイントを書き込み、その後終了します。トレーニングをクリーンに停止する際に使用されます。
- **SIGUSR1**: 即時チェックポイントを書き込み、トレーニングを継続。v3のpolish pivot中に実行を中断せずに状態をキャプチャするために使用。
チェックポイント形式: `[int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]`。ステップカウンタ、ウェイト数、次にウェイトに続いてAdamのモーメント。ビット単位で正確に再開。proxyはpolish時に`.samples.json` & `.state.json`を別々にアーカイブ;`.loss.json`はアーカイブされない(完全なトレーニング履歴を蓄積)。