RFC-0032: 对齐 MaxText 损失归一化与 Megatron（calculate_per_token_loss 逐样本语义）

概述

MaxText 已有与 Megatron 同名的 calculate_per_token_loss 配置项，但目前仅在梯度累积（GA>1）时改变聚合方式，在 Ling3 参考跑的 GA=1 场景下完全不生效——导致 lm loss 的**训练目标（梯度）**仍是全局 token 加权，与 Megatron 参考跑（calculate_per_token_loss=False、micro_batch_size=1，逐样本等权）不一致。本提案补全该配置在任意 GA 下的语义：=False 时 lm loss 目标改为逐样本归一、样本等权，同时保持日志端 learning/lm_loss 的全局 token 平均不变，使 MaxText 与 Megatron 在「目标逐样本、日志全局」的双通道行为上完全同构。

背景与动机

Megatron 的双通道归一化（参考跑实测配置）

ling3-plus 参考作业的 config dump 已确认 calculate_per_token_loss = False、micro_batch_size = 1。在该配置下，Megatron 对梯度和日志走两条独立通道：

通道	lm loss	mtp loss	router z-loss
梯度（`=False`，参考跑）	逐微批归一 + 微批等权（`schedules.py` `forward_step_calc_loss` else 分支：`output/=num_tokens_mb; /=num_microbatches`）	逐微批归一（`multi_token_prediction.py::_loss_forward` `sum(CE·mask)/num_tokens_mb`）× `loss_scale·cp/num_microbatches`	逐微批 × `1/num_microbatches`（`MoEAuxLossAutoScaler`）
梯度（`=True`，对照）	原始 sum 透传，`finalize_model_grads` 统一 ÷ 全局有效 token 数（DP all-reduce 后）	原始 sum 注入（scale 不除 `num_microbatches`），随全体梯度 ÷ 全局 lm token 数——归一化子是主 loss 的 token 数，而非 mtp 自身 rolled mask 的 token 数	原始注入，随全体梯度 ÷ 全局 lm token 数
日志（与 flag 无关）	恒为全局 Σsum/Σcount（`(sum, count)` dict 通道）	恒为逐微批均值（tracker × `1/num_microbatches`）	恒为逐微批均值（tracker）

关键结构：calculate_per_token_loss 只切换梯度端，完全不影响日志端——lm 日志恒为全局 token 平均、mtp 日志恒为微批平均。代码层面，mtp 的 tracker 写入（multi_token_prediction.py:1585-1597，用未加权的 raw_mtp_loss / num_tokens_mb）发生在 calculate_per_token_loss 分支（:1599-1615）之前，两条通道物理隔离。因此 =True 下 Megatron 的 mtp 会出现与 =False 相反的不一致（梯度全局、日志逐微批）——即两条日志估计量的差异不能靠切 flag 消除。

参考跑（=False）下 mbs=1，微批=单样本，梯度端所有 loss 都是「一样本一票」，梯度端内部是自洽的；不一致的是日志端的 lm loss——它走的是另一条协议通道：

上层训练框架 antllm 的 loss_func（antllm/training/megatron_train_utils.py:489）自己不做除法，返回 (loss_sum, num_tokens, {'lm loss': (sum, count)})（:542-544、:597-605）；
Megatron-core 的 train_step（megatron/training/training.py:2058-2073）对 (sum, count) 对跨微批先累加分子分母、最后才除一次，跨 DP 的归约同样是 sum-of-sums / sum-of-counts（上游 pretrain_gpt.py:267-268 在 loss_func 内做 DP all-reduce；Ant 栈由 atorch/dlrover 日志侧完成），得到全局 token 平均；
而 MTP / MoE aux 这类在模型 forward 中途产生的 loss 只能靠 autograd 侧信道（MTPLossAutoScaler/MoEAuxLossAutoScaler）注入，注入时必须已完成归一化，日志 tracker 存的也是微批内已除好的值——没有机会做「先累加后除」。

即：日志 lm loss 的估计量（全局）≠ 它自己训练目标的估计量（逐样本），这是协议结构使然，不是有意设计。

MaxText 现状与差距

