為什麼一張 GPU 永遠不夠：記憶體牆與資料平行

記憶體究竟去了哪裡

在擴展任何東西之前，你得先知道是什麼塞滿了裝置。對一個有 P 個參數的模型，訓練一步會持有四個大張量：權重（weights）、梯度（gradients）、優化器狀態（optimizer states），以及為反向傳播保留的激活值（activations）。在混合精度下用 Adam，光是前三者每個參數就約需 16 個位元組——bf16 權重 2 個、bf16 梯度 2 個，再加上 fp32 主副本與兩個 Adam 動量共 12 個。因此一個 70 億參數的模型光是「存在」就要約 112 GB，這還沒算進任何激活值記憶體。這個數字正是為什麼單張 80 GB 加速器無法訓練它，也是本軌每一個決策的起點。參見混合精度训练（mixed-precision training），了解為何要同時保留低精度與高精度副本。

M_{\text{train}} \approx \underbrace{2P}_{\text{fp16 weights}} + \underbrace{2P}_{\text{fp16 grads}} + \underbrace{12P}_{\text{Adam fp32 states}} = 16P \text{ bytes}

混合精度 Adam 每个参数约占 16 字节——权重、梯度与优化器状态之和。

資料平行：複製並平均

用 N 張裝置最簡單的方式就是資料平行（data parallelism）：在每張裝置上放一份完整的模型副本，各餵入全域批次的不同切片，反向傳播後把所有人的梯度平均，讓每份副本走出相同的一步。在數學上你計算的就是一次大批次的隨機梯度下降更新，只是把每筆樣本的梯度分散到不同機器上。這是分散式訓練的骨幹，即使是最大型的訓練也以它作為預設起點。

平均這一步是一種叫 all-reduce 的集體通訊（collective communication）：每張裝置最後都持有加總後的梯度。若天真地做（把所有東西送到單一主節點），那個主節點就成了瓶頸，流量隨 N 增長。解法是環狀全歸約，它把裝置排成一個邏輯環，讓梯度分塊沿環傳遞，使每條連結承載的量恆定、與裝置數無關——達到頻寬最佳，也正是資料平行能擴展到數千張 GPU 的原因。

\bar{g} = \frac{1}{N}\sum_{i=1}^{N} g_i

all-reduce 先把每个副本的梯度相加，再除以 N 得到平均后的更新。

for micro in range(accum_steps):
    loss = model(next_batch()) / accum_steps
    loss.backward()              # grads accumulate in .grad
all_reduce(model.grads, op=SUM)  # ring all-reduce across N ranks
model.grads /= world_size        # turn the sum into an average
optimizer.step(); optimizer.zero_grad()

一次資料平行步驟：累積、all-reduce、平均、更新。

不增加記憶體也能放大批次

你想要的批次（為了穩定的優化）往往比實際放得下的批次更大。梯度累積解決了這點：依序跑數個微批次（micro-batches），讓它們的梯度堆疊在 `.grad` 緩衝區裡，直到最後一個才呼叫優化器。有效批次變成 `micro_batch × accum_steps × world_size`，而峰值激活記憶體仍停在單一微批次的大小。理想的目標批次大小並非隨意——梯度雜訊尺度會告訴你，在多大的批次之前，多加樣本仍能改善梯度估計、之後便不再有用。

B_{\text{eff}} = N \cdot K \cdot b_{\mu}

在 N 个副本上累积 K 个微批次，可得到很大的有效批量而无需额外显存。

把網路藏在數學背後

每一次 all-reduce 都在網路線上耗時，若 GPU 閒著等它，你昂貴的叢集大半只是在給機房供暖。決定性的優化是通訊與計算重疊：在最後幾層的梯度一算出來的瞬間就開始歸約，而反向傳播此時還在較前面的層努力。由於反向傳播是一層一層產生梯度，幾乎整個 all-reduce 都能藏在本來就要發生的計算之下。

把參數分桶，讓數層的梯度合成一次 all-reduce，攤平每次呼叫的啟動開銷。
用反向掛勾（backward hook）在每個桶一裝滿時就觸發它的 all-reduce。
在獨立的 CUDA 串流上執行歸約，使它與進行中的反向計算重疊。
只在優化器步驟之前同步一次，此時每個桶都必須完成。

U = \frac{\text{achieved tokens/s}}{\text{peak tokens/s}}

粗略诊断：实测吞吐量除以理论峰值，衡量重叠隐藏通信的效果。