Reduce MimoOptimizer update-success across the world for cross-grid consensus

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()

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()