梯度累积与 Global Batch：小步拼大更新，账要算在 loss 和同步点上

读完这节应知道：global batch 是优化与统计单位，micro-batch 是显存与执行单位；两者之间隔着 DP 和累积步数。

训练一步要存的不只是参数，还有前向激活、反向中间量和梯度。对固定模型，激活峰值随一次前向的样本数增长；参数、梯度和 optimizer state 则与 micro-batch 大小基本无关。于是「一步看多少样本」和「一次放进 GPU 多少样本」必须拆成两个变量。

Megatron 的训练配置文档把它写成 “local batch size times data parallel size times number of micro batches”（megatron/training/config/training_config.py:13-14）。

三个乘数买的东西不同。micro_batch_size 决定单次前向的激活峰值，也决定 BN、dropout 等按 micro-batch 统计的层看到多少样本；data_parallel_size 决定模型副本数和集合通信宽度；gradient_accumulation_steps 决定一个 optimizer step 前串行跑多少个 micro-batch。把 GBS 固定在 128，可以让 MBS=4、DP=8、GA=4，也可以让 MBS=8、DP=16、GA=1；样本数相同，显存、通信次数和每卡的处理节奏都不同。

因此「把 batch 调大」不是单一动作。减小 batch 会改变梯度噪声和优化步数；增加 DP 需要更多副本和带宽；增加 GA 则把并行时间换成串行时间。三者的账必须在同一个 GBS 下对齐，否则训练曲线和吞吐都不可比。

读完这节应知道：三种方案都能让模型「跑起来」，但它们改变的是不同资源，不能互相替代。

表 1 · 三种方案改的是不同资源（定性对照，不含实测数字）

一个常见混淆是把「PP 的 microbatch」当成第二个 gradient accumulation。Megatron 的 pipeline schedule 收到的 num_microbatches 就是 get_num_microbatches()（megatron/training/training.py:3574-3579；scheduled_num_microbatches 由 3460 赋值），它同时是 GA 计数；PP 只改变这些 micro-batch 在流水线中的调度方式，不额外乘一个 PP 因子。

读完这节应知道：no_sync 省掉的是每个 micro-batch 的集合通信，不是梯度显存；loss 的除法位置决定变长序列的加权。

3.1 先 backward，再 step：梯度在 .grad 里相加

loss.backward() 会把当前计算图的梯度累加到参数的 .grad 上，而不是覆盖。PyTorch 文档对 no_sync 的描述是 “gradients will be accumulated on module variables”（torch/nn/parallel/distributed.py:1700-1702）。因此不需要 losses += loss；需要的是别在 micro-batch 之间 zero_grad()。

Megatron 在每个训练迭代开始时清零一次梯度缓冲区，然后跑 num_microbatches 次前向/反向，最后只调用一次 optimizer.step()（megatron/training/training.py:3473-3478、3574-3580、3662）。DeepSpeed 在累积边界上用 (micro_steps + 1) % gradient_accumulation_steps == 0 判断是否更新（deepspeed/runtime/engine.py:3494），ZeRO-3 则在 micro_step_id == 0 时拷贝、之后 add_ 到 fp32 梯度分区（deepspeed/runtime/zero/stage3.py:2041-2047）。

这段示意只覆盖「等长 micro-batch + 每卡同 token 数」的情形；变长序列或 padding 不同时，loss / gas 不是 token-level 的正确缩放，见 §3.3。

3.2 DDP/FSDP 的 no_sync：把 all-reduce 推到边界

PyTorch DDP 的 no_sync 把 require_backward_grad_sync 置为 False，退出时恢复。文档明确：梯度先累积在模块变量上，在退出上下文后的第一次前向/反向同步；并且 forward 必须放在上下文内，否则梯度仍会同步（torch/nn/parallel/distributed.py:1696-1722）。DDP 默认的通信 hook 是「先除以组大小，再 all-reduce」（torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py:19-32），所以 DDP 的同步是跨 rank 的平均。

FSDP 的 no_sync 只能用在 root 实例上，并递归作用到子实例；它会在累积期间保留完整模型梯度而不是梯度分片，因此显存通常更高；使用 CPU offload 时，梯度也要等到最终同步后才下 CPU（torch/distributed/fsdp/fully_sharded_data_parallel.py:1044-1081）。FULL_SHARD 平时在反向结束后做 reduce-scatter；no_sync 期间不做这一步。

