feat(megatron): handle zero-token no-signal steps in the shared trainer - #285
feat(megatron): handle zero-token no-signal steps in the shared trainer#285xcosmosbox wants to merge 4 commits into
Conversation
…ection The original implementation unconditionally broadcast the zero-token decision over the pipeline-model-parallel group, but that group does not contain global rank 0 in PP=1 topologies, causing a collective mismatch deadlock (rank 0 ValueError + other ranks stuck in optimizer all-reduce) that tripped the global restart path. With PP=1 every rank is the last pipeline stage and joins the same DP+CP all-reduce, so the local count is already consistent and no broadcast is needed. Add regression tests for PP=1 (no broadcast) and PP>1 (last stage broadcasts, non-last stages join).
|
All changes are ready, awaiting your review to merge @NINGBENZHE |
|
@codex review following the repository AGENTS.md and skills/code-review/SKILL.md |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8cad9cb47f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| torch.distributed.broadcast( | ||
| signal, | ||
| src=pp_size - 1, | ||
| group=mpu.get_pipeline_model_parallel_group(), |
There was a problem hiding this comment.
When PP is combined with TP or DP, pp_size - 1 is only the pipeline-local stage index and is generally not the global rank of the last member of each pipeline group; torch.distributed.broadcast(..., src=...) interprets src as a global rank. Consequently, most pipeline groups either reject this source as a non-member or wait for a rank that never participates, hanging every PP>1 training step. Resolve the group's last rank with dist.get_global_rank(pp_group, pp_size - 1) or Megatron's pipeline-last-rank helper.
Useful? React with 👍 / 👎.
| num_tokens_local = sum(x["values"][0] for x in losses_reduced) # type: ignore[index] | ||
| torch.distributed.all_reduce( | ||
| num_tokens_local, | ||
| group=mpu.get_data_parallel_group(with_context_parallel=True), |
There was a problem hiding this comment.
When --calculate-per-token-loss is disabled, loss_function() stores num_samples rather than the token count in values[0]. A genuine fully masked batch therefore yields a positive value here, so this helper returns false and the optimizer and scheduler still advance with zero gradients—the exact parameter drift this change is intended to prevent. This affects the default response-mean mode used by many existing training scripts; pass an explicit effective-token count instead of reusing the metric denominator.
Useful? React with 👍 / 👎.
| num_tokens_local, | ||
| group=mpu.get_data_parallel_group(with_context_parallel=True), | ||
| ) | ||
| signal[0] = 1 if num_tokens_local.item() == 0 else 0 |
There was a problem hiding this comment.
This helper is called on every training step, and num_tokens_local.item() synchronizes the CUDA stream on every last-stage rank; signal.item() then introduces another host synchronization on every rank before returning. This adds a pipeline-wide stall to the hot training path even for ordinary nonempty batches. Keep the zero-test tensor-side and avoid redundant host materialization in accordance with the repository's hot-path rule.
AGENTS.md reference: AGENTS.md:L38-L40
Useful? React with 👍 / 👎.
| "Training step %d has zero effective loss tokens globally; skipping optimizer and LR scheduler updates.", | ||
| step_id, | ||
| ) | ||
| update_successful = True |
There was a problem hiding this comment.
On a zero-token critic step no optimizer update occurs, but setting update_successful = True causes maybe_verify_critic_value_head_movement() to count it as an eligible successful update. After several consecutive zero-token critic batches, the runtime check emits a false value-head warning and marks itself verified, so it will not validate a later real update. Preserve a distinct skipped state or pass False to the movement checks for this branch.
Useful? React with 👍 / 👎.
2 participants
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"存在两处语义不一致与一个数值风险：
get_cp_local_num_tokens()在 CP=1 时对每个样本使用历史 per-sampleclamp_min(loss_mask.sum(), 1)（空 response 计 1 个 token），而 CP>1 路径按真实 unmasked token 计数（空 response 计 0）。同一份数据在 CP=1 与 CP>1 下会得到不同的 loss denominator，从而产生不同的 loss/gradient。0 * logits.sum()零连接，梯度恒为 0；但 optimizer.step() 仍会通过 Adam momentum / weight decay 移动参数，并推进 LR scheduler，在没有任何训练信号的情况下改变模型状态。这些是共享 Megatron 基础设施问题，不限于 RLOO 算法；本 PR 将其作为独立基础设施改动处理。
How
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 不一致），不再静默除零。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 路径的任何行为。get_pipeline_model_parallel_group()在 PP=1 拓扑下不含 global rank 0（relax 既有代码hf_weight_iterator_direct.py同样以pp_size > 1保护后才使用该组），无条件广播会导致 rank 0ValueError+ 其余 rank 卡在 optimizer all-reduce 的 collective 失配死锁（真实双卡训练首个 step 即复现，已修复并由回归测试锁定）。relax/backends/megatron/loss.py）：get_responses()对response_length == 0返回空切片（logits[0:0]/tokens[0:0]），避免tokens[-0:]取到完整 prompt 的边界错误。Testing
pre-commit runpasses（本次修改文件：ruff、ruff-format、docformatter、check-conflict-markers 等全部通过）验证摘要（本地，基于最新 main）：
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 加入广播。_is_global_zero_token_step（PP=1 路径），无 collective 失配；正常数据未误触发 zero-token 跳过（无zero effective loss tokens警告）。train_one_step集成：全局零 token step 断言 optimizer.step 与 LR scheduler 均未被调用、grad_norm == 0.0。response_length == 0的最终 scalar loss / gradient oracle，覆盖生产loss_function()的 reducer 与 Megatron token normalizer。get_responses()CP=1 空 response 返回匹配的空 chunks。Type of Change
Risk & Rollback
get_cp_local_num_tokens()是共享函数，CP=1 下空/全 mask response 的计数从1变为0，会影响所有算法在"存在空 response"时的 loss denominator；正常（非空）训练路径数值完全不变。Training step %d has zero effective loss tokens globally观测。Checklist