Important bugfixes in local CG implementation that were leading to loss curve gaps for latent MoE models - #4433
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
Signed-off-by: root <jiemingz@nvidia.com>
12a4724 to
edf80e8
Compare
|
/ok to test 5a1a920 |
deepakn94
changed the title
| o.cg_buffer_metadata.fwd_cudagraph_buffer = fwd_graph_out | ||
| fwd_buffer_reuse_ref_count += 1 | ||
|
|
||
| # if an input buffer requires a copy, and does not have metadata attached to it at this |
There was a problem hiding this comment.
Why is this being deleted?
There was a problem hiding this comment.
This was causing a convergence issue, due to buffers being reused too early and affecting the correctness of the backward pass.
This optimization was not actually valid since the tensors this optimization was targeting were not actually able to be freed, so I'm just deleting it here. Because it only targets tensors that require an explicit copy into the graph buffer, its also not used for inference.
|
/claude review |
|
Can you list the bugfixes in the PR description? |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Thanks for these bugfixes! The changes look good overall — backing up/restoring buffers across warmup, using torch.clone instead of torch.zeros_like for warmup inputs, and removing the premature pool return and incorrect frozen-layer eval mode all make sense.
One issue: the buffer restore (lines 1006-1008) is inside the if self.training and torch.is_grad_enabled(): block, but the backup (lines 800-802) is unconditional. See inline comment for details.
🤖 Generated with Claude Code
deepakn94
changed the title
Signed-off-by: root <jiemingz@nvidia.com>
|
/claude review |
svcnvidia-nemo-ci
added
Approved
|
/ok to test db9e078 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24909597688 |
…ss curve gaps for latent MoE models (NVIDIA#4433) Signed-off-by: root <jiemingz@nvidia.com>
…ss curve gaps for latent MoE models (NVIDIA#4433) Signed-off-by: root <jiemingz@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
…ss curve gaps for latent MoE models (NVIDIA#4433) Signed-off-by: root <jiemingz@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
…ss curve gaps for latent MoE models (NVIDIA#4433) Signed-off-by: root <jiemingz@nvidia.com>
5 participants
What does this PR do ?
The PR specifically is addressing the following issues that were seen to affect convergence:
cuda_graphs.pythat is removed..training()state of the router. During graph capture we set the module to.eval()in the case of its recomputed in the backward pass, causing the expert bias to be corrupted.Contribution process
Pre-checks
Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.