From 7efb9120fffb95922738099a3ea277deb8ebfec2 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Sat, 13 Jun 2026 16:35:25 -0700 Subject: [PATCH] mimo: reduce optimizer step success across the world (MIN) 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 Signed-off-by: ykarnati --- megatron/core/models/mimo/optimizer.py | 5 ++++ .../mimo/test_mimo_optimizer_consensus.py | 24 +++++++++++++++++++ 2 files changed, 29 insertions(+) create mode 100644 tests/unit_tests/models/mimo/test_mimo_optimizer_consensus.py diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 71500b5fcb6..821c3b93065 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/megatron/core/models/mimo/optimizer.py @@ -103,6 +103,11 @@ def step(self) -> Tuple[bool, Optional[float], Optional[int]]: num_zeros = self.count_zeros() if self.config.log_num_zeros_in_grad else None success = self.step_with_ready_grads() + # Reduce update success across the world (MIN) so disjoint-grid ranks agree. + success_tensor = torch.tensor([1 if success else 0], dtype=torch.int, device="cuda") + torch.distributed.all_reduce(success_tensor, op=torch.distributed.ReduceOp.MIN) + success = bool(success_tensor.item()) + return success, grad_norm, num_zeros @torch.no_grad() diff --git a/tests/unit_tests/models/mimo/test_mimo_optimizer_consensus.py b/tests/unit_tests/models/mimo/test_mimo_optimizer_consensus.py new file mode 100644 index 00000000000..1bba477ea7f --- /dev/null +++ b/tests/unit_tests/models/mimo/test_mimo_optimizer_consensus.py @@ -0,0 +1,24 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Distributed test for MimoOptimizer cross-grid step-success consensus.""" + +import pytest +import torch + +from megatron.core.models.mimo.optimizer import MimoOptimizer +from megatron.core.optimizer.optimizer_config import OptimizerConfig +from tests.unit_tests.test_utilities import Utils + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires >= 2 ranks.") +def test_step_success_is_world_min(): + """One rank's failed update must propagate to every rank via the MIN reduction.""" + Utils.initialize_distributed() + try: + opt = MimoOptimizer(module_infos={}, config=OptimizerConfig(log_num_zeros_in_grad=False)) + last_rank = torch.distributed.get_world_size() - 1 + opt.step_with_ready_grads = lambda: torch.distributed.get_rank() != last_rank + success, _, _ = opt.step() + assert success is False + finally: + Utils.destroy_model_parallel()