feat(megatron): handle zero-token no-signal steps in the shared trainer

[xcosmosbox]

What

为共享 Megatron 训练器增加 zero-token no-signal step 的完整处理：空/全 mask response 的 token 计数口径（CP=1 与 CP>1 一致）、全 batch 无有效 loss token 时的指标归零与 fail-fast，以及全局零 token step 跳过 optimizer 与 LR scheduler 更新。

这是此前在 #205 中实现、后按 review 意见回退的 shared 基础设施改动，现按要求建议独立成 PR，方便内部做全量算法 CE 验证。

Why

当前共享训练路径对"空/全 mask response"存在两处语义不一致与一个数值风险：

1. token 计数口径不一致：get_cp_local_num_tokens() 在 CP=1 时对每个样本使用历史 per-sample clamp_min(loss_mask.sum(), 1)（空 response 计 1 个 token），而 CP>1 路径按真实 unmasked token 计数（空 response 计 0）。同一份数据在 CP=1 与 CP>1 下会得到不同的 loss denominator，从而产生不同的 loss/gradient。
2. 除零风险：当整个 global batch 都没有有效 loss token（全部空/全 mask response）时，loss metric 的 denominator 为 0；当前代码直接做除法，可能产生 NaN/Inf 指标并掩盖 reducer 不一致。
3. optimizer 空转：全零 token step 的 loss 通过 0 * logits.sum() 零连接，梯度恒为 0；但 optimizer.step() 仍会通过 Adam momentum / weight decay 移动参数，并推进 LR scheduler，在没有任何训练信号的情况下改变模型状态。

这些是共享 Megatron 基础设施问题，不限于 RLOO 算法；本 PR 将其作为独立基础设施改动处理。

How

• token 计数（relax/backends/megatron/cp_utils.py）：
  • get_cp_local_num_tokens() 在 CP=1 时改为按真实 loss_mask.sum() 计数，空/全 mask response 贡献 0，与 CP>1 的 global valid-token denominator 语义一致。
• 指标归一化（relax/backends/megatron/loss.py、model.py）：
  • 新增 normalize_reduced_loss_metrics()：denominator 为 0 且所有 numerator 为 0 时按 no-signal step 报告全零指标；denominator 为 0 但存在非零 numerator 时 fail-fast（提示 reducer 不一致），不再静默除零。
• optimizer 语义（relax/backends/megatron/model.py）：
  • 新增 _is_global_zero_token_step()：在 pipeline last stage 上对 DP+CP 组 all-reduce 全 batch 有效 token 总数；PP=1 时所有 rank 均为 last stage、归约结果本地一致，直接本地判定；PP>1 时 last stage 沿 pipeline 组 broadcast 决策、非 last stage 加入 broadcast 完成 collective。各 rank 一致决定跳过 optimizer 与 LR scheduler 更新；跳过时 grad_norm=0.0 上报，不改变正常 step 路径的任何行为。
  • PP=1 不做 pipeline broadcast 是必须的：get_pipeline_model_parallel_group() 在 PP=1 拓扑下不含 global rank 0（relax 既有代码 hf_weight_iterator_direct.py 同样以 pp_size > 1 保护后才使用该组），无条件广播会导致 rank 0 ValueError + 其余 rank 卡在 optimizer all-reduce 的 collective 失配死锁（真实双卡训练首个 step 即复现，已修复并由回归测试锁定）。
• 空 response 切片（relax/backends/megatron/loss.py）：
  • get_responses() 对 response_length == 0 返回空切片（logits[0:0]/tokens[0:0]），避免 tokens[-0:] 取到完整 prompt 的边界错误。

Testing

• pre-commit run passes（本次修改文件：ruff、ruff-format、docformatter、check-conflict-markers 等全部通过）
• 新增/相关测试通过
• 文档已更新

验证摘要（本地，基于最新 main）：

• RLOO 相关 + zero-token 测试：65 passed。
• 新增测试覆盖：
  • _is_global_zero_token_step 纯函数：全零 vs 非零 token 的判定（monkeypatch 分布式原语）。
  • _is_global_zero_token_step 分布式契约回归：PP=1 断言 broadcast 零调用（all-reduce 后本地判定）；PP>1 断言 last stage 广播（src=pp_size-1）、非 last stage 加入广播。
• 正式双卡 60-rollout 配对实验（2×A100，RLOO → GRPO 同一 Ray head 串行完成）：
  • RLOO 60/60 steps、GRPO 60/60 steps，两臂 exit 0，同 Ray head（GCS/raylet PID 不变）验证通过；无 Global Restart、OOM、worker crash、非有限值。
  • 每 step 均真实执行 _is_global_zero_token_step（PP=1 路径），无 collective 失配；正常数据未误触发 zero-token 跳过（无 zero effective loss tokens 警告）。
  • 两臂耗时与既有基线一致（RLOO ~1h40m、GRPO ~1h38m），无性能回归。
  • train_one_step 集成：全局零 token step 断言 optimizer.step 与 LR scheduler 均未被调用、grad_norm == 0.0。
  • CP=2/4：token denominator 与 numerator 与未切分 reference 严格一致，包含空 response 用例。
  • 不等长、fully-masked、response_length == 0 的最终 scalar loss / gradient oracle，覆盖生产 loss_function() 的 reducer 与 Megatron token normalizer。
  • get_responses() CP=1 空 response 返回匹配的空 chunks。
  • 全零 loss-mask 返回零 normalizer 与零 numerator。
• 正常训练路径回归：非 zero-token 场景的行为与 main 完全一致（改动均为条件分支，默认路径零影响）。

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change
• Documentation update

Risk & Rollback

• 影响面：get_cp_local_num_tokens() 是共享函数，CP=1 下空/全 mask response 的计数从 1 变为 0，会影响所有算法在"存在空 response"时的 loss denominator；正常（非空）训练路径数值完全不变。
• 跨 rank 一致性：PP>1 时跳过决定依赖 pipeline broadcast，所有 rank 必须进入同一 collective；PP=1 时无 broadcast（所有 rank 参与同一 DP+CP all-reduce，本地判定一致）。
• 触发条件：全 batch 无有效 loss token 属于数据管线异常（正常数据不应出现）；本 PR 使该异常不再产生 NaN 指标与静默参数更新，并可通过日志 Training step %d has zero effective loss tokens globally 观测。
• 回退：单文件 revert 即可，不影响 【NO.28】实现 RLOO advantage estimator #205 的 RLOO 算法实现。

Checklist

• Diff 仅包含本任务必要改动
• 新增/相关测试全部通过
• 文档、脚本和默认值已更新
• 不含密钥、数据集、checkpoint 或机器隐私信息

[xcosmosbox]

All changes are ready, awaiting your review to merge @NINGBENZHE