Fix gradient corruption with layerwise param all-gather overlap - #4609
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Fix looks correct. grad_data is reused as the all-gather receive buffer (line 355) and must be zeroed after params are copied out, otherwise main_grad (a view into grad_data) accumulates on stale data. Both sync and async paths are covered, and the pattern matches the existing mxfp8 zeroing at line 490.
Left one suggestion about adding a multi-step regression test — the existing unit tests only do a single step and wouldn't have caught this.
|
/claude review |
There was a problem hiding this comment.
Clean, well-scoped bug fix. The root cause analysis is solid — grad_data reused as the all-gather receive buffer leaves stale data that corrupts gradients via the main_grad view. Zeroing in both sync and async paths is the right fix.
One minor suggestion: the async path through finish_param_sync also got the fix but test_overlap_param_gather_async_dispatch_and_finish doesn't assert that grad_data is zeroed afterward. Adding the same assertion there would guard both code paths against regressions.
38f9830 to
b9c58ed
Compare
|
/claude review |
b9c58ed to
f14e20a
Compare
|
/claude review |
…orruption The layerwise param sync path reuses grad_data as the all-gather receive buffer. After copying gathered params to model params, grad_data was not zeroed, so it still contained the all-gather results. Since main_grad is a view into grad_data, accumulation into main_grad would start from the result of the latest parameter all-gather instead of zero, corrupting gradients when using --overlap-param-gather with the layerwise distributed optimizer (e.g., Muon). Fix: zero bucket.grad_data after copying gathered params in both the synchronous (start_param_sync) and async (finish_param_sync) paths, matching the existing pattern used for the mxfp8 path. Add grad_data zeroing assertions to test_overlap_param_gather_multi_iteration and test_overlap_param_gather_async_dispatch_and_finish. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
f14e20a to
c486e79
Compare
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25403344061 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25408640123 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25415024224 |
…ync-2026-05 Upstream tip: 932d9ee (Add periodic GPU sniff tests, 2026-05-07). Conflict resolution: - 13 modify/delete conflicts in .github/* and skills/*: kept deletion (matches policy from c037427 removing upstream-NVIDIA CI machinery not relevant to this fork). - AGENTS.md content conflict: kept fork version (skills/contributing sections from upstream don't apply since the directories were removed). - New upstream additions matching the same removal policy were dropped: .claude/settings.json, .github/workflows/nightly-sync-main-to-dev.yml, skills/{cicd,linting-and-formatting,nightly-sync,run-on-slurm,testing, update-golden-values}/SKILL.md. - SECURITY.md (new from upstream): kept (generic security policy). Verified post-merge: - pretrain_gpt.py logging-patch hook intact (lines 405-408). - megatron/core/optimizer/__init__.py emerging_optimizers + muon logic intact. - megatron/training/arguments.py custom muon flags intact (lines 2297-2309). Notable upstream changes pulled in: - Removed legacy transformer + legacy GPT modules (NVIDIA#4207, NVIDIA#4322). - Docker base image bump to 26.04-py3 (NVIDIA#4611). - Inference fixes (vLLM grouped GEMM NVIDIA#4566, FlashInfer sampling NVIDIA#2456, EP sync NVIDIA#4607, MoE dispatcher fixes NVIDIA#4576). - Gradient corruption fix with layerwise param all-gather overlap (NVIDIA#4609). - Hybrid model + Flextron + GPU sniff tests + named validation sets + fault injection + InJob restart.
…IA#4609) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
…IA#4609) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
…IA#4609) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
…IA#4609) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
6 participants
Summary
--overlap-param-gatherwith the layerwise distributed optimizer (e.g., Muon).start_param_sync) reusesbucket.grad_dataas the all-gather receive buffer. Afterfinish_param_synccopies gathered params to model params,grad_datawas not zeroed — leaving stale all-gather data in the buffer.param.main_gradis a view intograd_data, the backward pass accumulates gradients (main_grad.add_(param.grad.data)) on top of this stale data, corrupting gradients.bucket.grad_dataafter copying gathered params in both the synchronous (start_param_sync) and async (finish_param_sync) paths, matching the existing pattern used for the mxfp8 path.Test plan
--overlap-param-gather --overlap-grad-reduceand verify loss curves match non-overlapped baseline🤖 Generated with Claude Code