diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 0b139e70443..f52a0a79162 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -443,9 +443,22 @@ def _apply_global_aux_loss( fused=self.config.moe_router_fusion, ) + # Compensate for DP dilution so global_aux_loss_coeff produces the same router + # gradient as aux_loss_coeff under identical input. + # global_aux_loss reduces tokens_per_expert/total_num_tokens over tp_dp_cp_group, but + # `probs` in switch_load_balancing_loss_func stays local to each rank (autograd cannot + # cross DP). The result is that for the same input, global_aux_loss is dp_size smaller + # than aux_loss, and so is its backward gradient. Subsequent DDP gradient averaging + # over dp_cp_group does not recover this factor. Multiply by dp_size here so the + # gradient strength matches the aux_loss baseline; the same dp_size is propagated to + # `normalize_scale` below so the logged value (aux_loss / normalize_scale) is + # unchanged and dashboards stay comparable across this fix. + dp_size = self.tp_dp_cp_group.size() // self.tp_cp_group.size() + global_aux_loss = global_aux_loss * dp_size + probs = self.attach_and_log_load_balancing_loss( probs, - global_aux_loss_coeff, + global_aux_loss_coeff * dp_size, global_aux_loss, "global_load_balancing_loss", self.tp_dp_cp_group, @@ -457,7 +470,7 @@ def _apply_global_aux_loss( def attach_and_log_load_balancing_loss( self, activation: torch.Tensor, - aux_loss_coeff: float, + normalize_scale: float, aux_loss: torch.Tensor, aux_loss_name: str, reduce_group: torch.distributed.ProcessGroup, @@ -468,7 +481,12 @@ def attach_and_log_load_balancing_loss( Args: activation (torch.Tensor): Activation tensor to attach the aux loss to. - aux_loss_coeff (float): Coefficient for the aux loss. + normalize_scale (float): Divisor applied to aux_loss before logging so that the + recorded metric stays on the coefficient-free scale. Callers must pass + `aux_loss_coeff` multiplied by any external scaling already applied to + `aux_loss` (e.g. the dp_size compensation in `_apply_global_aux_loss`), so + that `aux_loss / normalize_scale` reproduces the unit-coefficient loss + regardless of how the gradient strength was rescaled. aux_loss (torch.Tensor): Computed aux loss. aux_loss_name (str): Name of the aux loss for logging. reduce_group (torch.distributed.ProcessGroup): Process group for reduction. @@ -502,22 +520,24 @@ def attach_and_log_load_balancing_loss( get_moe_metrics_tracker().record( aux_loss_name, - aux_loss / aux_loss_coeff, + aux_loss / normalize_scale, layer_number, num_layers, reduce_group=reduce_group, 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. - # Use valid_token_count (excluding padding) if provided, otherwise use total tokens. + # Align with the non-per-token baseline gradient scale. + # In per-token mode, `MoEAuxLossAutoScaler` is seeded with `loss_scale` only + # (schedules.py drops the `cp / num_microbatches` factor), and finalize_model_grads + # divides every gradient by the global token count, all-reduced over dp_cp_group. + # That global divisor does not include the TP/CP factor that splits the local + # sequence (CP slices tokens directly; SP slices via `activation.shape[0]`), so + # attaching just `local_num_tokens` leaves the aux gradient 1/tp_cp_size smaller + # than the non-per-token baseline. Multiply by tp_cp_group.size() here to cancel + # that factor. The same correction applies to z_loss in `apply_z_loss`. num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0] + num_tokens = num_tokens * self.tp_cp_group.size() activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens) else: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss) @@ -541,15 +561,12 @@ 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. - # Count valid tokens: sum of inverted mask (False -> True = valid) + # Mirror the per-token attach scaling in `attach_and_log_load_balancing_loss`: + # finalize_model_grads divides by global token count reduced over dp_cp_group, + # which omits the TP/CP factor that splits the local sequence. Multiply by + # tp_cp_group.size() so the z_loss gradient stays on the non-per-token baseline. num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0] + num_tokens = num_tokens * self.tp_cp_group.size() logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens) 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 3b4697fc71a..383cc4d1d6c 100644 --- a/tests/unit_tests/transformer/moe/test_aux_loss.py +++ b/tests/unit_tests/transformer/moe/test_aux_loss.py @@ -467,6 +467,73 @@ def test_global_aux_loss(self, tp_size, ep_size, cp_size): assert router.ga_steps == 0 assert torch.all(router.global_tokens_per_expert == 0) + @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_global_aux_loss_gradient_scale(self, tp_size, ep_size, cp_size): + """Test that global_aux_loss produces the same router gradient as aux_loss + when all DP ranks receive identical input. + + See: https://github.com/NVIDIA/Megatron-LM/issues/3672 + """ + 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) + + aux_loss_router = self.new_router( + moe_router_load_balancing_type="aux_loss", + moe_aux_loss_coeff=1.0, + moe_router_dtype="fp64", + tensor_model_parallel_size=tp_size, + expert_tensor_parallel_size=ep_size, + context_parallel_size=cp_size, + ).cuda() + global_aux_loss_router = self.new_router( + moe_router_load_balancing_type="global_aux_loss", + moe_aux_loss_coeff=1.0, + moe_router_dtype="fp64", + tensor_model_parallel_size=tp_size, + expert_tensor_parallel_size=ep_size, + context_parallel_size=cp_size, + ).cuda() + + # Set identical weights + with torch.no_grad(): + global_aux_loss_router.weight.copy_(aux_loss_router.weight) + + # Create identical input across all DP ranks using a fixed seed + torch.manual_seed(0) + hidden_states = torch.randn( + (32, 1, aux_loss_router.config.hidden_size), + device=torch.device("cuda"), + dtype=torch.bfloat16, + ) + + # Forward + backward for aux_loss router (zero out main grad, isolate aux loss grad) + clear_aux_losses_tracker() + aux_loss_router.weight.grad = None + scores1, _ = aux_loss_router(hidden_states) + scores1.backward(torch.zeros_like(scores1)) + grad1 = aux_loss_router.weight.grad.clone() + + # Forward + backward for global_aux_loss router + clear_aux_losses_tracker() + global_aux_loss_router.weight.grad = None + scores2, _ = global_aux_loss_router(hidden_states) + scores2.backward(torch.zeros_like(scores2)) + grad2 = global_aux_loss_router.weight.grad.clone() + + assert torch.equal(grad1, grad2), ( + f"global_aux_loss gradient should match aux_loss gradient with identical input. " + f"Max diff: {(grad1 - grad2).abs().max().item()}" + ) + clear_aux_losses_tracker() + @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize(