大模型训练中的梯度累积深度解析：从数学原理到 Megatron-LM 流水线并行的工程实现

引言：为什么梯度累积是大模型训练的”隐形基础设施”

在 ChatGPT 横空出世的三年后，大模型训练的”显学”似乎已经从单卡微调转移到了千卡预训练。各类技术文章铺天盖地讨论着 3D 并行、ZeRO 优化、通信重叠，但有一个看似平淡无奇的技术细节却在每一次迭代中默默运转——梯度累积（Gradient Accumulation）。

梯度累积的本质极其简单：把一个大 batch 拆成若干 micro-batch，分别做前向和反向，把梯度攒起来，最后一步再更新参数。这个操作在 PyTorch 中不过是

之后不加

，等攒够了再调。但就是这么一个”一行代码”的逻辑，当它与流水线并行、序列并行、ZeRO 交织在一起时，工程复杂度会急剧上升。

本文将从梯度累积的数学动机出发，逐步展开到 Megatron-LM 中 1F1B 调度策略的实现细节，再到与 DeepSpeed ZeRO 的交互，最后讨论梯度累积对训练稳定性的隐性影响。这不是一篇入门教程，而是给已经了解并行训练基本概念的工程师准备的”下一层”深度分析。

数学动机：显存约束下的等效大 Batch

从 Batch Size 与泛化的关系说起

大模型训练对 batch size 的需求不是任意的。OpenAI 的 GPT-3 训练使用了 3.2M tokens 的 batch size，Meta 的 LLaMA 系列也保持在 4M tokens 的量级。这并非随意选择——过小的 batch size 会导致梯度噪声过大，训练不稳定；过大的 batch size 则会降低泛化性能（所谓”generalization gap”问题）。

但问题在于：一块 A100-80G 能放下的激活值是有限的。以 70B 模型为例，单卡的模型参数（使用 ZeRO-3 分片后）加上优化器状态就已经捉襟见肘，一个 4M tokens 的 batch 根本不可能一次塞进显存。

梯度累积的数学等价性

设全局 batch size 为

，拆分为

个 micro-batch，每个 micro-batch 大小为

。第

个 micro-batch 的损失为

，则：

1
2
3
4
5
6
7
8
9
10
11
12
# 不使用梯度累积（单步大 batch）

L = (1/N) * sum(L_i)

dL/dtheta = (1/N) * sum(dL_i/dtheta)



# 使用梯度累积（N 步小 batch）

# 步骤 1~N-1：仅 backward，不 step

for i in range(N):

    loss_i = forward_backward(micro_batch_i)

    # 梯度自动累加到 .grad

# 步骤 N：执行 step

optimizer.step()  # 使用累加后的梯度

optimizer.zero_grad()

关键等价条件：当损失函数是各样本损失的算术平均时，累加 N 个 micro-batch 的梯度等价于一次大 batch 的梯度除以 N。这意味着梯度累积在数学上严格等价于使用全局 batch size 进行训练——前提是你正确处理了 BN 统计量（BatchNorm 的 running mean/var）和学习率缩放。

何时等价性会破裂？

梯度累积的等价性并非无条件的。以下场景会导致实际行为偏离”理想大 batch”：

• BatchNorm 层：每个 micro-batch 独立计算均值和方差，而非使用全局统计量。大模型几乎都用 LayerNorm，所以这不是主要问题，但如果你在视觉模型上做迁移学习就需要注意。
• Dropout：不同 micro-batch 会产生不同的 dropout mask，但这种随机性差异实际上可忽略——它只是增加了梯度估计的方差。
• 序列并行中的激活重计算：当使用 selective 重计算策略时，不同 micro-batch 的重计算粒度可能不同，导致浮点误差的微小差异。

与流水线并行的深度耦合：1F1B 调度策略

梯度累积在单卡场景下只是”多跑几次 backward”，但在流水线并行（Pipeline Parallelism, PP）中，它决定了整个调度策略的设计。Megatron-LM 提出了 1F1B（One Forward One Backward） 调度，这是目前几乎所有主流训练框架的默认选择。