DeepSpeed 的 no_sync 由 inside_no_sync_ctxt 标记；Engine.step() 里直接断言不能在该上下文内调用（deepspeed/runtime/engine.py:3606-3608）。ZeRO stage ≥ 2 在反向期间重叠梯度分区，stage 0/1 只在累积边界做 allreduce_gradients（deepspeed/runtime/engine.py:3094-3112）。

3.3 loss 归约：三次平均，还是按 token 加权

Megatron 的 GPT loss function 返回的是当前 micro-batch 的 token loss 之和，同时返回非 padding token 数：loss = torch.sum(losses * loss_mask)，num_tokens = loss_mask.sum()（pretrain_gpt.py:288-293）。它不是先做 mean 再返回。

默认路径（calculate_per_token_loss=False）在 schedule 里做两次除法：先 output_tensor /= clamp(num_tokens, min=1)，再 output_tensor /= num_microbatches（megatron/core/pipeline_parallel/schedules.py:343-349）。随后 DDP 在 DP 组内做平均。合起来是三层平均：micro-batch 内按 token 平均 → micro-batch 之间等权平均 → DP rank 之间等权平均。如果每个 micro-batch 的 token 数相同，这与全局 token-level mean 相同；不同则不等。

逐 token 路径（calculate_per_token_loss=True）保留 loss 的和，在 finalize_model_grads 里把最后一个 PP stage 的 num_tokens 广播到 PP 组、在 DP(+CP) 组内 all-reduce（dp_cp_group），然后用 1 / total_global_tokens 缩放所有梯度（megatron/core/distributed/finalize_model_grads.py:805-822）。这条路径下 DDP 不做 1/dp 预缩放（megatron/core/distributed/distributed_data_parallel.py:268-273），分母只出现一次。

L_m 是第 m 个 micro-batch 的 token loss 之和，n_m 是它的非 padding token 数。两式相等当且仅当所有 n_m 相同（或权重相同）；变长序列、padding 比例不同、packing 后各 micro-batch 长度不同，都会让它们分开。

这不是数值误差，而是梯度权重的差异：mean-of-means 给每个 micro-batch 相同权重，短 micro-batch 的每个 token 被相对放大；token-level 给每个 token 相同权重。Megatron 的配置注释也说明，默认行为是「假设所有 token 都非 padding」，逐 token 路径才按 global batch 的实际非 padding token 数计算（megatron/core/transformer/transformer_config.py:324-326）。

3.4 clip 和 optimizer step 放在哪里

梯度裁剪必须在完整累积、DP 归约之后、参数更新之前。Megatron 的 optimizer.step() 依次做 prepare_grads()、clip_grad_norm(config.clip_grad)、step_with_ready_grads()（megatron/core/optimizer/optimizer.py:1078-1109）；裁剪用的 grad norm 在 grad_stats_parallel_group 上归约（megatron/core/optimizer/optimizer.py:545-551），看到的是全局梯度，不是某个 micro-batch 的局部梯度。

DeepSpeed ZeRO-3 同样在 step() 里按子组做 unscale_and_clip_grads，再调用 optimizer（deepspeed/runtime/zero/stage3.py:2913-2960、3019-3028）。如果在 micro-batch 循环里裁剪，你裁剪的是尚未累加的偏梯度，算法已经变了；这不是「同一个 clip 换个位置」。

zero_grad() 放在 optimizer step 之后、下一个累积循环之前。Megatron 在每次迭代开头清一次（megatron/training/training.py:3473-3478）；DeepSpeed 在边界 step 内部/之后处理。不要在 micro-batch 之间清零。

读完这节应知道：GBS 由 MBS、独立数据分区数和 GA 相乘；lr 的线性缩放只在论文列出的条件下成立，不能从 GBS 直接推。

4.1 账本与整除约束

TGB 是固定长度、无 padding、无 packing 时的 token 数；真实训练要按非 padding token 统计。

Megatron 在初始化时断言 global_batch_size % (micro_batch_size × data_parallel_size) == 0，并计算 num_micro_batches = global_batch_size // (micro_batch_size × data_parallel_size)（megatron/core/num_microbatches_calculator.py:381-388）。如果 GBS 没显式给，默认取 micro_batch_size × batch_parallel_size（megatron/training/arguments.py:641-642），也就是 num_microbatches=1。

这里的 DP 是「独立数据分区数」。Megatron 把 get_batch_parallel_size(args) 传给 num-microbatches calculator（megatron/training/global_vars.py:225-234），而 batch_parallel_size = world_size / (TP × PP × CP)（megatron/training/arguments.py:494-499）。TP/PP/CP 不乘 GBS；PP 的 microbatch 计数就是 GA，见 §07。

表 2 · 四组账本示例（按 GBS = MBS × DP × GA 计算；无 drop_last、无 packing、固定 seq_len）

