Repository navigation
Conversation
|
@claude review always |
|
@codex review |
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: 5c3260caa2
ℹ️ 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".
| logits = logits - logits_max.unsqueeze(-1) | ||
|
|
||
| partition_vocab_size = logits.size(-1) | ||
| vocab_start = group.rank() * partition_vocab_size |
There was a problem hiding this comment.
When Megatron is absent and compute_log_probs(..., process_group=None) is used, this dereference raises AttributeError before computing the loss. None is an explicitly supported value in the public signature and existing callers/tests use it to mean the default process group; the fallback should derive rank 0 when distributed is uninitialized or use dist.get_rank(group) so that None retains its established meaning.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Done in df4214c: None now means the default group, as in torch.distributed. Thanks.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Nothing blocking. The comments below are optional suggestions. There is no need to push a fix for them before merging.
Beyond the inline finding, I checked one more thing and it holds up. The test sets sys.modules["megatron"] = None and patches _MEGATRON_INSTALLED, but both stay inside the workers that run_multiprocess starts with mp.spawn, so they don't leak into the pytest process or later tests.
Extended reasoning...
The test patches module-level state, and it runs inside workers started by mp.spawn (tests/fast/dist_utils.py:23), so the change is confined to that test's own workers.
5c3260c to
1fa6b28
Compare
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
|
@claude review |
1fa6b28 to
df4214c
Compare
There was a problem hiding this comment.
Nothing blocking. The comments below are optional suggestions. There is no need to push a fix for them before merging.
Still open from earlier reviews (1):
- Unresolved: 1 minor or pre-existing.
Additional findings (outside the current diff — GitHub can't attach inline comments there):
-
🟣
miles/backends/training_utils/loss/hub/math_utils.py— Pre-existing: torchtitan runs with TP>1 and--true-on-policy-modestill crash without Megatron, so the fallback does not cover everycompute_log_probspath. The true-on-policy branch returns at math_utils.py:297, before the new_MEGATRON_INSTALLEDcheck. With TP>1 it goes through_ReplicatedLossAllGatherLastDim.forward, which importsmegatron.core.tensor_parallel.mappingsat math_utils.py:452. Fix: have every Megatron import reachable fromcompute_log_probswork without Megatron. Here, calldist.all_gather_into_tensordirectly, which is whatdist_all_gather_funcwraps.Why this was flagged
The trigger is a torchtitan run from the pip wheels with no Megatron installed,
tensor_model_parallel_size> 1 (passed to torchtitan at torchtitan_utils/config.py:76) andtrue_on_policy_modeset. Herecompute_log_probstakes the branch at math_utils.py:296-299._gather_true_on_policy_full_logits(math_utils.py:404-408) sees tp_size > 1 and calls_ReplicatedLossAllGatherLastDim.apply. Its forward then runsfrom megatron.core.tensor_parallel.mappings import dist_all_gather_funcat math_utils.py:452 and raisesModuleNotFoundError: No module named 'megatron'on the first log-prob step. The same path also runs from the fused true-on-policy helper at math_utils.py:1138. The base branch fails the same way, so this is pre-existing. The new_MEGATRON_INSTALLEDguard does not help because it sits only on the non-true-on-policy path at math_utils.py:301-304. FSDP is not affected because its TP group always has size 1 (fsdp_utils/parallel.py:87-91).Verification: Pre-existing. Trigger: a torchtitan install with no Megatron,
--tensor-model-parallel-size> 1 and--true-on-policy-mode. The log-prob step still raisesModuleNotFoundError: No module named 'megatron'. The base branch fails the same way, and the PR does not change line 452.compute_log_probs's true-on-policy branch (lines 296-299) returns before the new_MEGATRON_INSTALLEDcheck at line 301.
b72912f to
0d7bf56
Compare
df4214c to
f8cd0e7
Compare
f8cd0e7 to
2bf5e05
Compare
2bf5e05 to
2b84a30
Compare
HJSang
changed the title
There was a problem hiding this comment.
I reviewed the latest push and didn't find any bugs. A human should still look, because commit 5eda13c changes the scope of this PR: the plain-torch cross entropy now runs on every install, including Megatron runs that used to get Megatron's fused kernel. The PR description still says it is used only when megatron.core is absent.
What was reviewed:
_VocabParallelCrossEntropy: the MAX/SUM all-reduces, the shard range fromdist.get_rank(group), and the masked gather and one-hot scatter in backward.- The in-place forward on fp32 input. Every caller goes through
_upcast_chunk_to_fp32(copy=True), and the debugdetach()path also gets that private copy, so no caller tensor is mutated. Nothing earlier in the graph saves that tensor for backward. - The monkeypatch target in the tests.
compute_log_probsresolvesvocab_parallel_cross_entropyas a module global at call time, so patchingmath_utilsreaches it.
Extended reasoning...
The change replaces Megatron's lazily imported fused_vocab_parallel_cross_entropy in compute_log_probs (shared by every training backend) with a new plain-torch autograd Function, and updates and adds tests. It touches no security surface. The latest commit makes the torch path unconditional. Megatron-backend log-prob and gradient numerics now differ from the old fused kernel in the last bits (the commit message says the fused kernel also rounded the gradient to bf16), and performance changes too. That is a design decision for the owners. The changed paths are under /miles/backends/ in .github/CODEOWNERS.
There was a problem hiding this comment.
I reviewed the two commits pushed since my last review and didn't find any bugs. A human should still sign off, because this change replaces Megatron's fused cross-entropy kernel with plain torch on every backend, Megatron included, and two loss snapshots still need regenerating.
What was reviewed:
- The new
None-group behavior:group.rank()now raises onNoneinstead of reducing over the world group. Every caller passes a real TP group: FSDP usesself_group, torchtitan usesnew_groupor the mesh group, and Megatron uses the mpu group. One note: the PR description still says the function acceptsgroup=None, which is no longer true. - Forward and backward of
_VocabParallelCrossEntropyagainst Megatron's unfused version: the MAX/SUM all-reduces, target masking, and the in-placesoftmax - one_hotgradient match. - Swapping in
all_gather_into_tensorin_ReplicatedLossAllGatherLastDim.
Extended reasoning...
The change rewrites the shared log-prob path in miles/backends/training_utils/loss/hub/math_utils.py, and adds a new test file plus a help-text tweak in miles/utils/arguments.py. It touches no security-sensitive surface. The PR description reports that the Megatron backend's default path loses about 9% speed and changes in its last bits, since the fused kernel is gone. Two loss snapshots must be regenerated in another repo before merge. Whether that trade-off is acceptable is a design call for a maintainer, not a code-correctness question.
|
cc cmt [P2] The two snapshots listed under "Snapshots" have not been pushed yet, and this PR's CI has not run the test that compares them. In Suggested merge gate: push the regenerated snapshots, then comment |
compute_log_probs, which FSDP and torchtitan share with Megatron, always called Megatron's fused_vocab_parallel_cross_entropy, so neither backend could train on an install without Megatron (a pip wheel). When Megatron is absent it now uses a plain-torch port of the same vocab-parallel math; with Megatron installed the fused kernel runs exactly as before, and a broken Megatron install still fails at its import. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Spawning dominates the runtime; two spawns keep the file inside the fast CPU budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
megatron-energon installs a bare megatron namespace without megatron.core,
which find_spec("megatron") counted as Megatron. The fallback also took
group.rank() on a None group, which compute_log_probs callers pass for the
default group; dist.get_rank(None) gives torch.distributed's meaning.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The shift, exp and gradient each allocated another full-size fp32 logits buffer, so a run without Megatron could OOM where the fused kernel fits. One copy keeps the caller's logits untouched; the rest works in place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Without Megatron the local name fused_vocab_parallel_cross_entropy held the plain-torch fallback; both branches now bind cross_entropy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
One implementation instead of a switch on whether megatron.core is installed, so the log-probs no longer depend on what happens to be importable: with Megatron installed, the FSDP and torchtitan backends ran Megatron's fused kernel, and on a pip install the torch one. The torch version is bit-identical to Megatron's unfused vocab-parallel cross entropy; its fused kernel differs in the last bits (torch.compile's exp and summation order, up to 1.9e-6 in the loss on H200) and rounds the gradient to bf16. On H200 at 16k x 151936 the torch version is about 9% faster than the fused kernel with the same peak memory. The forward works in place on its fp32 input, as Megatron's does: every caller passes a fresh copy (_upcast_chunk_to_fp32), so a second copy would add an fp32 [T, V] buffer to every Megatron run. The fake-kernel tests patch the miles function, and the Megatron comparison is its own test, skipped without Megatron. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…og-prob path - The torch kernel took dist.get_rank(None), i.e. the world group, where callers elsewhere mean no tensor parallelism by None: with 8 ranks and no TP every rank but 0 masked its targets and the SUM mixed data-parallel ranks. It now takes group.rank(), as Megatron's kernel did, so None or a non-member group raises. Every caller passes a group. - The tensor-parallel true-on-policy gather imported Megatron's dist_all_gather_func, which is torch.distributed.all_gather_into_tensor; call that directly, so a wheel install can use it. - --debug-unified-grad-fused-logprob's help no longer mentions a fused kernel. Found by an adversarial review of this PR. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
319eb84 to
e6ded55
Compare
a53d9a6 to
9022cb2
Compare
guapisolo
left a comment
guapisolo
left a comment
There was a problem hiding this comment.
cc cmt
Approving #4315 at 9022cb2, with one merge gate. The review ran at a53d9a6; the rebase only moves the --debug-unified-grad-fused-logprob help-text edit into miles/utils/args/configs/debug.py, and the branch now sits on #4308's head.
- [P2] Merge gate, still open: the two snapshots listed in the body are still the 2026-08-31 versions in
yueming-yuan/miles-artifacts, and this PR's CI has not runtest_loss_snapshot.py(precisionlabel). Merging first would make every laterprecisionrun fail those two cases. Push the regenerated snapshots and get/rerun-test tests/fast/backends/training_utils/loss/test_loss_snapshot.pygreen before merging (#4315 (comment)).
Verification: read-only review of this increment at a53d9a6 with its green PR Test run (40553182119), and a range-diff to this head. Not covered: the snapshot comparison and a Megatron GPU end-to-end run.
Approved on the condition that the snapshot gate is met before merge.
|
With the plain-torch autograd Function, both grad-enabled and no-grad calls now execute the same forward implementation. Could we remove debug_unified_grad_fused_logprob and its plumbing, and rename the fused_cross_entropy fixtures/tests to match the new implementation? |
Stack (merge in order): #4308 → #4315 → #4316 → #4311 (needs radixark/miles-wheels#4).
Problem
compute_log_probs, which every backend shares, always imported Megatron'sfused_vocab_parallel_cross_entropy(the code carried a TODO for this). Without Megatron, FSDP and torchtitan die at their first log-prob step withNo module named 'megatron', as an H200 run from the wheels did.Change
vocab_parallel_cross_entropy: Megatron's vocab-parallel cross entropy in plain torch (fp32, MAX all-reduce, one batched SUM all-reduce,softmax - one_hotgradient on the owning rank). It returns the gradient in the input's dtype and takesgroup=Noneas the default group.compute_log_probsno longer imports Megatron at all. Choosing by what is installed would have made the same FSDP or torchtitan run give different log-probs in theradixark/miles:devimage (Megatron installed) and from the wheels._upcast_chunk_to_fp32), and the backward writes into the saved softmax.How it compares with Megatron (H200, bf16 logits, 16k tokens × 151,936):
vocab_parallel_cross_entropy: bit-identical loss and gradient, GPU TP=1 and TP=2 and CPU.torch.compile, whoseexpand summation order differ in the last bits: loss within 1.9e-6, and about 6 in a million bf16 gradient entries differ by one ULP. It also rounds its gradient to bf16 even for fp32 input.Snapshots
Two loss snapshots record the fused kernel's rounding with fp32 logits and need regenerating in
yueming-yuan/miles-artifactsbefore this merges:grpo_chunked_temp_b2andsft_chunked_b2(logits_graddiffers by up to 2.9e-4, the fused kernel's bf16 rounding; the other ten snapshots pass unchanged):Testing
At head
5eda13cf8b, on an H200 devbox:tests/fast/backends/training_utils/loss_hub/: 12 passed. With Megatron hidden (sys.modules["megatron"] = None): 10 passed, the two Megatron comparisons skipped.test_vocab_parallel_cross_entropy.py(gloo, TP=1 and TP=2): matches a full-vocabulary log-softmax on every shard in fp32 and bf16; bit-identical to Megatron's unfused cross entropy;compute_log_probswith a group and withNone.tests/fast/backends,tests/fast/utils/test_arguments.py,tests/fast/utils/debug_utils: 3,762 passed. Failing: the two snapshots above, and five that also fail onmain(Qwen3.5 packing intest_fsdp_hf_compat.pyandtest_fsdp_hub_kernels.py, and a Ray fixture).🤖 Generated with Claude Code