通过融合反向传播与优化器步骤,显著降低大模型训练内存占用。
FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

- 将优化器更新嵌入反向传播,逐块在寄存器中完成,避免梯度显式存储。
- 训练内存减少超50%,小批量下速度提升约1.5倍,8B模型微调时可扩容4倍批处理量。
- 适用于任意线性层与元素级优化规则,不依赖特定架构或优化器。
反向模式微分需先计算所有权重梯度并写入内存,再由优化器读取,此两阶段流程导致训练内存上限:在阶段交界处,所有层的梯度同时驻留内存。我们指出这种显式梯度是计算调度的副产品,而非学习必需;提出FORGE方法,将优化器步骤融入反向传播,逐块在寄存器中执行,使每个梯度块生成即被消耗,永不成为张量。该融合仅改变更新时机,不改变更新内容:全精度下融合步骤在数学上等价于原优化器更新,且在张量与序列并行下仍保持精确;在bf16与8位实际场景中,结果虽非比特完全一致,但偏差有界,权重存储通过随机舍入实现无偏。因梯度块始终在寄存器内生成与消费,无需转为bf16存储与重载,故保留了全精度保真度,而标准bf16/8-bit优化器会因转换损失精度。该方法不绑定特定架构或优化器,适用于任意线性层与元素级规则。实测表明,FORGE将优化器步骤内存减超50%,小批量下提速约1.5倍;集成至张量并行Megatron-LM后,相同GPU下可支持4倍于常规优化器的微批大小,实现8B模型训练。
原文摘要 · Abstract (English)
Reverse-mode differentiation computes every weight gradient, writes it to memory, and only then lets the optimizer read it back. This two-phase schedule sets the memory ceiling of modern training: at the seam between the phases, every layer's gradient is live at once. We argue that this materialized gradient is an artifact of how differentiation is staged, not a quantity that learning requires -- and we eliminate it. FORGE folds the optimizer step into the backward pass and applies it one tile at a time, entirely in registers, so each gradient tile is consumed the instant it is produced and never becomes a tensor. The fusion changes only when the update happens, not what it computes: in full precision the fused step is provably exact -- the identical optimizer update, for every element-wise rule -- and that exactness survives tensor- and sequence-parallel sharding; in the bf16 and 8-bit regimes used in practice it is faithful rather than bit-identical, its deviation bounded and, for the weight store, rendered unbiased by stochastic rounding. Because each gradient tile is born and consumed in the same registers, it is never converted down to bf16 to be stored and read back; FORGE thus preserves the full-precision fidelity that both bf16 and 8-bit optimizers lose to that conversion. Nor is the method tied to one architecture or one optimizer: linear layers are ubiquitous, and FORGE reclaims the gradient memory of any of them under any element-wise rule. Empirically FORGE more than halves the memory of an optimizer step and, at the small batch sizes typical of fine-tuning and continued pretraining, runs about 1.5x faster; integrated into tensor-parallel Megatron-LM it fits 8B training at four times the micro-batch a standard optimizer allows on the same GPUs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。