「梯度同步 / step」列写 1 是指 no_sync / 延迟归约把 DP 集合通信放在累积边界；如果每个 micro-batch 都同步，次数是 GA。A 组 DP=1，没有跨 rank 通信，所以实际同步 0 次。两种做法的样本账相同，通信账不同。表里的 1 个 optimizer step 是定义，不是实测。

4.2 学习率缩放：linear scaling rule 只在它的条件下成立

Goyal et al. §2.1 的规则是：minibatch size 乘以 k，learning rate 也乘以 k。§5.1 给出参考学习率 η = 0.1 · kn/256，baseline 是 k=8 GPUs × n=32 images/worker，即 kn=256、η=0.1。条件必须一起读：其他超参数保持不变，包括 weight decay（论文用 1e-4，且不作用于 BN 的 γ/β）和 BN 的 per-worker batch n=32；规则假设小 batch 步上的梯度变化不大，在训练早期会失效；§5.3 的 batch sweep 显示 8k 之后 error 上升。

Warmup 不是可选装饰。Goyal §2.2 给出 warmup 方法，§5.1 给出取值：从 η=0.1 在 5 epochs 内线性升到 3.2；constant warmup（先 0.1 保持 5 epochs 再跳）会在跳变后出现 error spike，训练无法完全恢复。Goyal Table 1 的三种 8k 设置里，gradual warmup 的 top-1 error 是 23.74，no warmup 是 24.84（差 1.10 个百分点），constant warmup 是 25.88（比 gradual 差 2.14 个百分点，甚至比 no warmup 还高 1.04 个百分点）。这些数字的条件是 ResNet-50 / ImageNet / 90 epochs / SGD Nesterov m=0.9 / 5 次独立运行。

Smith et al. 给出另一条路径：不衰减 lr，而是在 decay 点增大 batch。§5.2 的 WRN 16-4 / CIFAR-10 实验里，SGD m=0.9、ghost BN 128；原始 schedule 是 lr 0.1、batch 128、约 80,000 次更新、94.3% 测试准确率；把 batch 从 128 起每个衰减点 ×5（上限 5120）、lr 保持 0.1，约 29,000 次更新、94.4%（单次；五次追加运行的中位数 94.2%）。Smith §5.1 的总结句是 “it is the noise scale which is relevant, not the learning rate.” 这不是「batch 越大越好」：同一节里把 momentum 提到 0.98、初始 batch 3200 的设置只到 93.3%（中位数 93.5%），更新少于 2,500 次。

Megatron SC21 报告的是系统层面的 batch 选择：§5.1 的 175B 训练时间估算按 1024 张 A100、300B tokens、B=1536、140 TFLOP/s/GPU 计算；§5.4.3 明确 GPT-3 训练收敛时用的就是 B=1536。530B 对比在 560 GPUs 上放不下时，Megatron Table 2 caption 说明这是 ZeRO-3、microbatch size 4 的条件，作者把 GPU 数提到 640、B 提到 2560。Megatron §5.1 固定 seq_len=2048，所以 1536×2048 = 3,145,728 token/batch。SC21 报告的是 B 的取值；linear scaling rule 的表述来自 Goyal，不来自 SC21。

一个容易忽略的边界：GBS 相同不代表更新等价。MBS 改变 BN/dropout 的 per-step 统计，GA 改变串行顺序，DP 改变 rank 间的平均；即使 loss 归约写成 token-level，数据顺序、dropout mask 和 optimizer 的动量轨迹也不会与「一次大 batch」逐位相同。只有在固定模型、无 BN/dropout 随机性、相同数据顺序且 loss 归约一致的条件下，累积梯度才与大 batch 梯度相等。

读完这节应知道：同步、裁剪、step 各自只发生一次；micro-batch 循环里发生的是前向、反向和梯度累加。

把顺序拆成可核对的五步：

• 清零一次：迭代开始时 zero_grad_buffer() / optimizer.zero_grad()，之后整个 micro-batch 循环不再清零（megatron/training/training.py:3473-3478）。
• 前向/反向 GA 次：每次取 MBS 条样本；前 GA−1 次在 no_sync 或延迟归约下只累加，最后一次触发 DP 归约。
• 归约一次：DDP 默认 all-reduce（先除以组大小）；FSDP 用 reduce-scatter；Megatron 在 finalize_model_grads 里做 all-grads-sync，并在逐 token 路径上除以全局 token 数（megatron/core/distributed/finalize_model_grads.py:805-822）。
• 裁剪一次：optimizer.step() 内部先 clip_grad_norm，再更新参数（megatron/core/optimizer/optimizer.py:1078-1109）。
• 更新一次：参数更新后进入下一个迭代；zero_grad 在下一轮开头。

