Trainer floor-divides valid-token counts under model parallelism, silently mis-scaling gradients and optimizer updates

快速结论：该问题出现在使用 Transformers Trainer 进行模型并行（TP/CP/SP）训练时，有效 token 数无法被并行度整除会被地板除截断，导致 loss、梯度范数和优化器更新被静默放大；优先检查 _get_num_items_in_batch() 中的整除逻辑，并核对有效 token 数与并行度是否存在非整除关系。

适用环境：Transformers 5.23.0.dev0（上游 main 提交 ca8c871ede62b521a75a2ddfc5e1968ac9c8e097）；PyTorch 2.11.0+cu128；Python 3.13；Linux；2 × NVIDIA Tesla T4；FSDP2 + 上下文并行（cp_size=2）；FP32；SGD，学习率 1e-3。

最快修复方案：暂无确认的一步修复方案；Issue 讨论中提出的方向是保留分数 token 计数（改用精确浮点除法），并为可整除/不可整除两种情况补充回归测试，相关 PR 已进入合并流程。

注意事项：修复需确认不影响普通单卡训练；上下文并行负载均衡关闭时的诊断结果不能直接推广到其他并行组合；升级前请先在小型复现脚本上验证分母输出。

问题场景

用户在 Transformers Trainer 中启用 FSDP2 与上下文并行（cp_size=2）训练一个随机初始化的极简 LlamaForCausalLM（1 个解码层、hidden size 32），序列长度 64，使用 FP32 和 SGD 执行优化器更新时触发该问题。当有效目标 token 数为 3 等不能被并行度整除的值时，训练 loss、梯度范数和参数更新量被错误放大。

报错原文

Trainer floor-divides valid-token counts under model parallelism, silently mis-scaling gradients and optimizer updates

num_items_in_batch = num_items_in_batch // (
    self.get_tp_size() * self.get_cp_size() * self.get_sp_size()
)

valid=4 actual=2.0 expected=2.0
valid=3 actual=1.0 expected=1.5
valid=1 actual=0.0 expected=0.5

原因分析

最可能的原因是 Trainer._get_num_items_in_batch() 在按张量并行、上下文并行、序列并行因子归一化有效 token 数时使用了地板除 //。当有效 token 数不能被并行度因子整除时，分母被截断（例如 3 个有效 token、cp_size=2 时得到 1 而不是 1.5），截断后的分母传入 causal LM 交叉熵，造成 loss 和梯度被过度缩放。在非整除案例中，实测 loss、梯度范数和 SGD 参数更新范数均恰好放大 1.5 倍；可整除的对照案例完全一致。

环境排查

确认 Transformers 版本及是否位于含该整除逻辑的提交（Issue 中上游 main 为 ca8c871ede62b521a75a2ddfc5e1968ac9c8e097）。
确认 PyTorch 版本为 2.11.0+cu128，Python 3.13。
确认 GPU 数量与型号：2 × NVIDIA Tesla T4。
确认分布式配置：FSDP2、上下文并行 ParallelismConfig(cp_size=2)，以及 TP/SP 大小。
确认 average_tokens_across_devices=False、是否启用梯度裁剪、上下文并行负载均衡是否开启。
统计每个 batch 的有效目标 token 数，检查是否能被 tp_size * cp_size * sp_size 整除。

解决步骤

定位 src/transformers/trainer.py 中 _get_num_items_in_batch() 的并行度归一化表达式。
在修改前先用 Issue 提供的最小复现脚本统计分母：构造 labels，将前部设为 -100、尾部设为有效 token，分别传 4、3、1 个有效 token，打印 actual 与 expected。
可优先尝试将地板除 // 改为精确浮点除法，使分母保留分数值（如 3 / 2 = 1.5），避免截断。
为可整除与非整除两种情况补充回归测试，同时验证普通单卡训练路径不受影响。
应用修复后重新运行最小复现脚本及双 GPU FSDP2 + 上下文并行训练步骤，比较 loss、梯度范数和优化器更新范数。

验证方法

运行最小复现脚本，有效 token 数为 4 时输出 actual=2.0、3 时输出 actual=1.5、1 时输出 actual=0.5；在双 GPU 实验中，非整除案例（3 个有效 token）的分母应为 1.5，loss、梯度范数和 SGD 更新范数与精确除法参考实现比值应为 1.000000，可整除案例应继续完全一致。同时确认两个 GPU 进程均完成训练步骤，且单卡训练无行为变化。