为什么朴素调度不可行？

最直观的调度是”全前向再全反向”（All-Forward-All-Backward）：先把 N 个 micro-batch 全部跑完前向，再全部跑反向。问题在于：峰值显存。每个 micro-batch 的前向激活值都需要保存到反向传播时使用，N 个 micro-batch 的激活值同时驻留在显存中，峰值激活显存是单 micro-batch 的 N 倍。

以 4 级流水线、8 个 micro-batch 为例，朴素调度下 Stage 0 需要保存 8 份激活值。对于 70B 模型的一级流水线阶段，单份激活可能就达到 2-4 GB，8 份就是 16-32 GB——直接撑爆显存。

1F1B 调度的精妙设计

1F1B 的核心思想是：在流水线”预热”阶段，逐个执行前向传播直到流水线填满；然后进入稳态，每执行一个前向就紧跟一个反向，前向和反向交替进行；最后是”冷却”阶段，处理剩余的反向传播。

1
2
3
4
5
6
7
8
# 4 级流水线，8 个 micro-batch 的 1F1B 调度

# 时间 ->

# Stage 0: F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7

# Stage 1:    F0 F1 F2 B0 F3 B1 F4 B2 F5 B3 F6 B4 F7 B5 B6 B7

# Stage 2:       F0 F1 B0 F2 B1 F3 B2 F4 B3 F5 B4 F6 B5 F7 B6 B7

# Stage 3:          F0 B0 F1 B1 F2 B2 F3 B3 F4 B4 F5 B5 F6 B6 F7 B7

#

# F = Forward, B = Backward, 下标 = micro-batch 编号

1F1B 的显存优势来自一个简单的事实：完成反向传播后，对应的激活值即可释放。稳态阶段，前向产生一份新激活的同时释放一份旧激活，峰值激活数从 N 降至 PP 大小（流水线级数），通常远小于 N。

Megatron-LM 中的实现细节

在 Megatron-LM 的源码中，1F1B 调度由

中的

函数实现。几个关键细节：

1. 通信与计算的 Overlap：每个 stage 在前向完成后，需要通过 point-to-point 通信把激活值传给下一个 stage。Megatron 使用 NCCL 的 send/recv 操作，并通过

与计算流重叠。但注意——1F1B 调度中的通信并不能完全被计算覆盖，因为 stage 之间有严格的时序依赖。

2. 虚拟流水线阶段（Interleaved Pipeline）：Megatron-LM 提出了一种”交错”策略——每个物理 GPU 负责多个不连续的虚拟 stage。例如 4 卡 8 级流水线中，GPU 0 负责 Stage 0 和 Stage 4，GPU 1 负责 Stage 1 和 Stage 5，以此类推。这减少了流水线气泡，但也让梯度累积的调度更加复杂——同一个 GPU 上不同虚拟 stage 的梯度需要分别累积。

3. 梯度同步时机：在数据并行（DP）组中，All-Reduce 必须在所有 micro-batch 的梯度累积完成后才能执行。Megatron-LM 在

回调中处理这一逻辑——默认是立即 All-Reduce，但也可以延迟到 checkpoint 边界再同步，以进一步减少通信次数。

梯度累积与 ZeRO 优化的交互

DeepSpeed 的 ZeRO（Zero Redundancy Optimizer）通过分片优化器状态、梯度和参数来降低显存开销。但梯度累积在 ZeRO 下需要额外的工程处理。

ZeRO Stage 1/2 下的梯度累积

ZeRO-1 分片优化器状态，ZeRO-2 额外分片梯度。在 ZeRO-2 中，每个 GPU 只保存自己负责的那部分梯度分片。梯度累积期间，每个 micro-batch 的反向传播只产生局部梯度，这些梯度需要先 Reduce-Scatter 到对应的分片 owner，然后累加——而不是先在本地累加再一次性 Reduce-Scatter。