calculate_per_token_loss 定义于 src/maxtext/configs/types.py:1539（默认 True），ling2 / ling3-tiny / ling3-plus / ling3-flash 的模型 YAML 均已显式设为 false——但它目前只出现在 GA raw-sum 判断中（src/maxtext/trainers/pre_train/train.py:469-473），GA=1 时被旁路；
lm loss 训练目标恒为全局 jnp.sum(xent) / total_weights（src/maxtext/trainers/pre_train/train.py:418, 458, 477），等价于 Megatron =True 的语义；
日志端 learning/lm_loss = 全局 Σ/Σ（src/maxtext/trainers/pre_train/train.py:741-745），与 Megatron 日志一致 ✅；
MTP loss 经 maxtext#637 已改为逐样本归一（sow 粒度 [batch]，mtp_per_layer_loss_norm=True 分支层×样本等权），与 Megatron 的梯度、日志均已对齐 ✅；
剩余差距 = lm loss 的梯度：全局 token 加权 vs 逐样本等权。两个估计量之差为 Cov(样本 token 数, 样本平均 loss) / E[token 数]，在 ling3 数据 blend 上约 0.1% 的样本重加权。该差异在日志上不可见（两边日志同为全局估计量），因此在 #637 排查 mtp +0.1% 常数偏移之前一直未被发现。

收益

训练目标与 Megatron 严格同构，把最后一个已知的估计量差从 loss-parity 的假设列表中移除——后续残余漂移分析（如 plus 的 router 分叉问题）不再受此干扰；
ling 系列 YAML 里已存在的 calculate_per_token_loss: false 从「仅 GA 生效的部分语义」变为与 Megatron 同名同义，消除配置误导。

目标

calculate_per_token_loss=False 时（任意 GA 取值），lm loss 训练目标 = 逐样本归一（max(count, 1) 兜底）后样本等权平均，与 Megatron =False + mbs=1 数学等价；
日志端 learning/lm_loss 保持全局 Σ/Σ 不变；loss-parity 比较脚本零改动，lm/mtp 曲线在既有容差（lm 0.02% / mtp 0.2%）内零漂移;
calculate_per_token_loss=True（默认）路径行为逐 bit 不变，保证上游模型零影响；
单测可验证：手算逐样本归一（含零 token 样本）、两估计量方向性、GA=1 与 GA>1 等价性、=True 回归。

非目标

不改 Megatron / antllm 侧任何代码；
MTP loss 在 =False 模式的归一化（#637 已完成对齐）；=True 模式的 mtp 梯度已在实现阶段追加对齐（calculate_mtp_loss 按全局 lm token 数归一，见约束与注意事项）。仍不复刻的只剩 =True 模式下的 mtp 日志——Megatron 该模式日志仍为逐微批，MaxText 目标与日志耦合；ling3 全在 =False，此差异无实际影响；
router z-loss 的逐样本化：Megatron 端严格为逐微批，但系数仅 7e-5/模型（每层 ~1e-6），估计量差对目标的影响在 1e-7 量级——列为 Phase 2 可选项，先离线测量再决定；
load-balance aux loss、indexer loss：ling3 未启用（aux-free bias 路由、无 indexer），保持现行为；
eval loss 及 SFT / RL / DPO trainer 的损失路径；
日志端估计量的任何变化（日志两边本来就是对齐的）。

方案

核心思路：把 calculate_per_token_loss 的语义从「GA 聚合约定」提升为「损失归一化粒度」，与 Megatron 同名同义；同时将训练目标与日志解耦——这正是 Megatron 本身的双通道结构。

=True（默认）：现行为完全不变。GA>1 时走 raw-sum、由 GA 统一除全局 token 数；GA=1 时全局平均。lm 上等价 Megatron =True（mtp 不映射该模式，见约束与注意事项）；
=False：lm loss 逐样本归一、样本等权。GA>1 时每个 micro-step 内逐样本、step 间等权（每 step 样本数相同 ⇒ 总体仍为样本等权），现有 use_ga_raw_sum=False 路径天然兼容。等价 Megatron =False（mbs=1）；
日志：learning/lm_loss 继续用全局 Σ/Σ（与 Megatron 日志同估计量）；learning/loss（目标值）在 =False 下会比旧值低约 0.1% 量级（协方差项，token 多的样本 loss 偏高），属预期变化。

