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

un

guest
1 / ?
back to lessons

局部梯度相乘

前向與反向核心程式碼


前向傳播

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 權重,正向模式不可行。

為何使用反向模式

ANDREA-120M 有約 120M 個權重,並在每個訓練步驟產生單一純量損失。比較反向模式自動微分與前向模式。說明 (1) 哪種模式能在單一反向傳播中產生所有權重梯度;(2) 要計算所有 120M 個權重梯度需要多少次前向模式傳播;(3) ANDREA 使用哪種模式及原因。

每個前向操作都有反向雙生

配對原則

microgpt_cuda.cu 為每個操作提供兩個 CUDA 核心:一個計算前向輸出,一個在給定輸出梯度時計算輸入梯度。配對為一對一:


前向核心後向核心操作
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交叉熵損失

八個操作對涵蓋完整的 transformer。加上幾個實用核心:k_grad_norm_partialk_grad_norm_finalk_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 完成一個 transformer 區塊的前向傳播。追蹤該相同區塊的反向傳播(pre-norm 結構:`x = x + Attention(LN(x))` 然後 `x = x + MLP(LN(x))`)。依序命名反向 kernel,並說明每個與哪個前向 kernel 配對。至少涵蓋 4 個 kernel。

梯度在記憶體中的位置

每個權重張量對應一個梯度張量

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× 記憶體。

調整梯度緩衝區大小

ANDREA-120M 擁有約 120,000,000 個權重,並在每個訓練步驟中使用 4 個微批次進行梯度累積。計算:(a) FP16 下的梯度緩衝區大小(MB);(b) FP16 下權重 + 梯度 + Adam m + Adam v 的總記憶體;(c) 每個訓練步驟觸發多少個獨立的 `forward()` + `backward()` 呼叫。請顯示您的計算過程。

完全控制記憶體與精度

通用框架的代價

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 永不存檔(它累積完整的訓練歷史)。

為什麼不用 PyTorch

ANDREA 本可使用 PyTorch 的 autograd 而非手寫 `microgpt_cuda.cu`。給出兩個不同的工程原因,解釋為什麼 ANDREA 選擇自訂 CUDA。其中一個原因應參考記憶體或精確度控制;另一個應參考可重現性、框架依賴或長期維護。