diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 6c61bb3ed61..e4591ce3acf 100644 --- a/megatron/core/transformer/moe/router.py +++ b/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 diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py index de4fa906821..c8c7bf0dd0f 100644 --- a/tests/unit_tests/transformer/moe/test_aux_loss.py +++ b/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,