You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[OMNIML-6061] Keep the expert-parallel GPU test's routing clear of near-ties - #3198
tests/gpu/torch/utils/test_expert_parallel_load.py compares a model loaded expert-parallel with the same model loaded whole: its logits, its gradients, and its quantized-training gradients. On two RTX 6000 Ada GPUs, the backward and QAT comparisons failed in layer 1 only, with gradients off by 3-16% (reported in #2742's review). The failure was the test's, not the load's:
One of 64 tokens took a different pair of experts in layer 1 in the sharded model than in the whole model. Layer 0 routed identically.
The tiny model's random router leaves near ties: a token's third-best score came within 0.0002 of its second. Ada's bf16 GEMMs put layer 0's output a little further from the whole model's than other GPUs do, enough to flip one.
A token routed to other experts sends its gradient to other weights, which no tolerance should absorb.
A router that decides clearly. The tiny model routes each token by a code its embedding carries: token t goes to expert t % 8 and one more its id picks. Every token's top two scores now clear the third by over 2 in both layers, and the top expert's weight stays at 0.45-0.62, so gradients still flow through the router.
Routing checked before gradients. The backward and QAT tests check that the sharded and whole models route every token to the same experts in every layer before they compare gradients. A flip fails as routing differs: layer L sends N of M tokens to other experts than the whole model, not as the gradient errors it causes.
The 3e-2 gradient bound and the transformers floor are unchanged.
Usage
N/A: a test-only change.
Testing
The router's margin, from a CPU probe of the tiny model on the test's batches:
Router
Smallest gap between the 2nd and 3rd score
Top expert's weight
Random (before)
0.0000 (layer 0), 0.0002 (layer 1)
0.13-0.19
Clear (this PR)
2.28 (layer 0), 2.11 (layer 1)
0.45-0.62
The GPU test file on 2 RTX 6000 Ada GPUs, with torch 2.11.0 (CUDA 13.0, NCCL 2.28.9), on transformers 5.9.0 and 5.18.0 alike:
2: backward, routing differs: layer 1 sends 1 of 32 tokens to other experts than the whole model
With this PR, the largest gradient error on Ada is 0.0067 for the backward cases, 0.0107 for fp8 and 0.0081 for nvfp4, against the 3e-2 bound.
The skipped cases need 4 GPUs.
On Ada the fp8 cases ran without ModelOpt's FP8 CUDA extension, which the host's CUDA 13.3 toolchain could not build against torch 2.11.
GB200: pending; added here when it lands.
Pre-commit, mypy included, passes against main.
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).
Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).
Is this change backward compatible?: ✅ Test only.
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
Did you write any new necessary tests?: ✅ The routing check.
On two RTX 6000 Ada GPUs, the backward test failed in layer 1 only, with
the same numbers on transformers 5.9.0 and 5.18.0 and on main's own
replicated path: one of 64 tokens took a different pair of experts in the
sharded model than in the whole model. The tiny model's random router
leaves near ties (a third-best score within 0.0002 of the second), and the
GPU's bf16 rounding flips one. A token routed elsewhere sends its gradient
to other weights, which no tolerance should absorb.
The tiny model now routes each token by a code its embedding carries, so
every token's top two scores clear the third by over 2 in both layers.
Before the gradient comparisons, the backward and QAT tests check that the
sharded and whole models route every token alike, and a flip fails as
"routing differs" with its count. The 3e-2 bound is unchanged. Test only.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
No actionable comments were generated in the recent review.
Walkthrough
The test fixture now sets deterministic routing for the tiny Qwen3-MoE model. New helpers capture and compare expert routes. Backward and quantized-training workers check route equality before comparing gradients.
The fixture sets token embeddings and layer gate weights for deterministic routing before saving the model.
Route comparison before backward tests/gpu/torch/utils/test_expert_parallel_load.py
New helpers capture sorted expert IDs and compare routes. Backward and quantized-training workers check routes before backward and gradient comparison.
The change adds routing checks to the GPU gradient test. No concrete merge-blocking risk was established.
Pre-merge checks: 5 · 1
Failed checks: 1 warning
Check name
Status
Explanation
Resolution
Docstring Coverage
Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…
Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed checks: 5 passed
Check name
Status
Explanation
Description Check
Check skipped - CodeRabbit’s high-level summary is enabled.
Title check
The title names the main change: the expert-parallel GPU test's router is made to avoid near-tie routing decisions. The changeset replaces the random router in the tiny Qwen3-MoE fixture with a determ…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns
The PR changes only tests/gpu/torch/utils/test_expert_parallel_load.py, not the modelopt package, examples, dependency manifests, or requirements files. The added code contains no `torch.load(...,…
Full details: Title check
Explanation
The title names the main change: the expert-parallel GPU test's router is made to avoid near-tie routing decisions. The changeset replaces the random router in the tiny Qwen3-MoE fixture with a deterministic one whose top scores are well separated, and adds routing-equality checks before gradient comparisons. The title is specific, uses the ticket ID, and is one sentence of 10 words.
Full details: Security Anti-Patterns
Explanation
The PR changes only tests/gpu/torch/utils/test_expert_parallel_load.py, not the modelopt package, examples, dependency manifests, or requirements files. The added code contains no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval/exec, or # nosec bypass.
Full details: Docstring Coverage
Explanation
Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 5 functions across 1 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Finishing Touches
Generate docstrings
Commit to this branch
Create a new PR
Generate unit tests (beta)
Commit to this branch
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
CodeRabbit Verify
Run CodeRabbit Verify to test the changes in this pull request.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 7 remain after this review.
✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.41%. Comparing base (b002695) to head (1425b0b). ⚠️ Report is 2 commits behind head on main.
The reason will be displayed to describe this comment to others. Learn more.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM; the fixture change removes unstable routing without loosening assertions, and the justified test edit needs human sign-off.
Needs action:
Sign off on the fixture adjustment in tests/gpu/torch/utils/test_expert_parallel_load.py: real expert dispatch and router backpropagation remain exercised, with existing logits, calibration, and gradient bounds unchanged.
No action needed:
Routing checks preserve global expert IDs, remove hooks in finally, and gather mismatches before any rank starts backward.
Static review only; GPU runs and numerical probes were not reproduced. Guidance used: supplied main baseline; its commit SHA was unavailable.
Preview removed because the pull request was closed.
2026-11-03 21:07 UTC
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
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
meenchen
left a comment
What does this PR do?
Type of change: Bug fix (test only)
tests/gpu/torch/utils/test_expert_parallel_load.pycompares a model loaded expert-parallel with the same model loaded whole: its logits, its gradients, and its quantized-training gradients. On two RTX 6000 Ada GPUs, the backward and QAT comparisons failed in layer 1 only, with gradients off by 3-16% (reported in #2742's review). The failure was the test's, not the load's:This PR changes the test only:
tgoes to expertt % 8and one more its id picks. Every token's top two scores now clear the third by over 2 in both layers, and the top expert's weight stays at 0.45-0.62, so gradients still flow through the router.routing differs: layer L sends N of M tokens to other experts than the whole model, not as the gradient errors it causes.Usage
N/A: a test-only change.
Testing
The router's margin, from a CPU probe of the tiny model on the test's batches:
The GPU test file on 2 RTX 6000 Ada GPUs, with torch 2.11.0 (CUDA 13.0, NCCL 2.28.9), on transformers 5.9.0 and 5.18.0 alike:
-k denserouting differs: layer 1 sends 1 of 32 tokens to other experts than the whole modelBefore your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
A follow-up to #2742, whose review asked about the Ada failures. Part of expert-parallel PTQ for Hugging Face MoE models.
Summary by CodeRabbit