Repository navigation
Conversation
Signed-off-by: wiederse <1477618732@qq.com>
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes.
 |
| Layer / File(s) | Summary |
|---|---|
Scale gradients and validate distributed loss modelopt/torch/distill/plugins/megatron.py, tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py, CHANGELOG.rst |
For multiple tensor-parallel ranks, the loss adjustment scales gradients by the inverse of the tensor-parallel size while preserving the reported loss. Distributed tests compare shard losses and gradients to a full-vocabulary reference across top-P, reverse-KL, and temperature settings. The changelog records the fix. |
Priority:
Estimated code review effort: 2 (Simple) | ~15 minutes
Change: Bug fix
Suggested reviewers: aanoosheh
Merge Risk: 
5ec73
-
The change appears mergeable after normal checks. The reported distributed tests cover the intended gradient behavior, though they were not run independently here.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files. (1 skipped: 1… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title names the main change: Top-K KL distillation under tensor parallelism now scales gradients by 1/TP, which the raw summary and PR objectives both describe. It is specific and concise, and it … 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 PASS. The pull request changes only CHANGELOG.rst,modelopt/torch/distill/plugins/megatron.py, and a Megatron test. The added Python code contains notorch.load(..., weights_only=False), `numpy.…Full details: Title check
-
Explanation
The title names the main change: Top-K KL distillation under tensor parallelism now scales gradients by 1/TP, which the raw summary and PR objectives both describe. It is specific and concise, and it fits the changelog entry, the megatron.py loss change, and the updated test. No unrelated changes were reported.
Full details: Security Anti-Patterns
-
Explanation
PASS. The pull request changes only
CHANGELOG.rst,modelopt/torch/distill/plugins/megatron.py, and a Megatron test. The added Python code contains notorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or added# noseccomments. Nopyproject.tomlor requirements file changed, so no new dependency licensing issue applies.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: backward-only scaling corrects replicated TP gradients without changing forward losses; the justified reference-test rewrite needs human sign-off.
Needs action:
- Confirm the test rewrite in
tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py: it retains forward and probability-mass checks while adding shard-gradient comparisons, Top-P truncation, temperature scaling, and reverse KL.
|
/claude review |
|
/ok to test 5ec737a |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| parallel_state.initialize_model_parallel(tensor_model_parallel_size=size) | ||
| cfg = SimpleNamespace(tensor_model_parallel_size=size) | ||
| device = torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu") | ||
| student, teacher = _make_loss_inputs() |
There was a problem hiding this comment.
[SUGGESTION] _make_loss_inputs() defaults to vocab=16, and size here is torch.cuda.device_count() (via dist_workers), so the sharding below only behaves well when the GPU count divides 16.
torch.chunk(16, n, dim=-1) is not "split into n equal parts" — it is "chunks of ceil(16/n)":
n=5→[4,4,4,4], only 4 chunks, so.chunk(size, -1)[rank]raisesIndexErroron rank 4.n=6→[3,3,3,3,3,1], solocal_top_k = min(k, shard)is 3 on ranks 0-4 but 1 on rank 5, and thedist_nn.functional.all_gatherinsideTopLogitsKLLossgets mismatched shapes across ranks.n=3→[6,6,4], which runs but exercises an unequal vocab sharding Megatron never produces.
The PR was validated at TP=1/2/4 (and 8 divides 16), so this is latent rather than currently failing — but the failure mode on an odd runner is an opaque IndexError/NCCL shape error rather than a clear skip. Cheapest fix is to size the vocab off size so shards are always equal and >= k:
student, teacher = _make_loss_inputs(vocab=4 * size)(or keep vocab=16 and add assert 16 % size == 0).
| ): | ||
| parallel_state.initialize_model_parallel(tensor_model_parallel_size=size) | ||
| cfg = SimpleNamespace(tensor_model_parallel_size=size) | ||
| device = torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu") |
There was a problem hiding this comment.
[SUGGESTION] The torch.device("cpu") fallback is unreachable here. This helper only ever runs through dist_workers, whose pool is built with backend="nccl" and world_size=torch.cuda.device_count() (tests/gpu_megatron/conftest.py), so CUDA is always available. This looks like a leftover from the temporary Gloo TP=1/2/4 fixture mentioned in the PR description; CONTRIBUTING's "Remove dead code" applies.
| device = torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu") | |
| device = torch.device("cuda", rank) |
| @pytest.mark.parametrize( | ||
| ("top_p", "reverse", "temperature"), | ||
| [(None, False, 1.0), (0.5, False, 2.0), (0.5, True, 1.0)], | ||
| ) | ||
| def test_topk_logits_kl_loss_numerics_ghost_token_reference( | ||
| dist_workers, top_p, reverse, temperature | ||
| ): |
There was a problem hiding this comment.
[SUGGESTION] This is the only regression test for the new tp_size > 1 branch, but dist_workers sizes the pool to torch.cuda.device_count(). On a single-GPU runner size == 1, the tp_size > 1 branch in TopLogitsKLLoss.forward is never entered, and all three parametrizations pass vacuously — green CI without having tested the fix. (The top_p/reverse/temperature numerics they also assert are already covered single-process by test_logits_kl_losses_temperature_scaling and test_topk_logits_kl_loss_top_p_masks_tail.)
dist_workers_size_2 / dist_workers_size_4 exist for exactly this case — they pytest.skip("Need at least N GPUs") rather than silently degrade, so a 1-GPU run reports "skipped" instead of "passed". Consider requesting dist_workers_size_2 (keeping dist_workers as an additional case if you also want TP=device_count coverage), so the branch this PR fixes can never be reported as passing untested.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2988 +/- ##
==========================================
- Coverage 79.82% 79.11% -0.71%
==========================================
Files 661 661
Lines 74928 74930 +2
==========================================
- Hits 59810 59284 -526
- Misses 15118 15646 +528
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Claude review — PR #2988Scope: full review (trigger comment was a bare Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3 (all test-only, all posted inline) The fix is correct — I verified the
|
What does this PR do?
Type of change: Bug fix
With Top-K logits distillation and tensor parallelism enabled, each TP rank evaluates the same global KL loss. The autograd all-gather and full-vocabulary normalizer all-reduce sum gradients from these replicated losses, multiplying student gradients by the TP size even though the reported loss is correct.
Scale only the Top-K loss's backward contribution by
1 / TP. Forward loss values, loss-balancer weights, and dense KL behavior remain unchanged. Extend the existing ghost-token reference test to compare each vocabulary shard's gradients against a full-vocabulary reference, including Top-P truncation, non-default temperature, and reverse KL.Usage
No API or command-line changes. Existing
logit_kl_topk/--logit_kl_top_kconfigurations automatically receive the corrected gradients.Testing
pre-commit run --files modelopt/torch/distill/plugins/megatron.py tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py CHANGELOG.rst: passed.GPU test selection:
Run with 1, 2, or 4 visible GPUs.
Before your PR is "Ready for review"
CONTRIBUTING.md: N/A.Additional Information
No additional dependencies or communication collectives.
Summary by CodeRabbit