[training] fix: respect module-local MIMO loss normalization (#5555)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/megatron_mimo_parallel_utils.py

@@ -24,6 +24,7 @@
 from megatron.core.distributed.finalize_model_grads import finalize_model_grads as _finalize_model_grads
 from megatron.core.models.mimo import MimoModel
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
+from megatron.core.utils import get_model_config
 
 from megatron.bridge.models.megatron_mimo.megatron_mimo_provider import MegatronMIMOInfra
 
@@ -258,9 +259,9 @@ def _find_module(grid):
         # pass num_tokens=None for them to skip the broken normalization,
         # then broadcast the correct total from LLM and apply it manually.
         #
-        # With gradient_scaling_factor=1.0 (calculate_per_token_loss=True),
-        # DDP does a plain SUM.  After dividing by the global token count
-        # the gradient is correct — no DP compensation factor needed.
+        # A per-token module's DDP does a plain SUM, while a non-per-token
+        # module's DDP averages over its DP ranks. Account for that module-local
+        # choice before dividing by the global token count.
 
         # Phase 1: gradient all-reduce for each module.  Only the LLM gets
         # num_tokens so _finalize_model_grads can PP-broadcast + DP-all-reduce
@@ -295,7 +296,9 @@ def _find_module(grid):
             if module is not None and is_current_rank_in_grid(grid):
                 module_name, _ = _find_module(grid)
                 if module_name != MIMO_LANGUAGE_MODULE_KEY and num_tokens > 0:
-                    module.scale_gradients(1.0 / num_tokens.float().item())
+                    module_dp = _get_dp_size_from_grid(grid)
+                    ddp_mean_correction = 1.0 if get_model_config(module).calculate_per_token_loss else module_dp
+                    module.scale_gradients(ddp_mean_correction / num_tokens.float().item())
     else:
         # calculate_per_token_loss=False path.
         #
@@ -316,8 +319,10 @@ def _find_module(grid):
                     )
 
                     module_dp = _get_dp_size_from_grid(grid)
-                    if module_dp != llm_dp:
-                        module.scale_gradients(float(module_dp) / float(llm_dp))
+                    ddp_mean_correction = 1.0 if get_model_config(module).calculate_per_token_loss else module_dp
+                    module_scale = float(ddp_mean_correction) / float(llm_dp)
+                    if module_scale != 1.0:
+                        module.scale_gradients(module_scale)
 
 
 def zero_grad_buffer_for_multimodule(module_to_grid_tuple: List[Tuple]):

tests/unit_tests/training/megatron_mimo/test_megatron_mimo_parallel_utils.py

@@ -376,6 +376,7 @@ def test_per_token_loss_path(self, mock_finalize, mock_dp, mock_in_grid, mock_di
         s = _build_two_module_setup(llm_dp=2, encoder_dp=4)
         mock_in_grid.return_value = True
         mock_dp.side_effect = lambda grid: s.dp_by_grid[id(grid)]
+        s.encoder_module.config.calculate_per_token_loss = True
 
         num_tokens = torch.tensor(100)
 
@@ -405,6 +406,58 @@ def test_per_token_loss_path(self, mock_finalize, mock_dp, mock_in_grid, mock_di
         s.encoder_module.scale_gradients.assert_called_once_with(1.0 / 100)
         s.llm_module.scale_gradients.assert_not_called()
 
+    @pytest.mark.parametrize(
+        ("language_per_token", "encoder_per_token", "expected_gradient"),
+        [
+            (True, False, 4.0),
+            (False, True, 200.0),
+        ],
+    )
+    @patch(f"{MODULE}.dist")
+    @patch(f"{MODULE}.is_current_rank_in_grid")
+    @patch(f"{MODULE}._get_dp_size_from_grid")
+    @patch(f"{MODULE}._finalize_model_grads")
+    def test_mixed_per_token_configs_preserve_encoder_gradient_normalization(
+        self,
+        mock_finalize,
+        mock_dp,
+        mock_in_grid,
+        mock_dist,
+        language_per_token,
+        encoder_per_token,
+        expected_gradient,
+    ):
+        """Module-local DDP scaling must not change the global encoder gradient."""
+        from megatron.bridge.training.megatron_mimo_parallel_utils import finalize_model_grads_multimodule
+
+        s = _build_two_module_setup(llm_dp=2, encoder_dp=4)
+        mock_in_grid.return_value = True
+        mock_dp.side_effect = lambda grid: s.dp_by_grid[id(grid)]
+        s.llm_module.config.calculate_per_token_loss = language_per_token
+        s.encoder_module.config.calculate_per_token_loss = encoder_per_token
+
+        # Simulate the value after MCore DDP reduction for a global encoder
+        # gradient sum of 400. Non-per-token DDP averages over encoder DP;
+        # per-token DDP leaves the global sum intact.
+        encoder_gradient = torch.tensor(400.0)
+
+        def finalize_with_module_local_ddp(model, **kwargs):
+            if model[0] is s.encoder_module and not encoder_per_token:
+                encoder_gradient.div_(s.dp_by_grid[id(s.encoder_grid)])
+
+        mock_finalize.side_effect = finalize_with_module_local_ddp
+        s.encoder_module.scale_gradients.side_effect = encoder_gradient.mul_
+
+        num_tokens = torch.tensor(100) if language_per_token else None
+        finalize_model_grads_multimodule(
+            [MagicMock()],
+            num_tokens,
+            infra=s.infra,
+            module_to_grid_tuple=s.module_to_grid_tuple,
+        )
+
+        assert encoder_gradient.item() == pytest.approx(expected_gradient)
+
     @patch(f"{MODULE}.dist")
     @patch(f"{MODULE}.is_current_rank_in_grid")
     @patch(f"{MODULE}._get_dp_size_from_grid")
@@ -420,6 +473,8 @@ def test_non_per_token_loss_path_applies_dp_compensation(self, mock_finalize, mo
         s = _build_two_module_setup(llm_dp=2, encoder_dp=4)
         mock_in_grid.return_value = True
         mock_dp.side_effect = lambda grid: s.dp_by_grid[id(grid)]
+        s.llm_module.config.calculate_per_token_loss = False
+        s.encoder_module.config.calculate_per_token_loss = False
 
         finalize_model_grads_multimodule(
             [MagicMock()],