perf(megatron): skip zero-advantage backward

[xiaoh1024]

Summary

• Closes [Cohort-1 Task 23] 跳过零 advantage microbatch 的 backward #252
• 解决了什么问题：GRPO 中 advantage 全零的 microbatch 不产生参数梯度，但 Megatron 路径仍会构建并遍历完整 backward graph。
• 为什么采用该方案：保留 forward、loss 和指标，只省略经严格判定为零梯度的 backward graph；功能默认关闭，并对支持配置做 fail-fast 校验。

Changes

• relax/backends/megatron/model.py
  • 仅当所有有效 response token 的 advantage 都为零时，以 no_grad 执行该 microbatch 的模型 forward。
• relax/backends/megatron/loss.py
  • 在跳过前确认完整 policy loss 是有限零值。
  • 返回独立可求导标量，保持 Megatron forward/backward schedule 和指标归约行为。
• relax/backends/megatron/arguments.py、relax/utils/arguments.py
  • 新增默认关闭的 --skip-zero-advantage-backward。
  • 对当前未验证的训练模式、loss、并行和模型组合做 fail-fast 校验。
• relax/backends/megatron/actor.py、relax/utils/training/train_metric_utils.py
  • 汇总跳过 token 数与比例，并按实际执行的 forward/backward 工作量修正训练 FLOPs 统计。
• tests/backends/megatron/test_zero_advantage_backward.py、tests/utils/test_train_metric_utils.py
  • 覆盖零/非零 advantage 判定、跳过与正常梯度、loss 安全检查、配置约束和指标计算。
• CLI/兼容性：默认行为不变；当前支持同步 colocate、Megatron bridge、纯文本 Dense policy loss、PP1/CP1/EP1、非 LoRA，且不启用额外 entropy/KL loss。

Verification

• 环境、硬件、commit：

  • Base：5b2301110db97f38234402dc90e70ae6fb063cde
  • Candidate：39ddc2d2b661df9e083556161ba7208e406e0514
  • Image：ghcr.io/redai-infra/relaxrl:dev-20260715-8325919e@sha256:3fa8ce578acda6c829b83016bde42c38fa892681e4f36ca330f545616fe578e2
  • Python 3.12.3；PyTorch 2.11.0+cu129；CUDA 12.9；SGLang 0.5.12.post1；Ray 2.56.0；Transformers 5.6.0
  • 8 × NVIDIA L20Z 80 GiB；128 vCPU；1024 GiB memory
  • Model：Qwen/Qwen3-4B@1cfa9a7208912126459214e8b04321603b3df60c
  • Train data：zhuzilin/dapo-math-17k@2e65612930298bde4c5d58fd97b3f23a483aaff9
  • Eval data：zhuzilin/aime-2024@1c625e328db94ec7ef7ff169016b097c468d60b9

• 可复制命令：

  export MODEL_DIR=/path/to/models
  export DATA_DIR=/path/to/data
  export NUM_ROLLOUT=21
  
  # Baseline
  export EXP_DIR=/path/to/output/baseline
  bash scripts/training/text/run-qwen3-4B-8xgpu.sh
  
  # Candidate：同一 recipe，仅在临时副本的 SGLANG_ARGS 中加入开关
  export EXP_DIR=/path/to/output/candidate
  candidate_recipe=$(mktemp scripts/training/text/.run-qwen3-4B-zero.XXXXXX.sh)
  trap 'rm -f "$candidate_recipe"' EXIT
  cp scripts/training/text/run-qwen3-4B-8xgpu.sh "$candidate_recipe"
  sed -i '/--rollout-num-gpus-per-engine 8/a\   --skip-zero-advantage-backward' "$candidate_recipe"
  bash "$candidate_recipe"

• 单元测试：

  python -m pytest -q \
    tests/backends/megatron/test_zero_advantage_backward.py \
    tests/utils/test_train_metric_utils.py
  

• 端到端结果：每次运行统计 rollout 1–19；rollout 0 含启动开销，rollout 20 含评测，均不计入主指标。

  Run	Total tokens/s	Response tokens/s	Step time	Actor train time	Train time	Skip token fraction	GPU util	Peak VRAM
  Baseline 1	9,548.35	9,328.51	177.45 s	66.86 s	79.56 s	—	87.50%	69.70 GiB
  Baseline 2	9,565.10	9,349.78	179.29 s	67.70 s	80.63 s	—	87.67%	73.04 GiB
  Candidate 1	10,997.76	10,743.34	153.55 s	42.50 s	55.09 s	59.70%	84.72%	71.71 GiB
  Candidate 2	11,008.76	10,760.74	156.19 s	44.07 s	56.95 s	62.09%	84.59%	71.03 GiB
  Mean / delta	9,556.72 → 11,003.26 (+15.14%)	9,339.15 → 10,752.04 (+15.13%)	178.37 → 154.87 s (−13.18%)	67.28 → 43.28 s (−35.67%)	80.10 → 56.02 s (−30.06%)	60.89%	87.59% → 84.65% (−2.93 pp)	71.37 → 71.37 GiB

  两次独立运行的吞吐均值标准差为 baseline 11.84 tokens/s、candidate 7.78 tokens/s（CV 0.12% / 0.07%）。GPU 利用率下降是减少 backward 计算后的预期结果：rollout 与角色切换等未改变阶段在更短的 step 中占比上升。模型、样本和有效 token 口径不变；仅省略经校验不会产生参数梯度的 backward。峰值显存按两次运行各自峰值的均值报告，基本不变。

