[Cohort-1 Task 23] 跳过零 advantage microbatch 的 backward

[xiaoh1024]

背景

GRPO 中，同组样本 reward 相同时，归一化后的 advantage 全为零。此时对应 microbatch 的 policy loss 和参数梯度均为零，但当前 Megatron 训练路径仍会构建并遍历完整 backward graph。

计划增加默认关闭的 --skip-zero-advantage-backward：保留完整 forward、loss 计算和指标统计；仅当 microbatch 中所有有效 response token 的 advantage 都为零，且完整 policy loss 经运行时确认是有限的零值时，跳过模型 backward graph。

代码基线

• 上游仓库：https://github.com/redai-infra/Relax
• 开发仓库：https://github.com/xiaoh1024/Relax
• 基线 commit：5b2301110db97f38234402dc90e70ae6fb063cde
• 开发分支：perf/task23-zero-advantage-backward

环境与配置

• 镜像：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 内存
• 模型：Qwen/Qwen3-4B@1cfa9a7208912126459214e8b04321603b3df60c
• 训练数据：zhuzilin/dapo-math-17k@2e65612930298bde4c5d58fd97b3f23a483aaff9
• 评测数据：zhuzilin/aime-2024@1c625e328db94ec7ef7ff169016b097c468d60b9
• Recipe：scripts/training/text/run-qwen3-4B-8xgpu.sh
• GRPO；training TP2/DP4；rollout TP8；rollout batch size 32；n=8；global batch size 128；max response length 8192；temperature 1；seed 1234

基线结果

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

Run	Total tokens/s	Response tokens/s	Step time	Actor train time	Train time	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
Mean	9,556.72	9,339.15	178.37 s	67.28 s	80.10 s	87.59%	71.37 GiB

改动范围

• relax/backends/megatron/model.py::train_one_step：识别有效 advantage 全零的 microbatch，并以 no_grad 执行其 forward。
• relax/backends/megatron/loss.py::loss_function：校验完整 policy loss 为有限零值，并向 Megatron schedule 返回独立标量以保持调度语义。
• relax/backends/megatron/arguments.py、relax/utils/arguments.py：增加 opt-in 参数并拒绝未经验证的组合。
• relax/backends/megatron/actor.py、relax/utils/training/train_metric_utils.py：记录跳过 token 数、命中率，并按实际执行工作量修正训练 FLOPs 统计。
• 单元测试覆盖判定、零/非零梯度、loss 安全检查、参数约束和性能指标。

支持范围为同步 colocate、Megatron bridge、纯文本 Dense policy loss、PP1/CP1/EP1、非 LoRA，且不启用额外 entropy/KL loss。

运行命令

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"

验收口径

• perf/step_token_per_s 稳定窗口均值提升至少 5%，并报告 actor_train_time、train_time、GPU 利用率和峰值显存。
• baseline 与 candidate 使用相同模型、数据、硬件、batch、序列长度、采样设置和统计窗口，各完成两次运行。
• 有效样本和 token 口径一致；loss、reward、grad norm 正常；无 OOM、NaN、Inf 或数据丢失。
• 默认行为不变；单元测试覆盖跳过路径、正常 backward 路径和不支持配置的 fail-fast。

交付方式：Draft PR，包含实现、测试、验证命令和原始结果。

请维护者确认该优化方向、支持范围和验收口径是否合适。