备选方案

新增独立开关（如 per_sample_loss_norm）——否决：与既有同名 flag 语义重叠，ling YAML 已通过 calculate_per_token_loss: false 表达了意图；两个开关还会组合出无意义状态；
无条件改为逐样本（照搬 #637 的 MTP 做法）——否决：#637 的 sow 粒度变更在关闭 per-layer-norm 时数学不变、无条件安全；而 lm 直接改目标会影响所有上游模型，必须由开关门控；
反向对齐：让 Megatron 参考改跑 =True——否决：参考日志是既成工件，生产 Megatron 配置归属另一团队，本 fork 的目标是复现参考。

约束与注意事项

两个开关的分工与合法组合：MaxText 中 calculate_per_token_loss 决定 lm 归一化及（=true 时）mtp 的全局归一化；mtp_per_layer_loss_norm 只在 =false 下决定 mtp 粒度。与 Megatron 的模式映射为：
- Megatron =False ↔ MaxText calculate_per_token_loss=false + mtp_per_layer_loss_norm=true：ling3 全系现配置，本 RFC 对齐的组合——落地后梯度端（lm/mtp 均逐样本）与日志端（lm 全局、mtp 逐样本）四象限与 Megatron 完全一致；
- Megatron =True ↔ MaxText calculate_per_token_loss=true：mtp 梯度同样映射——calculate_mtp_loss 在 =true 时按全局 lm token 数归一（覆盖 mtp_per_layer_loss_norm，与 Megatron 的 finalize ÷ 全局 lm token 数一致），GA raw-sum 的 mtp_loss × total_weights 补偿随之从近似变为精确；唯一不映射的是该模式下的 mtp 日志（Megatron 恒逐微批，MaxText 与目标耦合），见非目标；
Megatron 的微批 = mbs 个样本，本方案的「逐样本」只在 mbs=1 时与「逐微批」等价。参考跑与生产 ling3 均为 mbs=1；若未来参考侧 mbs>1，需扩展为按 mbs 分组归一（不在本期）；
零 token 样本：Megatron 对 num_tokens=0 的微批贡献 0 但仍计入 1/num_microbatches（schedules.py 中 output*=0 分支 / MTP 的 max(num_tokens,1)）；MaxText 用 jnp.maximum(w, 1) 等价实现（分子被 mask 后为 0）；
vocab tiling（num_vocab_tiling>1）路径只返回标量 loss sum，暂不支持逐样本；=False 且 tiling>1 时在 pyconfig 显式报错（ling3 未使用 tiling）；
xent 已含 lm-head z-loss 项（z_loss_multiplier，ling3 为 0），逐样本归一跟随 xent 整体进行，两种取值下语义自洽。

设计细节

改动集中在 src/maxtext/trainers/pre_train/train.py::loss_fn（Linen 与 NNX 两分支同改）：

python

# 现状 (src/maxtext/trainers/pre_train/train.py:469-477)
use_ga_raw_sum = (
    config.gradient_accumulation_steps > 1
    and config.calculate_per_token_loss
    and not config.use_tunix_gradient_accumulation
)
if use_ga_raw_sum:
  loss = total_loss
else:
  loss = total_loss / (total_weights + EPS)

改为：

python

if use_ga_raw_sum:
  # =True 且 GA>1：raw sum，GA 结束后统一除全局 token 数（现状，不变）
  loss = total_loss
elif config.calculate_per_token_loss:
  # =True 且 GA=1：全局 token 平均（现状，不变）
  loss = total_loss / (total_weights + EPS)
else:
  # =False：逐样本归一 + 样本等权（新增，对齐 Megatron mbs=1）
  per_sample_xent = jnp.sum(xent, axis=-1)                                   # [B]
  per_sample_weights = jnp.sum(data["targets_segmentation"] != 0, axis=-1)   # [B]
  loss = jnp.mean(per_sample_xent / jnp.maximum(per_sample_weights, 1))

实现要点：

