局部梯度相乘
前向傳播
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)
每個因子都是局部梯度:當其輸入小幅變化時,一個操作的輸出如何變化。透過圖形向後乘以局部梯度,將損失信號傳播到每個權重。
反向模式微分
反向傳播以反向順序計算梯度:從 dL/dlogits = 1 開始,然後向後穿越交叉熵、輸出投影、層歸一化、十二個 Transformer 區塊,然後是嵌入。在每個步驟中,將傳入梯度乘以局部 Jacobian。
當輸出是單一純量(損失)且有許多輸入(權重)時,反向模式很有效率。一個向後傳遞即可產生模型中每個權重的梯度。正向模式需要每個權重一個傳遞;對於 ANDREA-120M 擁有約 120M 權重,正向模式不可行。
為何使用反向模式
每個前向操作都有反向雙生
配對原則
microgpt_cuda.cu 為每個操作提供兩個 CUDA 核心:一個計算前向輸出,一個在給定輸出梯度時計算輸入梯度。配對為一對一:
| 前向核心 | 後向核心 | 操作 |
|---|---|---|
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 | 交叉熵損失 |
八個操作對涵蓋完整的 transformer。加上幾個實用核心:k_grad_norm_partial、k_grad_norm_final、k_grad_scale 用於梯度裁剪(見活動 75)。
反向核心的工作
給定從後層流入的梯度(grad_output),反向核心計算:
1. grad_input:與操作輸入張量相關的梯度。這會進一步向後傳遞。
2. grad_weight:與操作中可學習參數相關的梯度。這會進入優化器狀態。
兩者都在單一 kernel 啟動中計算。CUDA 執行緒並行合作處理梯度張量的圖塊。
保存的張量
反向計算通常需要前向傳播的值。例如,k_layernorm_bwd 需要前向傳播中計算的均值與變異數;k_mlp_bwd 需要 GELU 預激活值。訓練引擎在前向傳播期間將這些儲存在專用緩衝區中,然後在反向傳播期間讀取它們。
記憶體成本:大致與每個保存張量的前向輸出形狀相同。對於 ANDREA-120M,batch=8、seq=1024、d_model=768,一個保存張量為 8 × 1024 × 768 × 4 bytes = 25 MB。橫跨 12 層及每層多個保存張量,激活值在訓練期間主導 VRAM(~5-10 GB 在 24 GB 卡上)。
追蹤單一反向步驟
梯度在記憶體中的位置
每個權重張量對應一個梯度張量
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]
加上 token embeddings、position embeddings,以及最終一層 layer norm。梯度緩衝區的總記憶體與權重記憶體相符:~120M 個浮點數,在 FP32 下約 480 MB,在 FP16 下約 240 MB。
跨微批次的累積
ANDREA 的 batch_size = 8 在 FP16 下適合 VRAM。更大的有效批次需要梯度累積:在小批次上執行多次前向+反向傳播,將梯度加總到同一個緩衝區,然後執行一次優化器步驟。
for microbatch in range(n_microbatches):
forward(微批次)
backward() # ADD 到梯度緩衝區,不會覆寫
scale_grads(1.0 / n_microbatches) # 跨微批次平均
optimizer_step()
zero_grads() # 重置以進行下一個訓練步驟
反向傳播核心使用 += 語義,而不是 =。每次呼叫都會將梯度貢獻 ADD 到現有緩衝區;緩衝區會持有運行總和,直到 zero_grads() 清空它。
優化器狀態
AdamW(活動 73)為每個權重額外持有兩個緩衝區:一階矩 m 和二階矩 v。訓練時的總記憶體:
權重: 1× 權重數量
梯度: 1× 權重數量
Adam m: 1× 權重數量
Adam v: 1× 權重數量
saved acts: ~2-4× 依據層數與批次而定
──────────────────────────────────────────
總計: ~6-8× 權重數量
ANDREA-120M at FP16: ~240 MB × 4 個緩衝區 (權重、梯度、m、v) + ~5-10 GB 激活值 = ~10-12 GB 總計。遠低於 RTX 4090 的 24 GB 上限。ANDREA-12M 在 1.4 GB 中訓練;10× 參數擴展帶來 ~10× 記憶體。
調整梯度緩衝區大小
完全控制記憶體與精度
通用框架的代價
PyTorch 與 JAX 讓 autograd 變得方便:撰寫 Python 程式碼,即自動取得梯度。代價:您的程式碼與 CUDA 之間有一個通用調度層。每個操作都需經過 Python 解釋器開銷、框架記帳,以及動態核心選擇。在單一 GPU 上訓練小型語言模型時,此開銷很重要。
ANDREA 避免的具體成本:
1. Python 解釋器延遲。 每個 PyTorch 操作都會跨越 Python/C++ 邊界。每個訓練步驟約有 ~100 次核心啟動,每分鐘約 9 個步驟,即每分鐘 ~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 獲得的優勢:
- 完全可重現性。 訓練管道僅由一個 C 二進位檔案加上一個 Python 代理組成。無 PyTorch 版本發行間的版本漂移,無與框架 wheel 的 CUDA 版本不匹配。
- 位元精確的恢復。 SIGTERM 觸發檢查點寫入,精確捕捉 GPU 所見的每個張量。恢復會接續執行時相同的損失軌跡。
- 可預測的記憶體。 ANDREA-120M 訓練 200K 步驟無 OOM。記憶體在引擎啟動時即已計算。
- 直接硬體存取。 Tensor core 圖塊大小、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
Bandit 規格:docs/FIREHOSE-BANDIT.md
模型文件:docs/ANDREA.md
硬體需求:一台具備 ≥8 GB VRAM 的 NVIDIA GPU(RTX 3060 或更好)。任何人都可以從這些工件重現 ANDREA-12M。自訂 CUDA 路徑就是原因之一:無框架版本凍結,五年後無依賴驚喜。
信號與檢查點
CUDA 訓練迴圈會回應兩個 POSIX 信號:
- SIGTERM:立即寫入檢查點,然後退出。用於乾淨停止訓練時。
- SIGUSR1:寫入立即檢查點,繼續訓練。用於 v3 的 polish pivot 中捕捉狀態而不中斷運行。
檢查點格式:[int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]。步驟計數器、權重計數,然後是權重後跟 Adam 動量。精確位元恢復。代理在 polish 上單獨存檔 .samples.json & .state.json;.loss.json 永不存檔(它累積完整的訓練歷史)。