DeepSpeed 的实现选择了一种折中方案：

• 连续累加模式：每个 micro-batch 的反向后，先在本地梯度缓冲区累加，等所有 micro-batch 完成后再执行 Reduce-Scatter。这节省了通信次数，但需要在本地保存完整的梯度缓冲区。
• 即时通信模式：每个 micro-batch 反向后立即执行 Reduce-Scatter，只保留本地分片。通信量增大但峰值显存更低。

DeepSpeed 默认使用连续累加模式，因为它对吞吐量更友好。但在极端显存受限的场景下，你可能需要手动切换到即时通信模式：

1
2
3
4
5
6
7
8
9
10
# DeepSpeed 配置中控制梯度累积的通信策略

{

  "gradient_accumulation_steps": 16,

  "gradient_clipping": 1.0,

  "zero_optimization": {

    "stage": 2,

    "overlap_comm": true,        # 通信与计算重叠

    "contiguous_gradients": true  # 使用连续梯度缓冲区

  }

}

ZeRO-3 下的特殊考量

ZeRO-3 进一步分片模型参数，这意味着每次前向/反向都需要 All-Gather 参数。在梯度累积场景下，参数的 All-Gather 在每个 micro-batch 都要执行一次，而不是只执行一次。这会带来 N 倍的参数通信开销。

缓解策略是使用 预取（Prefetch）：在当前 micro-batch 的计算过程中，提前 All-Gather 下一个 micro-batch 需要的参数。Megatron-DeepSpeed 的组合实现中，这一优化可以使 ZeRO-3 下梯度累积的额外通信开销降低约 40-60%。

梯度累积对训练稳定性的隐性影响

梯度累积并非没有代价。除了显存和通信的显性开销外，它还对训练稳定性有着容易被忽视的影响。

梯度裁剪的时机问题

大模型训练中，梯度裁剪（Gradient Clipping）是防止梯度爆炸的关键手段。标准做法是计算全局梯度范数，如果超过阈值就缩放。但在梯度累积场景下：

1
2
3
4
5
6
7
8
9
10
11
12
13
# 错误做法：每个 micro-batch 独立裁剪

for i in range(N):

    loss_i = forward_backward(micro_batch_i)

    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # 每次裁剪都基于部分梯度，不等价于全局裁剪



# 正确做法：累积完成后统一裁剪

for i in range(N):

    loss_i = forward_backward(micro_batch_i)

# 累积完成后才裁剪

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

optimizer.step()

optimizer.zero_grad()

在分布式训练中，全局梯度范数的计算本身也需要 All-Reduce（跨 DP 组）。Megatron-LM 在

中维护了一个持久的梯度缓冲区，累积期间梯度一直驻留其中，直到所有 micro-batch 完成后才计算范数并裁剪。

混合精度下的精度损失

在 FP16/BF16 混合精度训练中，梯度的累积精度直接影响训练质量。PyTorch 默认在 FP32 的

缓冲区中累加梯度（通过

管理），但有两个陷阱：

• 梯度缩放（Loss Scaling）：动态 Loss Scaling 在梯度累积期间可能调整缩放因子。如果某个 micro-batch 触发了 inf/nan 导致缩放因子减半，后续 micro-batch 的梯度就会基于新的缩放因子累加——这打破了等价性。解决方案是在梯度累积期间锁定缩放因子，只在 step 后检查并调整。
• BF16 的累加精度：BF16 只有 7 位尾数，远少于 FP16 的 10 位。虽然累加在 FP32 中进行，但如果你的自定义层直接操作 BF16 梯度（绕过了 PyTorch 的自动精度提升），累积误差可能不可忽略。

实战：调优梯度累积步数的工程决策

理论分析完毕，回到实际工程问题：梯度累积步数（N）该怎么选？

决策框架

选择 N 本质上是在显存、吞吐量和训练稳定性之间做权衡。以下是系统化的决策流程：

