Fix for sequence-level aux MoE loss being dependent on batch size

megatron/core/transformer/moe/router.py

@@ -495,7 +495,9 @@ def _apply_seq_aux_loss(
             aux_loss,
             "seq_load_balancing_loss",
             self.tp_cp_group,
-            valid_token_count=local_num_tokens,
+            # local_num_tokens is per-sequence (bsz folded into the expert dim above);
+            # * bsz recovers the micro-batch total, else per-token-loss scaling keeps a 1/MBS.
+            valid_token_count=local_num_tokens * bsz,
         )
         return probs
 

tests/unit_tests/transformer/moe/test_aux_loss.py

@@ -383,6 +383,79 @@ def test_seq_aux_loss(self, tp_size, ep_size, cp_size):
         torch.testing.assert_close(aux_loss, seq_aux_loss)
         torch.testing.assert_close(grad1, grad2)
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("with_padding", [False, True])
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_seq_aux_loss_mbs_invariant_per_token_loss(
+        self, tp_size, ep_size, cp_size, with_padding
+    ):
+        """seq_aux_loss gradient must be invariant to MBS under --calculate-per-token-loss.
+
+        The same global batch is processed as N micro-batches of size 1 (MBS=1) and as one
+        micro-batch of size N (MBS=N). Both cover the same tokens, so the finalize-time
+        1/total_tokens normalization is an identical constant and the accumulated
+        router-weight aux gradients must match. Before the fix (valid_token_count dropped the
+        bsz factor), the MBS=N gradient is scaled by 1/N and the assertion fails. The padding
+        case additionally checks the correction uses valid (non-padded) token counts.
+        """
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            expert_tensor_parallel_size=ep_size,
+            context_parallel_size=cp_size,
+        )
+        model_parallel_cuda_manual_seed(42)
+        clear_aux_losses_tracker()
+
+        router = self.new_router(
+            moe_router_load_balancing_type="seq_aux_loss",
+            moe_aux_loss_coeff=1.0,
+            moe_router_dtype="fp64",
+            calculate_per_token_loss=True,
+            # fp32 weights so the MBS=1 gradient (accumulated over N backward passes)
+            # is not degraded by bf16 rounding relative to the single MBS=N backward.
+            params_dtype=torch.float32,
+            bf16=False,
+            tensor_model_parallel_size=tp_size,
+            expert_tensor_parallel_size=ep_size,
+            context_parallel_size=cp_size,
+        ).cuda()
+
+        seq_len = 32
+        num_seqs = 4
+        with get_cuda_rng_tracker().fork():
+            hidden_states = torch.randn(
+                (seq_len, num_seqs, router.config.hidden_size),
+                device=torch.device("cuda"),
+                dtype=torch.float32,
+            )
+        padding_mask = None
+        if with_padding:
+            # True marks padding tokens (second half of each sequence).
+            padding_mask = torch.zeros((seq_len, num_seqs), dtype=torch.bool, device="cuda")
+            padding_mask[seq_len // 2 :, :] = True
+
+        def run(indices):
+            pmask = None if padding_mask is None else padding_mask[:, indices]
+            scores, _ = router(hidden_states[:, indices, :].contiguous(), padding_mask=pmask)
+            scores.backward(torch.zeros_like(scores))  # isolate the aux-loss gradient
+            clear_aux_losses_tracker()
+
+        # MBS=1: N micro-batches of size 1, accumulating the aux-loss gradient.
+        router.weight.grad = None
+        for b in range(num_seqs):
+            run(slice(b, b + 1))
+        grad_mbs1 = router.weight.grad.clone()
+
+        # MBS=N: a single micro-batch of size N.
+        router.weight.grad = None
+        run(slice(0, num_seqs))
+        grad_mbsN = router.weight.grad.clone()
+
+        torch.testing.assert_close(grad_mbs1, grad_mbsN)
+
     @pytest.mark.internal
     @pytest.mark.skipif(
         not torch.cuda.is_available() or not HAVE_ROUTER_FUSION,