PP 的情形：每个 micro-batch 还要穿过 pipeline stages，schedule 里的 num_microbatches 仍是 GA（megatron/training/training.py:3460、3574-3579）。interleaved 1F1B schedule 要求 num_microbatches 是 PP degree 的整数倍（SC21 §2.2.2），这是调度约束，不是账本里多一个乘数。PP 带来的是 bubble，不是额外的 GBS。

读完这节应知道：大 batch 的实测收益依赖 warmup、模型和优化器；论文里的「可行」不等于「无条件更好」。

表 3 · Goyal et al. Table 1：ResNet-50 / ImageNet / 90 epochs / SGD Nesterov m=0.9 / weight decay 1e-4 / 5 次独立运行的均值 ± 标准差（每次运行取最终 5 epochs 的中位数误差）

表 3 的读法：8k + gradual warmup 与 256 baseline 相差 0.14 个百分点（23.74 − 23.60），这是这篇论文条件下「大 batch 可行」的核心证据；但同一张表里 constant warmup 25.88 比 no warmup 24.84 还高 1.04 个百分点，说明 warmup 的形状本身会改变结论。8k 是 256 的 32 倍（8192/256），学习率 3.2 也是 0.1 的 32 倍；这个倍数关系只在这套设定和这些条件下成立。

表 4 · Smith et al. §5.2：WRN 16-4 / CIFAR-10 / SGD / ghost BN 128；更新次数与测试准确率

表 4 的读法：在相同的 epoch 预算下，把 lr 衰减换成 batch 增长，可以把参数更新从约 80,000 次压到约 29,000 次，单次准确率 94.4% 与原始 94.3% 相差 0.1 个百分点（94.4 − 94.3）。但这个替代不是无条件的：同一节里 momentum 0.98、初始 batch 3200 的设置只有 93.3%，说明噪声尺度、momentum 和 batch 上限要一起调。Smith §5.1 的表述是 “it is the noise scale which is relevant, not the learning rate.”

Megatron SC21 的 batch 选择更偏系统账：§5.1 的 175B 训练时间估算按 1024 张 A100、300B tokens、B=1536 计算，§5.4.3 明确 GPT-3 训练收敛时用的就是 B=1536；530B 对比在 560 GPUs 上放不下，Megatron Table 2 caption 说明这是 ZeRO-3、microbatch size 4 的条件，作者把 GPU 数提高到 640、B 提高到 2560。在 §5.1 固定的 seq_len=2048 下，1536×2048 = 3,145,728 token/batch。SC21 报告的是 B 与系统配置；本页不把它读成「更大的 batch 准确率更高」，它说明的是 batch 选择与显存、GPU 数、并行配置绑在一起。

把三组结果放在一起：Goyal 证明的是「在 ResNet-50 / ImageNet 的特定条件下，8k batch + 线性 lr + gradual warmup 可以接近 256 baseline」；Smith 证明的是「在 WRN / CIFAR-10 上，用 batch 增长替代 lr 衰减可以得到相近的 epoch 曲线，同时减少更新次数」；SC21 报告的是「大规模 GPT 训练实际使用的 B」。三者都不能单独外推成「GBS 越大越好」或「lr 必须线性缩放」。

读完这节应知道：梯度累积省的是激活，不是所有显存；它也不会自动修正 BN、loss 权重或流水线调度。

BN 看的是 per-worker batch，不是 global batch

BatchNorm 的均值和方差在每次前向的 micro-batch 上统计。梯度累积把多个 micro-batch 串行跑，但每个 micro-batch 仍然只看到 MBS 条样本，所以 BN 统计不会因为 GA 变大而变成 global batch 的统计。Goyal 在放大 batch 时保持每个 worker 的 n=32 不变，正是为了让 BN 的 loss function 不变（§2.3、§5.1）；Smith 的 CIFAR-10 实验用 ghost BN 128（§5.1、§5.2）。如果模型依赖 BN，想复现大 batch 的行为需要 sync BN、ghost BN 或重新调 BN 的 momentum，而不是只加 GA。

变长序列：默认是 mean-of-means，不是 token-level