第一步：确定目标全局 Batch Size。这通常由训练配方决定，而非工程选择。例如 LLaMA-2 70B 使用 4M tokens 的 batch size。

第二步：计算理论 micro-batch 大小。

1
2
3
4
5
6
7
8
9
10
11
12
# 计算公式

global_batch_tokens = 4_000_000  # 4M tokens

seq_len = 4096

dp_world_size = 64  # 64 路数据并行



# 每个 DP 副本的 batch tokens

tokens_per_dp = global_batch_tokens / dp_world_size  # 62,500



# 每个 micro-batch 的样本数

micro_batch_size = tokens_per_dp / seq_len / N  # N 是梯度累积步数



# 需要让 micro_batch_size 是正整数且能放进显存

第三步：通过显存 profiling 确定最大 micro-batch。在目标模型和硬件上，逐步增大 micro-batch 直到 OOM，取最大可行值。然后反推 N：

1
N = ceil(tokens_per_dp / (max_micro_batch * seq_len))

第四步：检查 N 对吞吐量的影响。N 越大，吞吐量通常越高（因为并行度更灵活），但流水线气泡比例也越高。经验法则：N 至少应是 PP 大小的 2-4 倍，才能让 1F1B 调度的气泡率低于 15%。

常见配置参考

动态梯度累积：一种进阶策略

固定 N 并非唯一选择。在训练的不同阶段，显存压力是变化的——早期 checkpoint 的激活值较小，后期随着模型收敛和优化器状态的膨胀，显存余量在减少。动态梯度累积允许根据当前显存状态自适应调整 N：

1
2
3
4
5
6
7
8
# 伪代码：动态梯度累积

current_free_mem = torch.cuda.mem_get_info()[0]  # 可用显存