xent 需以 [B, S] 形状存活到此处（现状已满足，mask 已乘）；
total_loss / total_weights 照旧计算并透传 aux——learning/lm_loss 日志（src/maxtext/trainers/pre_train/train.py:741-745）与 parity 流程零改动；
逐样本中间量沿 batch 维分片（activation_embed_and_logits_batch），额外开销可忽略；
MTP：mtp_loss 已是逐样本等权均值（#637），loss += mtp_loss 无需变动；其粒度由 mtp_per_layer_loss_norm 控制，与本 flag 正交（见约束与注意事项）；
moe_z_loss：Phase 1 保持全局归一（见非目标）；=False 下 use_ga_raw_sum=False，aux 加法路径不变；
indexer 稠密预热（total_loss=0.0）路径：=False 下同样置 0，无交互；
pyconfig 新增校验：calculate_per_token_loss=False and num_vocab_tiling>1 → 报错；
GA：代码零改动，等价性已逐行核对，见下节「GA 开启时的等价性」。

与 Megatron 的等价性论证

Megatron =False、mbs=1：单卡目标 = (1/N_mb) · Σ_mb [Σ(CE·mask)_mb / max(tok_mb, 1)]，N_mb = 样本数；DP 端梯度 all-reduce 取均值 ⇒ 全局所有样本等权。MaxText 新路径：mean_B [Σ(CE·mask)_b / max(tok_b, 1)]，jit 全局 batch 即全体样本 ⇒ 同一估计量。零 token 样本两边均「贡献 0、计入分母」。CP 维度 MaxText 无对应概念（无 ×cp_size 因子），序列不切分，天然一致。

GA 开启时的等价性（gradient_accumulation_steps > 1，已逐行核对）

本 fork 的 src/maxtext/utils/gradient_accumulation.py 已按 Megatron 语义定制了分母策略，与本 RFC 的 loss_fn 改动组合后无需再改：

梯度（=False）：GA 将全局 batch 均分为 K 个 micro-step（:138-146 reshape，每步恰好 B/K 个样本），每步调用同一个 _loss_fn 得到「该步样本的逐样本均值」（本 RFC 改动点）；梯度按步求和（:134），末尾 grad_divisor = gradient_accumulation_steps（:174）。合成 (1/K)·Σ_steps mean_{B/K}(逐样本) = 全体样本等权，与 Megatron =False（每微批 ÷ 自身 token 数 × 1/num_microbatches，DP 梯度取均值）同一估计量。Megatron 的「微批」与 MaxText 的「GA micro-step」粒度不同（前者 mbs=1 个样本，后者 B/K 个样本），但逐样本归一后两级平均（步内均值 → 步间等权）恰好坍缩为同一个「样本等权」估计量，粒度差异不再可见；
梯度（=True）：raw-sum 模式，梯度 ÷ 累计 total_weights（全局 token 数，:174）= Megatron =True 的 finalize_model_grads ÷ 全局 token 数。现状不变；
日志：lm_loss = Σ_steps CE_sum / Σ_steps weights（:180，即 sum-of-sums / sum-of-counts）恒为全局 Σ/Σ，与 flag 无关，对应 Megatron 日志 lm；raw_mtp_loss 按步累加 ÷K（:197）= 逐样本均值，对应 Megatron mtp tracker × 1/num_microbatches。两者均不受本 RFC 影响；
learning/loss 的混合定义：GA>1 时上报的 loss 为「全局 lm_loss + 各 aux/K」（:182-188），不是训练目标本身（既有设计，为保持上游上报语义）。因此本 RFC 引起的目标估计量变化只会体现在 GA=1 的 learning/loss 上（约 −0.1% 平移）；GA>1 的 learning/loss 不变。learning/lm_loss / learning/raw_mtp_loss 在任意 GA 下定义一致；
use_tunix_gradient_accumulation 仅用于 post-train HF 数据管线（SFT），预训练 GA 不经过该路径，不在本 RFC 范围。

测试计划

单元测试（tests/unit，CPU 可跑，复用 #637 的测试模式）：

