Repository navigation
fix(router-replay): allow routeless batches and token-free groups - #4900
Open
linj-glitch wants to merge 3 commits into
Open
linj-glitch wants to merge 3 commits into
linj-glitch wants to merge 3 commits into
Conversation
linj-glitch
changed the title
linj-glitch
added
the
CI:Lfast
Links 4585549-51 died at the first prev-logprob fetch: the batch mixed rows restored from the pre-R3 step_102 checkpoint and a token-free group, so the routed_experts column was absent and the worker raised. Treat that as the documented missing-route fallback for the whole batch (log it) rather than a schema failure. Test updated. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit 568ca7f5a5f24323492d4d1716579d23451e4864) Signed-off-by: Lin Jia <linj@nvidia.com>
…ively instead of raising With policy.router_replay.enabled=true three guards raised RuntimeError for a payload without routed_experts and took the run down: - TQReplayBuffer.commit: a prompt group whose every completion is token-free or masked (every sandbox failed before the model spoke) has no generation to route. It now commits with a one-line log; groups with trainable tokens still raise, listing the payload fields. BatchedDataDict is tensordict-backed, so the guard reads fields with .get, not 'in'. - SyncRolloutActor: the same check on the synchronous path, keyed on token_loss_mask. - The Megatron worker's _should_use_router_replay raised for a fetched batch without routes (rows restored from a checkpoint written before Router Replay was enabled); it now logs and uses the native router for that batch, as the documented per-token missing-route fallback does. Tests: test_tq_replay_buffer.py (token-free group commits), test_router_replay.py (worker fallback). Signed-off-by: Lin Jia <linj@nvidia.com>
linj-glitch
force-pushed
the
linj/router-replay-routeless
branch
from
October 28, 2026 23:41
2cd8623 to
0ea9540
Compare
Contributor
Author
|
/ok to test 0ea9540 |
Integrate upstream main for the required CI freshness gate. Match the token-free regression fixture to the upstream tensorization keyword contract. Signed-off-by: Lin Jia <linj@nvidia.com>
Contributor
Author
|
/ok to test d8620aa |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Current integration and validation (2026-10-29 UTC)
Merged upstream main
feb735f6132179d62634f515b2727952efb0dd14into this branch with a signed additive commit,d8620aa4eb2c399cad127ee36815f61b76c54b3d, to satisfy the required CI freshness gate. The production patch retains exactly its original added and removed lines. An independent integration review found no new production issues.tests/:pytest unit/single_controller/test_tq_replay_buffer.py unit/experience/test_sync_rollout_actor.py unit/models/megatron/test_router_replay.py -p no:cacheprovider -o 'addopts=-rA -q'.pre-commit run --all-files: formatting, Ruff, TOML, and recipe checks passed. Pyrefly reports 44 errors, with identical error messages, source paths, and multiplicities on pristine upstream main; there are no new diagnostics from this PR. The overall hook command therefore exits 1 for the baseline type errors.git diff --checkpassed; the commit contains current main and a DCO sign-off. Validation used the existing Python 3.13 environment and a private CPU-only Ray cluster, with module resolution verified against this PR checkout.The default CPU replay-buffer, sync-actor, and router-replay selection passed. Megatron-marked cases require the unavailable mcore stack and were not selected. The existing token-free fixture now accepts the
per_token_multimodal_fieldskeyword introduced on main; production behavior is unchanged by the integration.Repository CI must validate this exact head; earlier-head results do not establish that it passes. Existing reviews have been preserved.
What does this PR do ?
With
policy.router_replay.enabled=true, three places raiseRuntimeErrorwhen a batch or group carries norouted_experts, and take the whole run down. Two kinds of batch are legitimately routeless and are now accepted: a prompt group whose every completion is token-free or masked (nothing was generated, so there is nothing to route), and a fetched training/logprob batch without the field (rows restored from a checkpoint written before Router Replay was enabled). Megatron routes those tokens natively, exactly as the documented missing-route fallback does per token.Issues
None.
Usage
No configuration change.
Before your PR is "Ready for review"
Pre checks:
Additional Information
Where it failed.
TQReplayBuffer.commit(replay_buffer.py): an agentic group in which every sandbox failed before the model spoke is committed as masked placeholders.pack_payloadproduces norouted_expertsfor it, and the guard raised at the first such group.SyncRolloutActor(sync_rollout_actor.py): the same check on the synchronous path._should_use_router_replay(megatron_policy_worker.py): a batch restored from a pre-Router-Replay checkpoint has no routes; the worker raised instead of training on native routing.Fix. A group counts as token-free when its
token_masksums to zero or every row hasmask_sampleset (masked placeholders can still carry assistant tokens). Such groups commit with a one-line log instead of raising; groups with trainable tokens still raise, now listing the payload fields. The worker falls back to native routing with a log line.BatchedDataDictis tensordict-backed, so the guard reads fields with.getrather thanin.Tests.
tests/unit/single_controller/test_tq_replay_buffer.py::test_commit_allows_a_token_free_group_without_routed_experts(new);tests/unit/models/megatron/test_router_replay.pyupdated for the worker fallback. Both files pass locally.🤖 Generated with Claude Code