Skip to content

[Dev] fix(moe): align MoE router loss scale semantics #3768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Victarry wants to merge 2 commits into from
+104 −20
Closed

[Dev] fix(moe): align MoE router loss scale semantics #3768

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
25642dd
fix(moe): align MoE router loss scale semantics
Mar 10, 2026
868d268
docs(moe): explain gradient scale invariants in router aux loss attac…
May 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 37 additions & 20 deletions megatron/core/transformer/moe/router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
67 changes: 67 additions & 0 deletions tests/unit_tests/transformer/moe/test_aux_loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down