feat(sc): support router replay with TQ - #3378
Conversation
github-actions
Bot
added
the
Documentation
zyzhou5
added
CI:Lfast
5da6705 to
695b5fc
Compare
github-actions
Bot
added
the
Documentation
|
/ok to test 695b5fc |
|
Awesome @zyzhou5 ! Thank you for the PR and all the experiments. LGTM. |
|
/ok to test 893c9fa |
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Static review of the SC + R3 wiring, run as a 6-agent pass (RL expert, bug finder, test reviewer, design reviewer, comment reviewer, plus adjudication). Overall LGTM — one real bug on an adjacent path, the rest are low-severity nits and process items. Nothing here should block merge.
The change is tightly scoped and mirrors the existing patterns well. Worth calling out explicitly:
tq_policy.py:521was the only one of five field-selection sites still hardcodinglist(DP_TRAIN_FIELDS)— the other four already usedfields_with_optional_routed_experts. That one-liner is the actual fix, and it's the right one.backfill_missing_routed_experts's "no-op when the batch carries no routes at all" is subtle and correct: it's what preserves the producer-side guard's ability to distinguish "replay off" from "capture failed". A naive always-fill would have silently defeated the guard this PR adds. The docstring says why, which is the right call.- The ordering comment at
payload.py:62-64is load-bearing and accurate — we verifiedextract_initial_prompt_messagesreally does reuse the same dicts. - Guard placement at
TQReplayBuffer.commitis the true choke point (both the native and Gym SC rollouts converge there before the TQ write) and it fires before any side effect.generate_and_pushremoves the reserved group onBaseExceptionand re-raises, so it can't strand a slot. - The new tests are Ray/GPU-free, land in the right L0 lane, and the partially-failed-group test encodes exactly the invariant the backfill exists for.
Checked and found clean, so nobody re-derives it: the recipe's defaults: chain resolves checkpointing.enabled to false (grpo_math_1B_megatron.yaml:31 overrides the true in grpo_math_1B.yaml), so the documented launch command works; routed_experts is in TOKEN_ALIGNED_FIELDS so zero-pad rows never reach _validate_replay_tensor; -1 is representable in all three signed carry dtypes; the R3 recipe sets enable_prefix_caching: false; SC + NeMo Gym fails loud rather than silently (nemo_gym.py:410 + the new commit guard); check_r3_trace.py accepts both new flags and hard-fails if trace_rollout_payload never fires; 10x8x60min = 80 GPU-h matches the 3480→3560 bump exactly and nothing references the renamed test function; all seven R3 env vars are documented.
Two things we deliberately did not raise, in case they come up: an early read that the recipe was missing checkpointing.enabled: false was wrong (disproved by the defaults: chain above), and a suggestion to tighten the nightly gate from < 1.02 to < 1.015 was dropped — it rested on eyeballing values off the screenshot and both sibling R3 nightlies use 1.02.
The convergence evidence holds up: the screenshot is a matched 4-arm study (SC-r3-on / SC-r3-off / legacy-r3-on / legacy-r3-off, 100 steps) where SC-r3-on tracks legacy-r3-on on both token_mult_prob_error and js_divergence_error while both r3-off arms separate cleanly — exactly the claim, and exactly what the guide's own "Effectiveness Check" asks for.
Not run: no tests were executed and no linter ran — uv refuses this host (the lockfile is sys_platform == 'linux' only), so this is entirely static review.
Generated by Claude Code
893c9fa to
b2b7d15
Compare
zyzhou5
changed the title
b2b7d15 to
4c4e96b
Compare
|
/ok to test 4c4e96b |
4c4e96b to
3fd854d
Compare
|
/ok to test 3fd854d |
3fd854d to
5b191ab
Compare
|
/ok to test 5b191ab |
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Backfill missing routed_experts for partial or failed generations before prompt flattening in SingleController payloads and the legacy sync and async drivers. Share the missing-route sentinel between the vLLM producer and Megatron consumer, and use the common data-plane field contract. Clear persistent nightly trace state, document runtime validation and fallback observability, add the recipe launch command, and direct legacy async TransferQueue users to the supported SingleController entrypoint. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
5b191ab to
d5be3d9
Compare
|
/ok to test d5be3d9 |
Adding `timeouts` as a required positional argument broke every direct construction of AsyncRolloutImpl. That went unnoticed because production always builds it through the config path, which passes the field -- the only direct construction is in tests. Upstream's new tests/unit/experience/test_rollout_manager_router_replay.py (from NVIDIA-NeMo#3378) is one, so after syncing onto current main it failed with TypeError: __init__() missing 1 required positional argument: 'timeouts'. The rebase was textually clean; the incompatibility is in the signature, not in any line either side edited, so nothing flagged it. RolloutTimeouts is a frozen dataclass whose fields all default to None, meaning "no deadline, wait indefinitely" -- the historical behaviour. So defaulting the parameter restores the old semantics for callers that do not ask for deadlines, rather than silently imposing one. Same reasoning already applied to RolloutRetryPolicy, whose defaults reproduce single-attempt behaviour so a directly-constructed manager does not silently gain retries. No separate regression test: upstream's file is the regression test, and it runs in our lane. Signed-off-by: asolergibert <asolergibert@nvidia.com>
3 participants
What does this PR do ?
This PR adds R3 router replay support to the SingleController native async GRPO + TQ path. It propagates vLLM routed-expert indices through async rollout message logs into the TQ training payload, validates that routes are present when router replay is enabled, and safely handles partially failed rollout groups by using the existing -1 fallback semantics. It also adds diagnostic tracing, unit coverage, an R3 SingleController recipe with a launch command, and a nightly end-to-end correctness test.
It closes #3327
Validation
Logs can be found in (Test 5) https://api.wandb.ai/links/nvidia-nemo-fw-public/6j8mkubs