Reduce MimoOptimizer update-success across the world for cross-grid consensus - #5331
Merged
yashaswikarnati merged 1 commit intoJul 23, 2026
Conversation
yashaswikarnati
commented
Jun 14, 2026
| num_zeros = self.count_zeros() if self.config.log_num_zeros_in_grad else None | ||
| success = self.step_with_ready_grads() | ||
|
|
||
| # Cross-grid consensus: encoder and LLM may live on disjoint grids, so the |
Contributor
Author
There was a problem hiding this comment.
comment here is very verbose
Contributor
Author
There was a problem hiding this comment.
Trimmed to one line: # Reduce update success across the world (MIN) so disjoint-grid ranks agree. (d254c96)
yashaswikarnati
commented
Jun 14, 2026
| @pytest.mark.skipif( | ||
| torch.cuda.device_count() < 2, reason="Cross-grid step consensus requires at least 2 ranks." | ||
| ) | ||
| def test_step_success_is_world_min_consensus(): |
Contributor
Author
There was a problem hiding this comment.
overall tests surface is overkill for what we added.
Contributor
Author
There was a problem hiding this comment.
Slimmed to the single MIN-consensus case (dropped the trivial all-ranks-agree test). (d254c96)
yashaswikarnati
force-pushed
the
ykarnati/upstream-mimo-optimizer-step-consensus
branch
from
June 14, 2026 00:38
0cf56a0 to
d254c96
Compare
This was referenced Jun 16, 2026
yashaswikarnati
marked this pull request as ready for review
June 17, 2026 00:26
svcnvidia-nemo-ci
added
Final Review
deepakn94
approved these changes
Jun 17, 2026
Phlip79
approved these changes
Jun 25, 2026
Contributor
Author
|
/ok to test d254c96 |
svcnvidia-nemo-ci
added
Approved
yashaswikarnati
enabled auto-merge
June 25, 2026 00:49
MimoOptimizer drives encoder and LLM optimizers that may live on disjoint process-group grids. found_inf is already reduced across the world with MAX before the early return, but the final update-success flag from step_with_ready_grads() was only an AND across this rank's per-module optimizers and not world-consistent. Encoder-grid and LLM-grid ranks could therefore disagree on whether the step succeeded, desynchronizing LR scheduling. Reduce success over the world with ReduceOp.MIN, mirroring the existing found_inf reduction, so every rank agrees. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>
yashaswikarnati
force-pushed
the
ykarnati/upstream-mimo-optimizer-step-consensus
branch
from
July 21, 2026 23:16
7dc3565 to
7efb912
Compare
Contributor
Author
|
/ok to test 7efb912 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29951552469 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29960717151 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29968252811 |
yashaswikarnati
deleted the
ykarnati/upstream-mimo-optimizer-step-consensus
branch
July 23, 2026 02:20
shjwudp
pushed a commit
to shjwudp/Megatron-LM
that referenced
this pull request
Jul 30, 2026
…onsensus (NVIDIA#5331) Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
terminator123
pushed a commit
to 021ai/Megatron-LM
that referenced
this pull request
Aug 3, 2026
…onsensus (NVIDIA#5331) Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
svcnvidia-nemo-ci
pushed a commit
to dimapihtar/Megatron-LM
that referenced
this pull request
Aug 4, 2026
…onsensus (NVIDIA#5331) Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
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.
MimoOptimizer.step() already reduces found_inf across the world (MAX), but the final update-success flag from step_with_ready_grads() was only ANDed across this rank's per-module optimizers. For encoder/LLM on disjoint grids that lets ranks disagree on whether the step succeeded, desynchronizing LR scheduling.
This adds a world ReduceOp.MIN on the success flag, mirroring the found_inf reduction, so every rank agrees. Uses the torch.distributed world group (no parallel_state reads).
Part of the NMFW-516 MIMO-on-stock-trainloop series.
cog: per-rank verified green on cw-dfw 8-GPU.