默认路径先除以当前 micro-batch 的 token 数，再除以 micro-batch 数（megatron/core/pipeline_parallel/schedules.py:343-349）。当序列长度不同、padding 比例不同或 packing 后各段长度不同，短 micro-batch 会被赋更高的每 token 权重。逐 token 路径把整个 global batch 的非 padding token 数作为分母（megatron/core/distributed/finalize_model_grads.py:805-822），并在 DDP 上关闭 1/dp 预缩放（megatron/core/distributed/distributed_data_parallel.py:268-273）。使用哪条路径要在训练配置里明确写出来；不要用「loss 除以 GA」去代替 token 加权。

梯度裁剪只在累积边界之后

裁剪的 grad norm 必须在完整累积和 DP 归约之后计算。Megatron 的 optimizer.step() 先 prepare_grads、再 clip_grad_norm、最后 step_with_ready_grads（megatron/core/optimizer/optimizer.py:1078-1109），grad norm 在 grad_stats_parallel_group 上归约（megatron/core/optimizer/optimizer.py:545-551）。DeepSpeed ZeRO-3 在 step() 里做 unscale_and_clip_grads（deepspeed/runtime/zero/stage3.py:3019-3028）。在 micro-batch 循环里裁剪会改变算法，不能当成等价实现。

rampup：在这个 Megatron 快照里已弃用

历史上的 rampup_batch_size 是三元组：起始 batch、增量、ramp-up samples。配置文档的例子是 [16, 8, 300000] 配合 global-batch-size 1024：从 16 开始，按 (1024−16)/8 = 126 个区间线性升到 1024，每个区间约 300000/126 = 2380 个样本（megatron/training/config/training_config.py:21-31）。但在这个 commit 里，rampup_batch_size 已标记 deprecated 且被忽略（megatron/core/num_microbatches_calculator.py:100-104、157-161），替代品是 step_batch_size_schedule，格式 "THRESHOLD:BS ..."，阈值可按样本或 token（megatron/training/config/training_config.py:33-40）。schedule 更新时会按 consumed samples 重新算 current_global_batch_size 和 num_micro_batches（megatron/core/num_microbatches_calculator.py:575-593）。

不要把 batch rampup 和 lr warmup 混为一谈：前者改变 GBS/GA，后者改变学习率。Goyal 的 warmup 是 lr 从 0.1 升到 3.2（§2.2）；rampup 是让账本本身随训练推进变化。前者改变第几步的 GBS，后者改变第几步的 lr；两者都会让「第几步的有效更新」不同。

PP 的 microbatch 不是第二个 GA

Megatron 的 pipeline schedule 收到的 num_microbatches 就是 get_num_microbatches()（megatron/training/training.py:3460、3574-3579），也就是 GBS / (MBS × DP)。因此不要写 GBS = MBS × DP × GA × PP：PP 不乘 GBS。PP 的 microbatch 是流水线调度单位，每个 microbatch 的梯度都累加到同一组缓冲区；PP 带来的额外约束是 interleaved 1F1B 要求 num_microbatches 是 PP degree 的整数倍（SC21 §2.2.2），以及 bubble 会降低利用率。想通过增加 PP microbatch 来减少 bubble，必须同时调 GBS、MBS 或 DP，否则你改变的是优化 batch，不只是调度。

它解决不了什么

• 不省参数、梯度和 optimizer state：GA 只把激活峰值压到一个 micro-batch；模型、梯度缓冲区和优化器状态仍然要完整存放。
• 不自动修正 BN：每个 micro-batch 的 BN 统计仍是局部的；大 batch 的 BN 行为要靠 sync BN / ghost BN 或调参。
• 不保证大 batch 精度：Goyal 的 8k 结果依赖 gradual warmup、固定 BN per-worker batch 和特定 lr 规则；Smith 的高 momentum 设置也会掉点。
• 不是逐位等价：dropout、数据顺序、BN 统计和 optimizer 动量轨迹都会让「累积」与「一次大 batch」不同；只有在固定模型、无随机性、相同数据顺序且 loss 归约一致时，梯度才相等。
• DDP 的 no_sync 不省显存：梯度仍然完整保存在 .grad；FSDP 的 no_sync 反而因为保留完整梯度而更费显存（torch/distributed/fsdp/fully_sharded_data_parallel.py:1053-1059）。

本页未覆盖

手算和交互假设固定 seq_len、无 padding、无 drop_last、无 packing、无 rampup、无 PP bubble；没有覆盖 TP/CP 的通信账、PP 的 bubble 公式、ZeRO 的分区细节、FSDP 的 shard 策略、以及不同框架对 no_sync 的边界差异。表 2 B/C/D 行的「1 次同步」是 no_sync / 延迟归约下的计数；A 组 DP=1 为 0 次。如果实现选择每个 micro-batch 都同步，计数会变成 GA，但样本账不变。