English· Español· Deutsch· Nederlands· Français· 日本語· ქართული· 繁體中文· 简体中文· Português· Русский· العربية· हिन्दी· Italiano· 한국어· Polski· Svenska· Türkçe· Українська· Tiếng Việt· Bahasa Indonesia

un

ゲスト
1 / ?

ローカル勾配の乗算

Forward & Backward Kernels


フォワードパス

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では前方モードは非現実的です。

なぜ逆モードか

ANDREA-120M は約1億2000万の重みを持ち、各トレーニングステップで単一のスカラー損失を出力します。逆モード自動微分と順モード自動微分を比較し、(1) どのモードが1回の後方パスで全重み勾配を生成するか、(2) 全1億2000万の重み勾配を計算するために必要な順モードのパス数、(3) ANDREA がどのモードを使用しその理由を述べなさい。

すべての順方向演算に後方双子が付与される

ペアリング規律

microgpt_cuda.cu はすべての操作に対して2つの CUDA カーネルを提供します:順方向出力を計算するものと、出力勾配から入力勾配を計算するもの。このペアリングは1対1です:


フォワードカーネルバックワードカーネル操作
k_embed_fwdk_embed_bwdトークン埋め込み検索
k_layernorm_fwdk_layernorm_bwdレイヤー正規化
k_attn_qkv_fwdk_attn_qkv_bwdQ, K, V 投影
k_attn_fwdk_attn_bwdスケールドドット積アテンション
k_attn_out_fwdk_attn_out_bwd出力投影 W_O
k_mlp_fwdk_mlp_bwdMLP (GELU付き)
k_residual_addk_residual_add_bwd残差接続
k_loss_fwdk_loss_bwdクロスエントロピー損失

8つの演算ペアがトランスフォーマー全体をカバーします。加えて、いくつかのユーティリティカーネル:k_grad_norm_partialk_grad_norm_finalk_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つのバックワードステップのトレース

ANDREA-120M が1つのトランスフォーマーブロックを通るフォワードパスを完了します。その同じブロックを通るバックワードパスで何が起こるかをトレースしてください(pre-norm構造: `x = x + Attention(LN(x))` 然后 `x = x + MLP(LN(x))`)。バックワードカーネルを発火順に名前を挙げ、それぞれがどのフォワードカーネルとペアリングされるかを述べなさい。少なくとも4つのカーネルをカバー。

メモリ内の勾配の保存場所

各重みテンソルごとに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×のメモリ。

グラディエントバッファのサイズ調整

ANDREA-120M は約120,000,000個のパラメータを持ち、トレーニングステップあたり4つのマイクロバッチにわたって勾配蓄積を使用します。計算: (a) FP16での勾配バッファサイズをMBで; (b) FP16での重み + 勾配 + Adam m + Adam v の総メモリ; (c) トレーニングステップあたりに発火する別々の `forward()` + `backward()` 呼び出し数を。算術を示してください。

メモリと精度の完全制御

汎用フレームワークのコスト

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`はアーカイブされない(完全なトレーニング履歴を蓄積)。

なぜPyTorchを使わないのか

ANDREAは手動で`microgpt_cuda.cu`を書く代わりにPyTorchのautogradを使うことができた。ANDREAがカスタムCUDAを選んだ2つの異なるエンジニアリング上の理由を挙げなさい。一つの理由はメモリまたは精度制御を参照;もう一つの理由は再現性、フレームワーク依存、または長期メンテナンスを参照。