diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 235a616dfbf..09de99bbebf 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -479,16 +479,36 @@ def attach_and_log_load_balancing_loss( needs_dp_avg=needs_dp_avg, ) if self.calculate_per_token_loss: - # Scale the aux_loss by the number of tokens. - # The expected final scaling for aux_loss gradients is 1/(num_micro_batches * dp_size). - # After commit 02648000, Megatron started using the number of total tokens to scale - # gradients under the argument of calculate_per_token_loss, - # which scales both the main_loss gradient and aux_loss gradient by - # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads function. - # To correct this scaling, we need to scale the aux_loss by num_local_tokens here. + # Target final scaling on aux_loss gradients: 1 / (num_micro_batches * dp_size), + # matching the !calculate_per_token_loss path. + # + # --calculate-per-token-loss already divides every parameter gradient by + # total_global_tokens (the global non-padded token count summed in + # finalize_model_grads). The router's `num_local_tokens` (= activation.shape[0]) + # is sequence-parallel sharded — the router weight is marked + # `sequence_parallel=True` in Router.reset_parameters (see + # `setattr(self.weight, 'sequence_parallel', ...)` above), so each TP rank + # computes a partial gradient on the router weight from its local sequence + # shard, and `_allreduce_non_tensor_model_parallel_grads` SUMS those partial + # gradients across the TP group. Re-expressing total_global_tokens in terms of the + # router's `num_local_tokens`: + # total_global_tokens + # = num_micro_batches * dp_cp_size * loss_func_local_tokens + # = num_micro_batches * dp_cp_size * tp_size * num_local_tokens + # = num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()) + # (using loss_func_local_tokens = tp_size * num_local_tokens, then regrouping + # dp_cp_size * tp_size as dp_size * tp_cp_group.size()). + # + # So pre-multiplying aux_loss by num_local_tokens * tp_cp_group.size() cancels + # that same factor in total_global_tokens above, leaving 1 / (num_micro_batches * + # dp_size) as the effective scaling on the aux_loss gradient — the target. # Use valid_token_count (excluding padding) if provided, otherwise use total tokens. - num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0] - activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens) + num_local_tokens = ( + valid_token_count if valid_token_count is not None else activation.shape[0] + ) + activation = MoEAuxLossAutoScaler.apply( + activation, aux_loss * num_local_tokens * self.tp_cp_group.size() + ) else: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss) return activation @@ -511,16 +531,24 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size() z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask) if self.calculate_per_token_loss: - # The expected final scaling for z_loss gradients is - # 1/(num_micro_batches * dp_size). - # After commit 02648000, Megatron started using the number of total tokens - # to scale gradients under the argument of calculate_per_token_loss, - # which scales both the main_loss gradient and z_loss gradient by - # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads(). - # To correct this scaling, we need to scale the z_loss by num_local_tokens here. + # Same derivation as in attach_and_log_load_balancing_loss: + # - Target final scaling on z_loss gradients: 1 / (num_micro_batches * dp_size). + # - In terms of the router's `num_local_tokens`, the total_global_tokens + # divisor that finalize_model_grads applies factors as + # num_micro_batches * dp_size * (num_local_tokens * tp_cp_group.size()). + # - Pre-multiplying z_loss by num_local_tokens * tp_cp_group.size() cancels + # that same factor in total_global_tokens, leaving + # 1 / (num_micro_batches * dp_size) as the effective scaling — the target. + # The /tp_cp_group.size() on moe_z_loss_coeff above is a separate forward-side + # correction: z_loss is computed independently on each TP+CP rank's local + # logits and must be averaged across TP+CP rather than summed. # Count valid tokens: sum of inverted mask (False -> True = valid) - num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0] - logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens) + num_local_tokens = ( + (~padding_mask).sum() if padding_mask is not None else logits.shape[0] + ) + logits = MoEAuxLossAutoScaler.apply( + logits, z_loss * num_local_tokens * self.tp_cp_group.size() + ) else: logits = MoEAuxLossAutoScaler.apply(logits, z_loss) diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py index ccd11bf29af..118203ee1a6 100644 --- a/tests/unit_tests/transformer/moe/test_aux_loss.py +++ b/tests/unit_tests/transformer/moe/test_aux_loss.py @@ -178,6 +178,62 @@ def test_a2a_dispatcher(self, tp_size, ep_size, cp_size): container.aux_loss_test(self.input, self.baseline_grad, "seq_load_balancing_loss") +class TestPerTokenAuxLoss: + """Regression test for the aux_loss TP/CP scaling fix under + --calculate-per-token-loss. Computes a baseline aux-loss input + gradient at (tp=1, cp=1) and asserts that each parametrized + (tp, ep, cp) config produces a matching gradient on each rank's + local input slice. Without the fix, the per-rank scale on aux_loss + would shrink with tp_cp_size and the assertion would fail at any + config with tp_size > 1 or cp_size > 1. + """ + + def setup_method(self, method): + baseline_container = AuxlossTestContainer( + tp_size=1, + ep_size=1, + pp_size=1, + cp_size=1, + num_moe_experts=8, + moe_router_topk=2, + moe_router_load_balancing_type="aux_loss", + moe_token_dispatcher_type="alltoall", + moe_aux_loss_coeff=0.1, + calculate_per_token_loss=True, + ) + moe_layer = baseline_container.moe_layer + self.input = torch.randn((32, 8, moe_layer.config.hidden_size)).cuda() + self.input.requires_grad = True + probs, indices = apply_module(moe_layer.router)(self.input) + probs.sum().mul_(0).backward() + self.baseline_grad = self.input.grad + self.input.grad = None + clear_aux_losses_tracker() + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @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_per_token_aux_loss_invariant_to_tp_cp(self, tp_size, ep_size, cp_size): + container = AuxlossTestContainer( + tp_size=tp_size, + ep_size=ep_size, + pp_size=1, + cp_size=cp_size, + num_moe_experts=8, + moe_router_topk=2, + moe_router_load_balancing_type="aux_loss", + moe_token_dispatcher_type="alltoall", + moe_aux_loss_coeff=0.1, + calculate_per_token_loss=True, + ) + container.aux_loss_test(self.input, self.baseline_grad, "load_balancing_loss") + + class TestRouterAuxLoss: def setup_method(self, method): Utils.initialize_model_parallel(1, 1) diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index 46cf639e059..9dae154fc33 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -93,6 +93,7 @@ def __init__( add_bias_linear=kwargs.get("add_bias_linear", False), moe_permute_fusion=kwargs.get("moe_permute_fusion", False), moe_flex_dispatcher_backend=kwargs.get("moe_flex_dispatcher_backend", None), + calculate_per_token_loss=kwargs.get("calculate_per_token_loss", False), ) # init moe layer