背景：显存省下来了，等效批次却可能已经变了
梯度累积（Gradient Accumulation） 常被理解为一种简单的显存技巧：把一个大批次拆成多个微批次，连续执行多次前向和反向传播，只在累积窗口结束时调用一次 optimizer.step()。
在定长、单卡、全精度的简单任务中，把每个微批次的 Loss 除以累积步数，通常能得到接近大批次训练的结果。但 LLM 微调经常同时存在以下条件：
- 每个样本的有效 Token 数量不同；
- Prompt、Padding 和被 Mask 的标签不参与 Loss；
- 多个 DDP Rank 上的有效 Token 数量不一致；
- 使用 BF16、FP16 或 GradScaler；
- 学习率、Warmup、日志和保存周期按 Step 推进；
- Epoch 末尾可能只有不完整的累积窗口。
这些因素叠加后，“微批次数相同”不等于”优化器看到的数据权重相同”。最隐蔽的问题不是程序报错，而是训练可以正常运行，Loss 也持续下降，但与预期的大批次基线不再数学等价。
核心原理：LLM 的有效批次应按可训练 Token 计算
样本数只是外壳，Loss 的真实分母是有效标签数
常见的等效批次公式是：
samples_per_update = micro_batch_size × accumulation_steps × data_parallel_world_size它适合描述吞吐和 DataLoader 行为，但对自回归语言模型而言还不够。Causal LM 的交叉熵通常只计算 labels != -100 的位置，因此一次更新的真实统计质量更接近：
valid_tokens_per_update = 所有 Rank、所有微批次中 labels != -100 的 Token 总数假设两个微批次分别有 100 个和 900 个有效 Token。若先对每个微批次取 Mean Loss，再把两个 Mean 相加后除以 2，那么两个微批次各占 50% 权重；正确的全批次 Loss 应让它们分别占 10% 和 90%。
因此，Token 级任务应优先采用：
loss = 所有有效 Token 的 loss_sum / 所有有效 Token 的总数而不是：
loss = 每个微批次 mean_loss 的平均值Hugging Face 在修复 Gradient Accumulation 问题时明确指出，Causal LM 等 Token 级任务需要把整个累积窗口的 Loss 总和除以全部非 Padding Token 数，而不能平均各微批次已经归一化的 Loss。
分母必须跨 Rank 汇总
DDP 默认会对各 Rank 的梯度做平均。若每个 Rank 的有效 Token 数差异较大，只按本地 Token 数归一化，会让不同 Rank 的梯度权重失真。
一种清晰的实现方式是：
- 每个 Rank 统计本累积窗口的本地有效 Token 数；
- 使用
all_reduce(SUM)得到全局有效 Token 数； - 每个微批次用
reduction="sum"计算本地 Loss Sum； - 考虑 DDP 最终会除以 World Size，对本地 Loss 乘以
world_size / global_valid_tokens； - 累积窗口最后一次反向传播触发梯度同步。
这样最终梯度对应的是全局所有有效 Token 的平均 Loss。
工程落地：把一次优化器更新作为最小事务
下面代码展示了 BF16 + DDP 场景的核心结构，强调边界和归一化逻辑：
from contextlib import nullcontext
import torch
import torch.distributed as dist
import torch.nn.functional as F
IGNORE_INDEX = -100
def causal_lm_loss_sum(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
shift_logits = logits[..., :-1, :].contiguous().float()
shift_labels = labels[..., 1:].contiguous()
return F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=IGNORE_INDEX,
reduction="sum",
)
def train_update(ddp_model, microbatches, optimizer, scheduler, max_grad_norm: float):
"""microbatches 共同形成一次 optimizer update。"""
optimizer.zero_grad(set_to_none=True)
device = next(ddp_model.parameters()).device
local_tokens = torch.zeros((), device=device, dtype=torch.long)
for batch in microbatches:
local_tokens += batch["labels"].ne(IGNORE_INDEX).sum()
global_tokens = local_tokens.clone()
if dist.is_initialized():
dist.all_reduce(global_tokens, op=dist.ReduceOp.SUM)
world_size = dist.get_world_size() if dist.is_initialized() else 1
if global_tokens.item() == 0:
raise ValueError("Accumulation window contains no trainable tokens")
denominator = global_tokens.to(torch.float32)
for index, batch in enumerate(microbatches):
is_last_microbatch = index == len(microbatches) - 1
sync_context = (
nullcontext() if is_last_microbatch else ddp_model.no_sync()
)
# DDP 要求 forward 也位于 no_sync 上下文内
with sync_context:
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
outputs = ddp_model(
input_ids=batch["input_ids"],
attention_mask=batch.get("attention_mask"),
)
loss_sum = causal_lm_loss_sum(outputs.logits, batch["labels"])
# DDP 会对 Rank 梯度取平均，因此乘 world_size 抵消该平均
loss = loss_sum * world_size / denominator
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(
ddp_model.parameters(), max_norm=max_grad_norm,
)
optimizer.step()
scheduler.step()
return {
"global_valid_tokens": int(global_tokens.item()),
"microbatches": len(microbatches),
"grad_norm": float(grad_norm),
}这段代码有三个关键约束：
- Loss 使用
reduction="sum"，统一在累积窗口级别归一化； - 中间微批次使用
no_sync()，最后一个微批次才触发 DDP 同步； optimizer.step()和scheduler.step()都只在完整更新边界执行一次。
注意：如果使用 Hugging Face Trainer、Accelerate、DeepSpeed 或 FSDP，应优先遵循框架自己的 Loss、同步和 Step 契约，不要把上述 World Size 补偿公式机械叠加到框架已经完成的归一化逻辑上。
no_sync 的正确边界
PyTorch DDP 的 no_sync() 会暂停跨进程梯度同步，并把梯度保存在本地，直到第一次退出该上下文的 Forward/Backward 才同步。
容易写错的地方是只把 backward() 放入 no_sync()，但把 Forward 留在外面。PyTorch 官方文档明确要求 Forward 也必须位于上下文内，否则梯度仍可能发生同步。
推荐规则：
| 微批次位置 | Forward | Backward | 梯度同步 |
|---|---|---|---|
| 前 N-1 个 | no_sync() 内 | no_sync() 内 | 不触发 |
| 第 N 个 | 正常上下文 | 正常上下文 | 触发累计同步 |
每个 Rank 必须以一致顺序进入和退出同步边界。
AMP、梯度裁剪和跳步必须处于同一更新边界
FP16 + GradScaler 下，累积窗口内的梯度应保持相同 Scale。PyTorch 的 AMP 文档要求：
- 中间微批次只执行
scaler.scale(loss).backward()； - 累积完成后再调用一次
scaler.unscale_(optimizer)； - 在 Unscale 后进行梯度裁剪；
- 最后调用
scaler.step(optimizer)和scaler.update()； - 一个有效批次内不能改变 Scale，也不能把已缩放梯度与未缩放梯度混合累积。
如果检测到 Inf/NaN，GradScaler 可能跳过本次 Optimizer Step。此时 Scheduler 也不应独立前进，否则学习率时间轴会领先于参数更新次数。生产代码应以”优化器是否真实更新”为 Scheduler、EMA、全局 Update ID 和 Checkpoint 命名的共同信号。
尾部窗口：不要继续除以固定累积步数
Epoch 最后经常只剩 1 到 accumulation_steps - 1 个微批次。若仍固定除以完整累积步数，尾部梯度会被额外缩小。
更稳妥的策略有两种：
- 动态窗口：按尾部实际有效 Token 总数归一化，并正常执行一次更新；
- 丢弃尾部：使用
drop_last或数据层补齐，使所有更新窗口固定，但要明确记录被丢弃或重复的数据。
对流式数据集还要额外注意：框架是否在 DataLoader 结束时自动触发同步和 Step。Accelerate 的 Gradient Accumulation Plugin 默认会与当前 DataLoader 结束状态同步；自定义无限流或跨 Epoch 累积时，需要显式配置该行为。
Scheduler 与训练预算：统一使用 optimizer_update_id
训练系统中至少存在三种 Step：
| Step 类型 | 触发频率 | 适用场景 |
|---|---|---|
| Micro Step | 每个微批次一次 | 仅用于调试、细粒度日志 |
| Optimizer Update | 每个累积窗口一次 | Warmup、LR Scheduler、EMA、Checkpoint |
| Token Step | 累计到一定有效 Token 数 | 吞吐统计、跨实验对齐 |
Warmup、LR Scheduler、权重衰减、EMA 和 Checkpoint 周期通常应基于 Optimizer Update；吞吐、数据预算和跨实验对齐则更适合同时记录累计有效 Token。
建议为训练日志固定以下字段：
micro_step_id
optimizer_update_id
epoch
global_valid_tokens_in_update
cumulative_train_tokens
learning_rate
optimizer_step_skipped
grad_norm_before_clip
grad_norm_after_clip不要让一个含义模糊的 global_step 同时承担三套语义。
适用场景
本文方案尤其适合：
- 变长对话数据的 Causal LM 或 SFT；
- Prompt 部分使用
-100Mask、只训练 Assistant 输出的任务； - Sequence Packing 后每个 Pack 的有效训练 Token 数不同；
- 多卡 DDP/FSDP，且各 Rank 的长度分布可能不完全一致；
- 需要比较不同 Micro Batch Size 或 Accumulation Steps 的可复现实验。
对于分类、排序或偏好学习等样本级目标，分母未必应是 Token 数。核心原则不是”所有任务都按 Token 除”，而是根据目标函数定义选择正确统计单位，并确保累积前后保持一致。
常见误区
误区一：Mean Loss 再除以累积步数一定正确
只有各微批次统计单位和有效数量相同时才成立。变长 Token 任务应对 Loss Sum 和有效 Token 总数做统一归一化。
误区二：每个微批次同步一次只是性能差一点
它通常不会改变数学结果，但会产生不必要的 All-Reduce，显著削弱梯度累积原本想减少通信频率的价值。
误区三：Scheduler 可以按 DataLoader Batch 推进
这会让学习率随 Micro Step 而不是参数更新推进。改变累积步数后，训练曲线会直接失去可比性。
误区四：每个微批次都做梯度裁剪
这会改变向量相加后的方向和幅度。梯度裁剪应作用于完整累积窗口的最终梯度。
误区五：尾部窗口继续使用固定除数
不完整窗口必须使用实际统计量，或明确丢弃；不能假装它仍包含完整数量的微批次。
上线检查清单
- 用同一随机种子比较”大批次单步”与”多个微批次累积一步”的参数更新差异；
- 构造长度悬殊的样本，确认短样本不会被等权放大；
- 检查全局有效 Token 数是否跨 Rank 汇总；
- Profile 每个 Optimizer Update 的 All-Reduce 次数，确认中间微批次未同步；
- 验证
no_sync()同时覆盖 Forward 和 Backward； - 确认尾部窗口、空标签窗口和 DataLoader 结束行为；
- FP16 下确认 Unscale → Clip → Step → Update 的顺序；
- 发生 Overflow 跳步时，确认 Scheduler、EMA 和 Update ID 不前进；
- 日志同时记录 Micro Step、Optimizer Update 和累计有效 Token；
- 改变
gradient_accumulation_steps后重新核对 Warmup 与总更新次数。