moss-tts-local-1.5 sft NaN bug

[ruby11dog]

我的训练脚本：
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TRAIN_DATA="xxx"
export OUT_DIR="xxx"
accelerate launch
sft.py
--model-path xx
--train-jsonl "$TRAIN_DATA"
--output-dir "$OUT_DIR"
--per-device-batch-size 8
--gradient-accumulation-steps 4
--learning-rate 2.0e-5
--warmup-ratio 0.0
--lr-scheduler-type constant
--mixed-precision bf16
--channelwise-loss-weight 1,32
--gradient-checkpointing
--skip-nonfinite-batches
会有大量的NaN，无法正常训练：
warning: Non-finite gradient norm at epoch=0 global_step=1: nan. first_nonfinite_grad=module.transformer.embed_tokens.weight bad=388956160/388956160 dtype=torch.bfloat16 shape=(151936, 2560); record_ids=['205785', '55181', '191931', '193623', '160327', '31329', '34026', '194913', '219003', '98766', '48350', '36495', '149423', '58690', '242286', '112645', '82931', '217513', '71839', '42900', '128648', '18118', '119128', '76251', '135686', '177750', '26527', '157316', '162412', '41533', '154624', '179435', '234848', '53256', '243554', '201511', '119731', '131094', '101268', '70420', '127925', '106988', '34870', '211582', '101706', '119840', '119165', '225646']; skipped
都是module.transformer.embed_tokens.weight这里出的NaN，
我的训练环境在moss-delay下训练都是没问题的，麻烦帮忙看看问题，感谢！

[YWMditto]

@codex Try to provide an answer.

[chatgpt-codex-connector]

Answer

我看了一下代码，module.transformer.embed_tokens.weight 出现在日志里不一定表示 embedding 本身是根因。在 v1.5 local 模型里，text_lm_head.weight 和 transformer.embed_tokens.weight 是 tie 在一起的：self.text_lm_head.weight = self.transformer.embed_tokens.weight。所以只要文本 head 的反向传播出现 NaN，最终也会显示到 transformer.embed_tokens.weight 这个共享参数上。相关代码在 moss_tts_local_v1.5/modeling_moss_tts.py。

更可能的方向是：当前 v1.5 local SFT 路径的 loss 计算方式和 moss-delay 的 loss 聚合方式不完全一样。v1.5 local 现在是逐 head 计算 cross_entropy，然后按 head 权重加权平均；它没有像 moss-delay 那样先按每个 channel 的有效 token 数聚合/归一化。v1.5 local 的实现见 compute_supervised_loss_from_hidden。moss-delay 则会统计每个 channel 的有效 token 数，并用 total_tokens_per_channel 做归一化后再加权。

因此我建议按下面顺序排查：

1. 先关掉 --gradient-checkpointing 跑一个小 batch 验证。
   你现在的命令里开了 --gradient-checkpointing。v1.5 local 默认 checkpointing scope 是 global，默认实现是 module_wrapper，并且默认是 non-reentrant；这些参数已经在脚本里暴露出来了。如果关掉 checkpointing 后 NaN 消失，问题优先看 activation checkpointing + bf16 + attention kernel 的组合。

2. 把 --per-device-batch-size 从 8 降到 1 或 2，再确认是否仍然在 step 1 NaN。
   你的全局 batch 是 8 * 4 * GPU数，第一步就出现全 embedding 参数 NaN，说明不是训练后期发散，更像某个 forward/backward 组合在当前 batch/精度下数值不稳定。

3. 临时去掉 --channelwise-loss-weight 1,32，或把音频总权重调低，例如 1,8。
   当前参数解析说明 1,32 表示 text head 权重为 1，总 audio loss 权重为 32，然后均分到所有 audio heads。虽然 v1.5 local 最终会除以参与的 total_weight，但它的聚合路径与 moss-delay 不同，建议先降低音频 loss 权重判断是否是某些 audio head 的梯度放大触发了 BF16 不稳定。

4. 把 attention 实现切到更保守的实现测试。
   如果你现在默认走 flash attention，可以尝试加：

   --attn-implementation sdpa

   或：

   --attn-implementation eager

   这能帮助判断是不是 flash attention / bf16 backward 的数值问题。

5. --skip-nonfinite-batches 只能绕过坏 batch，不建议作为最终解决方案。
   代码是在 loss finite 时继续 backward，然后在 clip_grad_norm_ 后检查 grad norm；你的日志正是这个路径，即 loss 可能是 finite，但 backward 后梯度已经 NaN。如果大量 batch 都被 skip，训练基本无法有效进行。

一个最小排查命令可以先改成：

accelerate launch sft.py \
  --model-path xx \
  --train-jsonl "$TRAIN_DATA" \
  --output-dir "$OUT_DIR" \
  --per-device-batch-size 1 \
  --gradient-accumulation-steps 4 \
  --learning-rate 2.0e-5 \
  --warmup-ratio 0.0 \
  --lr-scheduler-type constant \
  --mixed-precision bf16 \
  --channelwise-loss-weight 1,8 \
  --attn-implementation sdpa

如果这样稳定，再逐项恢复：

1. --channelwise-loss-weight 1,32
2. --per-device-batch-size 8
3. --gradient-checkpointing
4. flash attention / auto attention

这样可以比较快定位是 loss 权重/批大小、checkpointing，还是 attention kernel + bf16 backward 导致的。

Testing / Checks

• ✅ git status --short
• ✅ rg -n "channelwise|skip-nonfinite|Non-finite|embed_tokens|bf16|loss" -S .
• ✅ nl -ba moss_tts_local_v1.5/finetuning/sft.py | sed -n '133,166p;613,707p;1071,1108p'
• ✅ nl -ba moss_tts_local_v1.5/modeling_moss_tts.py | sed -n '102,117p;180,190p'
• ✅ nl -ba moss_tts_delay/modeling_moss_tts.py | sed -n '351,374p'

View task →