手算校验：构造 token 数不均 + 一个全 padding 样本的 batch，断言 =False 的 loss 与手算逐样本均值一致，零 token 样本贡献 0 且计入分母；
方向性：构造「token 多的样本 loss 更高」的 batch，断言 =False 估计量 < =True 估计量（协方差符号）；
回归：=True 下新旧代码 loss 逐 bit 一致；
GA 等价性：同一全局 batch 在 GA=1 与 GA=4 下（=False）梯度以及 lm_loss / raw_mtp_loss 指标一致（浮点容差内）；learning/loss 因 GA>1 的混合上报定义（见「GA 开启时的等价性」）允许不同，不作断言;
校验路径：num_vocab_tiling>1 + =False 报错；
配置断言：ling3 系列 model YAML 的 =False 实际进入逐样本分支。

端到端（TPU Pod 执行）：

scripts/ci_smoke_test.sh 通过；
flash 64chip 短跑：learning/lm_loss、learning/raw_mtp_loss 曲线与改动前一致（日志通道未动），learning/loss 出现约 −0.1% 量级的预期平移；
夜间 loss-validation（flash）：lm / mtp 相对差保持在既有容差（lm 0.02% / mtp 0.2%）内；
plus 长跑（择机）：观察残余漂移是否收窄——不作为验收门槛：本改动对梯度是 ~0.1% 的方向性重加权，低于观测噪声，验收以「构造正确（单测）+ 无回归（e2e）」为准。

影响范围

训练目标变化：所有设置 calculate_per_token_loss: false 的模型（ling2 / ling3-tiny / ling3-plus / ling3-flash）。默认 True 的上游模型零影响；
日志与 parity：learning/lm_loss、learning/raw_mtp_loss 及全部比较脚本零改动；learning/loss 的数值定义变化需周知（若有看板/告警引用其绝对值需检查阈值）——注意该平移仅出现在 GA=1（GA>1 时 learning/loss 是混合上报定义，不受影响，见设计细节）；
参考文件：scripts/reference_losses_*.json 与 Megatron 参考日志无需重新生成（日志估计量未变）；
checkpoint 兼容：无优化器/参数结构变化，可直接加载旧 ckpt；但目标函数改变，建议随新 run 生效而非 mid-run 热切；
文档：src/maxtext/configs/types.py 字段描述需重写（现描述只提 GA），wiki 的 accuracy-alignment 相关页面同步更新。

实施计划

Phase 1（核心，单 PR）：loss_fn 改动 + pyconfig 校验 + 全部单测；TPU Pod 跑 smoke + flash 短跑验证后合入；连续观察两晚 flash nightly loss-validation；
Phase 2（可选）：用现有日志离线测量 router z-loss 两估计量之差；若超出日志噪声再实施 z-loss 逐样本 sow（做法同 #637），否则在本 RFC 记录 wontfix 结论；
Phase 3：文档收尾——src/maxtext/configs/types.py 描述、accuracy-alignment 页、本 RFC 状态从 draft 推进。

风险与缓解

风险	缓解
改动训练目标引起 flash nightly 波动	效应量为 ~0.1% 的样本重加权（且 mtp 先例 #637 未引起可见 lm 变化）；连续观察 2 晚，超容差即 revert（单 PR，可干净回滚）
`=False` + GA>1 旧路径语义变化（原「逐 micro-step 全局」→ 新「逐样本」）	该组合当前无生产使用（参考跑 GA=1）；语义变化正是本 RFC 目的，PR 说明中显式注明
`learning/loss` 平移触发既有看板告警	合入前周知并更新阈值
零 token 样本兜底与 Megatron 边界行为不一致	单测 1 显式覆盖；两边语义均为「贡献 0、计入分母」，已逐行核对
未来 Megatron 参考改用 mbs>1	约束节已声明等价条件；届时扩展为按 mbs 分组归一

RFC-0032: 对齐 MaxText 损失归一化与 Megatron（calculate_per_token_loss 逐样本语义） ​

概述 ​

背景与动机 ​

Megatron 的双通道归一化（参考跑实测配置） ​

MaxText 现状与差距 ​

收益 ​

目标 ​

非目标 ​

方案 ​

备选方案 ​

约束与注意事项 ​

设计细节 ​

与 Megatron 的等价性论证 ​

GA 开启时的等价性（gradient_accumulation_steps > 1，已逐行核对） ​

测试计划 ​

影响范围 ​

实施计划 ​

风险与缓解 ​