大規模訓練
梯度累積(gradient accumulation)
大批次訓練常因穩定性與吞吐而受青睞,但你想要的批次未必能一次塞進記憶體。梯度累積把統計上的批次與實體上的批次解耦。你把想要的有效批次切成數個微批次,對每個各跑一次前向與反向,把得到的梯度加進一個緩衝區,先不執行優化器步。直到最後一個微批次處理完——此時累積的梯度等於一個大批次本會產生的梯度——才套用單一一次更新。
由於每個微批次都在處理完後、下一個之前就被釋放,記憶體峰值由微批次大小決定,而非有效批次大小,因此你能在固定硬體上達到任意大的有效批次,代價是按比例增加的循序步數。微妙之處在平均的正確性:損失必須被縮放(或梯度除以累積次數),使累積梯度是平均而非總和;而在資料平行訓練裡,除了最後一個微批次外,每個微批次都應抑制 all-reduce,以免過早地對部分梯度做平均。
梯度累積也是填滿管線的機制:它產生的微批次,正是流過管線階段以縮小氣泡的那些單元。它在記憶體上幾乎免費,在時間上卻不然——累積步數越多,每次更新就要做越多次前向/反向,因此它提升的是有效批次大小,而非硬體利用率。
g = \frac{1}{K}\sum_{k=1}^{K} \nabla_\theta \mathcal{L}(B_k),\quad \theta \leftarrow \theta - \eta\, g
在單次更新前對 K 個微批次梯度取平均,重現了大小為 |B_k|·K 的單一批次之梯度。
梯度累積只在「不含批次統計層」時才等價於真正的大批次:像批次正規化這類元件是逐微批次計算統計量的,因此無法重現大批次行為,不過 transformer 的 layer norm 不受影響。
又称
另见