• 正确性/质量护栏：

  Metric	Baseline	Candidate
  rollout/raw_reward	−0.6135	−0.5676
  rollout/response_lengths	6,514.22	6,498.04
  train/loss	0.011276	0.011337
  train/grad_norm	0.099801	0.097318
  AIME accuracy (240 samples/run)	0.2750	0.3917
  NaN / Inf / OOM	none	none

  AIME 仅作为未退化护栏；不把两次随机评测的均值差异解释为本优化带来的质量收益。四次运行均完成 21/21 rollouts、240-sample AIME eval 和 iteration-20 checkpoint。

Risk & Rollback

• 已知限制：当前仅支持同步 colocate、Megatron bridge、纯文本 Dense policy loss、PP1/CP1/EP1、非 LoRA，且不启用额外 entropy/KL loss；其他组合在参数校验阶段报错。
• 风险：若零 advantage 判定遗漏其他 loss 项，可能错误跳过有效梯度。实现限制支持范围，并在运行时确认完整 loss 是有限零值后才跳过。
• 关闭开关或回退方式：该功能默认关闭；不传 --skip-zero-advantage-backward 即恢复原有路径。

Checklist

• Diff 仅包含本任务必要改动
• 新增/相关测试全部通过
• CLI 帮助和默认值已更新
• 不含密钥、数据集、checkpoint 或机器隐私信息
• 已逐条回复 review comment

[Copilot]

Pull request overview

This PR introduces an opt-in Megatron training optimization to skip building/traversing the backward graph for microbatches whose effective advantages are all zero, while preserving forward/loss/metrics semantics and adding fail-fast validation for supported configurations.

Changes:

• Add --skip-zero-advantage-backward (default off) plus strict configuration validation in Megatron args.
• Detect zero-effective-advantage microbatches in the Megatron forward step and avoid autograd graph construction; enforce a runtime “complete loss is finite and zero” guard when skipping.
• Track skipped-backward token stats and adjust FLOPs/MFU reporting accordingly; add unit tests for skip logic, validation, and perf metrics.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
relax/backends/megatron/model.py	Detect all-zero effective advantages and run model forward under no_grad for those microbatches; record skipped seq lens.
relax/backends/megatron/loss.py	Add skip marker key, log skip metric weight, and replace skipped loss with a detached differentiable scalar after validating it’s finite zero.
relax/backends/megatron/arguments.py	Add fail-fast validation restricting the feature to a known-safe Megatron configuration subset.
relax/backends/megatron/actor.py	Gather and aggregate per-rank skipped seq lens for downstream perf logging.
relax/utils/training/train_metric_utils.py	Subtract skipped-backward FLOPs from actor-train TFLOPs/MFU and log skipped token counts/fraction.
relax/utils/arguments.py	Add CLI flag --skip-zero-advantage-backward (default disabled).
tests/backends/megatron/test_zero_advantage_backward.py	New tests covering zero-advantage detection, validation constraints, and loss/grad behavior for skipped vs non-skipped paths.
tests/utils/test_train_metric_utils.py	Add test ensuring FLOPs/MFU excludes skipped-backward work and logs skipped token metrics.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[xiaoliang0601]

✅ 验收通过，但是不合入。

还有几个问题，希望修下：

1. 实现的有点 ad-hoc，如果只支持 GRPO，建议入口参数阶段明确拒绝 critic/PPO，只保留 GRPO。
2. text-only 的 fail-fast 疑似实际会失效，因为 validator 检查 args.is_vl_model，但该字段直到 Actor 初始化、参数校验结束后才赋值。因此 VLM 会错误通过“不支持配置”检查。请再次确认这个逻辑。
3. count_nonzero(...).item() 仍在每个 microbatch 强制 GPU→CPU 同步；命中 skip 后 loss 又有第二次 .item()。这违反仓库“热路径禁止 GPU-CPU 同步”的硬规则。应在 rollout/microbatch 规划阶段批量生成 CPU skip metadata，而不是 forward 前逐批同步。

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[xiaoh1024]

✅ 验收通过，但是不合入。

还有几个问题，希望修下：

1. 实现的有点 ad-hoc，如果只支持 GRPO，建议入口参数阶段明确拒绝 critic/PPO，只保留 GRPO。
2. text-only 的 fail-fast 疑似实际会失效，因为 validator 检查 args.is_vl_model，但该字段直到 Actor 初始化、参数校验结束后才赋值。因此 VLM 会错误通过“不支持配置”检查。请再次确认这个逻辑。
3. count_nonzero(...).item() 仍在每个 microbatch 强制 GPU→CPU 同步；命中 skip 后 loss 又有第二次 .item()。这违反仓库“热路径禁止 GPU-CPU 同步”的硬规则。应在 rollout/microbatch 规划阶段批量生成 CPU skip metadata，而不是 forward 前逐批同步。

感谢 review，几个问题已经修复：

1. 参数校验现在明确限定为 GRPO，并拒绝 critic、PPO 及其他未验证配置。
2. VLM 检查已前移至 HF config 校验阶段，并在模型物化后再次检查，避免校验被绕过
3. 已移除热路径中的 GPU–CPU 同步。现在统一在 rollout 规划阶段生成 CPU metadata，并由 DataIterator 按实际 batch 调度传递

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.