if current_free_mem > 20 * 1024**3:  # > 20GB 空闲

    N_dynamic = max(N // 2, 1)  # 减少累积步数，增大 micro-batch

    micro_batch_size_dynamic = tokens_per_dp // seq_len // N_dynamic

else:

    N_dynamic = N

    micro_batch_size_dynamic = base_micro_batch_size

这种策略在 Megatron-LM 中并未原生支持，但一些团队（如 01.AI、DeepSeek）在内部训练框架中实现了类似逻辑。核心挑战是保证训练的全局 batch size 在动态调整下仍然保持一致——否则优化器的动量项和方差项会基于不同 batch size 计算，导致训练不稳定。

梯度累积与序列并行的协作

序列并行（Sequence Parallelism, SP）是 Megatron-LM 在 TP 组内对 LayerNorm/Dropout 等非张量并行算子做的一种分片优化。它把序列维度切到多个 TP rank 上，每个 rank 只计算一部分 token。当 SP 遇上梯度累积时，有几个有趣的交互：

激活重计算与累积步数的联合优化

SP 的核心收益是减少激活显存——每个 rank 只保存 1/TP 的激活。但激活重计算（Activation Checkpointing / Recomputation）可以进一步降低显存，代价是额外的前向计算。当使用梯度累积时，激活显存与 N 成正比（1F1B 调度下与 PP 成正比），因此：

1
2
3
4
5
6
7
8
9
10
11
# 峰值激活显存估算

peak_activation = (

    base_activation_per_micro_batch

    * pipeline_stages  # 1F1B 调度

    / tensor_parallel_size  # 序列并行分片

    * (1 - recompute_ratio)  # 重计算比例

)



# 总显存约束

total_mem = model_params + optimizer_states + peak_activation + gradient_buffer

assert total_mem < gpu_memory

在工程实践中，recompute_ratio 的选择与 N 是耦合的：如果 N 较大（意味着更多 micro-batch），激活峰值更高，需要更高的 recompute_ratio 来补偿；如果 N 较小，可以降低 recompute_ratio 以提升吞吐量。

SP 组内的梯度 All-Reduce

在 SP 中，LayerNorm 和 Dropout 的梯度是跨 TP rank 分片的，需要在反向传播时执行 All-Reduce 来聚合。在梯度累积场景下，这些 All-Reduce 在每个 micro-batch 都会触发——即使梯度尚未最终累积完成。Megatron-LM 使用了一种延迟 All-Reduce 策略：在累积期间，SP 组内的梯度先在本地缓冲，累积完成后统一 Reduce。但这要求额外的缓冲区显存，约等于一个完整梯度张量的大小。

面试高频问题：梯度累积的深层理解

以下是从各大 AI 公司面试中提炼出的梯度累积相关问题，以及深入解析：

Q1：梯度累积是否等价于使用更大 batch size 训练？给出精确的等价条件。

等价条件是：(1) 损失函数是各样本损失的算术平均；(2) 梯度裁剪在累积完成后统一执行；(3) 没有依赖 batch 统计量的层（如 BatchNorm）；(4) 混合精度训练中 Loss Scaling 在累积期间保持不变；(5) 不使用依赖 batch 内样本交互的损失函数（如对比学习中同一 batch 内的正负样本构造）。

Q2：1F1B 调度中，如果某个 stage 成为计算瓶颈，梯度累积会如何受影响？

在同步流水线中，每个 step 的耗时由最慢的 stage 决定。如果某个 stage 计算量偏大（通常是最深的层或 attention 密集的层），1F1B 调度不会改变这一点——它只是减少了峰值显存。但值得注意的是：计算不均衡会放大流水线气泡。理想情况下每个 stage 耗时相同，气泡率为 (PP-1)/(PP-1+N)。如果某个 stage 慢了 30%，有效气泡率会显著增加，因为其他 stage 在等待期间处于空闲状态。

解决方案包括：(1) 调整流水线切分点使各 stage 计算量均衡；(2) 使用 interleaved pipeline 减少气泡；(3) 在瓶颈 stage 上使用更激进的激活重计算来释放显存、增大 micro-batch size。

Q3：在 MOE 模型中，梯度累积有什么特殊考量？

MOE 模型的核心特点是 token 路由的动态性。每个 micro-batch 的 token 被路由到不同的专家，但不同 micro-batch 的路由结果不同——这意味着每个专家在不同 micro-batch 中收到的 token 数量波动很大。在梯度累积场景下，这种波动会导致：(1) 某些专家的梯度基于极少量样本，估计方差大；(2) All-to-All 通信的负载不均衡在累积期间被放大。实践中，MOE 训练通常使用更大的 N（32-64）来平滑路由波动带来的梯度噪声。

总结与展望

梯度累积是大模型训练中最不起眼却最不可或缺的工程组件之一。它的数学原理简单到可以用一行公式概括，但在与流水线并行、ZeRO 优化、序列并行、混合精度训练的交互中，衍生出了丰富的工程细节：

• 1F1B 调度将梯度累积与流水线气泡控制深度融合，是 Megatron-LM 训练效率的基石
• ZeRO-3 下的参数 All-Gather 在每个 micro-batch 都要执行，预取是关键优化
• 梯度裁剪时机、Loss Scaling 锁定、BF16 精度问题是训练稳定性的隐性守护
• MOE 路由波动要求更大的累积步数来平滑梯度噪声

展望未来，随着上下文长度从 128K 向 1M+ 演进，序列维度的切分将更加重要，梯度累积与 Ring Attention、Ulysses 等序列并行方案的交互还有大量工程优化空间。同时，异步梯度累积（不同 micro-batch 使用不同版本的参数进行前向计算）也是一个有潜力但尚未成熟的方向——它打破了严格的数学等价性，但可能带来显著的吞吐量提升。

理解梯度累积的这些深层细节，不仅能帮助你在训练集群上获得更好的性能，更能让你在面对”训练出了 NaN”、”显存 OOM”、”多卡 loss 不一致”等 debug 场景时，快速定位到根因。在大模型训练的战场上，魔鬼永远在细节里。