diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py index 51660d12c7d..12253e6c4b6 100644 --- a/megatron/core/distributed/finalize_model_grads.py +++ b/megatron/core/distributed/finalize_model_grads.py @@ -326,6 +326,9 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n or "global_aux_loss" in config.moe_router_load_balancing_type ) and hasattr(module, 'reset_global_aux_loss_tracker'): module.reset_global_aux_loss_tracker() + if getattr(module, 'qb_beta_accum', None) is not None: + module.qb_beta_accum.zero_() + module.qb_beta_count.zero_() def _update_router_expert_bias( @@ -368,6 +371,48 @@ def _update_router_expert_bias( expert_bias.copy_(updated_expert_bias) +def _update_router_qb_beta( + model: List[torch.nn.Module], + config: TransformerConfig, + dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, +): + """Update the quantile-balancing per-expert bias once per global batch. + + Averages each router's accumulated quantile (qb_beta_accum/qb_beta_count) across + DP, EMA-blends it with the current qb_beta, re-centers, and writes it back. + """ + qb_beta_list = [] + qb_beta_accum_list = [] + qb_beta_count_list = [] + for model_chunk in model: + for module in get_attr_wrapped_model(model_chunk, 'modules')(): + if getattr(module, 'qb_beta_accum', None) is not None and module.training: + qb_beta_list.append(module.qb_beta) + qb_beta_accum_list.append(module.qb_beta_accum) + qb_beta_count_list.append(module.qb_beta_count) + + if len(qb_beta_list) == 0: + return + + stacked_beta = torch.stack(qb_beta_list, dim=0) + local_avg_list = [ + accum / count.clamp(min=1).to(accum.dtype) + for accum, count in zip(qb_beta_accum_list, qb_beta_count_list) + ] + stacked_local_avg = torch.stack(local_avg_list, dim=0) + + torch.distributed.all_reduce( + stacked_local_avg, op=torch.distributed.ReduceOp.AVG, group=dp_cp_group + ) + + ema = config.moe_router_quantile_balancing_ema + stacked_new_beta = ema * stacked_beta + (1.0 - ema) * stacked_local_avg + stacked_new_beta = stacked_new_beta - stacked_new_beta.mean(dim=-1, keepdim=True) + + for qb_beta, new_beta in zip(qb_beta_list, stacked_new_beta): + qb_beta.copy_(new_beta) + + def _allreduce_non_tensor_model_parallel_grads( model: List[torch.nn.Module], config: TransformerConfig, @@ -542,6 +587,9 @@ def finalize_model_grads( ) _update_router_expert_bias(model, config, tp_dp_cp_group=tp_dp_cp_group) + if config.moe_router_load_balancing_type == "quantile_balancing": + _update_router_qb_beta(model, config, dp_cp_group=dp_cp_group) + reset_model_temporary_tensors(config, model) # normalize gradients for per-token loss normalization. diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 79b45156ed8..5c053adb6b1 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -200,6 +200,44 @@ def sinkhorn(cost: torch.Tensor, tol: float = 0.0001) -> torch.Tensor: return d1 * cost * d0.unsqueeze(1) +def qb_dual_update( + scores: torch.Tensor, k: int, beta: torch.Tensor, update_beta: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + """Dual coordinate-descent quantile-balancing routing assignment. + + Picks the top-k experts per token from ``scores - beta``. When ``update_beta`` is + True, also returns the raw column quantile of ``scores`` that drives each expert + toward ``m * k / n`` tokens. + + Args: + scores (torch.Tensor): Scores of shape ``[m, n]`` (tokens, experts). + k (int): Experts to select per token. + beta (torch.Tensor): Current per-expert bias of shape ``[n]``. + update_beta (bool): If False, return ``beta`` unchanged (eval/inference). + + Returns: + Tuple[torch.Tensor, torch.Tensor]: indices of shape ``[m, k]`` and either + ``beta`` (when ``update_beta`` is False) or the column quantile ``[n]``. + """ + num_tokens, num_experts = scores.shape + + topk_result = (scores - beta).topk(k + 1, dim=1) + indices = topk_result.indices[:, :-1] + + if not update_beta: + return indices, beta + + assert (num_tokens * k) % num_experts == 0, ( + "Quantile balancing requires the number of routed assignments " + f"({num_tokens} tokens * top-{k}) to be divisible by " + f"{num_experts} experts." + ) + col_target = num_tokens * k // num_experts + alpha = topk_result.values[:, -1:] + beta_local = (scores - alpha).topk(col_target + 1, dim=0).values[-1].contiguous() + return indices, beta_local + + def get_capacity( num_tokens: int, num_experts: int, capacity_factor: float, min_capacity: Optional[int] = None ) -> int: @@ -681,6 +719,7 @@ def topk_routing_with_score_function( fused: bool = False, router_replay: Optional['RouterReplay'] = None, dense_output: bool = False, + precomputed_indices: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Compute the routing probabilities and map for top-k selection with score function. @@ -705,6 +744,10 @@ def topk_routing_with_score_function( Defaults to None. dense_output (bool, optional): If True, return dense tensors [num_tokens, topk] instead of sparse tensors [num_tokens, num_experts]. Defaults to False. + precomputed_indices (torch.Tensor, optional): Top-k indices [num_tokens, topk] + selected by the caller. When given, the score function's + own top-k is bypassed and probs are computed at these + indices (e.g. for quantile balancing). Defaults to None. Returns: Tuple[torch.Tensor, torch.Tensor]: @@ -723,6 +766,9 @@ def topk_routing_with_score_function( """ assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}." num_tokens, num_experts = logits.shape + assert not ( + fused and precomputed_indices is not None + ), "precomputed_indices is not supported with the fused top-k score function." if fused: if not HAVE_TE or fused_topk_with_score_function is None: raise ValueError( @@ -793,16 +839,27 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None): if score_function == "softmax": if use_pre_softmax: scores = torch.softmax(logits, dim=-1, dtype=torch.float32) - probs, top_indices = compute_topk(scores, topk, num_groups, group_topk) + if precomputed_indices is not None: + top_indices = precomputed_indices + probs = torch.gather(scores, dim=1, index=top_indices) + else: + probs, top_indices = compute_topk(scores, topk, num_groups, group_topk) else: - scores, top_indices = compute_topk(logits, topk, num_groups, group_topk) + if precomputed_indices is not None: + top_indices = precomputed_indices + scores = torch.gather(logits, dim=1, index=top_indices) + else: + scores, top_indices = compute_topk(logits, topk, num_groups, group_topk) probs = torch.softmax(scores, dim=-1, dtype=torch.float32) elif score_function in ("sigmoid", "sqrtsoftplus"): if score_function == "sigmoid": scores = torch.sigmoid(logits.float()) else: scores = torch.nn.functional.softplus(logits.float()).sqrt() - if expert_bias is not None: + if precomputed_indices is not None: + top_indices = precomputed_indices + scores = torch.gather(scores, dim=1, index=top_indices) + elif expert_bias is not None: scores_for_routing = scores + expert_bias.float() _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) scores = torch.gather(scores, dim=1, index=top_indices) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 7414c8a7ab0..6273a520588 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -17,6 +17,7 @@ apply_router_token_dropping, compute_routing_scores_for_aux_loss, get_tokens_per_expert_and_token_count, + qb_dual_update, router_gating_linear, sinkhorn, switch_load_balancing_loss_func, @@ -216,6 +217,41 @@ def __init__( self.global_tokens_per_expert = None self.ga_steps = None + # Quantile balancing replaces the aux loss with a per-expert bias `qb_beta`. + # `qb_beta_accum`/`qb_beta_count` collect the per-microbatch quantile, reduced + # and reset each global batch. + if self.routing_type == "quantile_balancing": + assert not self.is_aux_loss_enabled(), ( + "Quantile balancing handles load balance via the bias update; " + "aux losses must be disabled (set moe_aux_loss_coeff to 0)." + ) + self.register_buffer( + 'qb_beta', + torch.zeros( + self.config.num_moe_experts, + dtype=torch.float32, + device=torch.cuda.current_device(), + ), + ) + self.register_buffer( + 'qb_beta_accum', + torch.zeros( + self.config.num_moe_experts, + dtype=torch.float32, + device=torch.cuda.current_device(), + ), + persistent=False, + ) + self.register_buffer( + 'qb_beta_count', + torch.zeros((), dtype=torch.long, device=torch.cuda.current_device()), + persistent=False, + ) + else: + self.qb_beta = None + self.qb_beta_accum = None + self.qb_beta_count = None + self.router_replay = None if self.config.moe_enable_routing_replay: self.router_replay = RouterReplay() @@ -230,6 +266,13 @@ def _maintain_float32_expert_bias(self): if hasattr(self, 'expert_bias') and self.expert_bias is not None: if self.expert_bias.dtype != torch.float32: self.expert_bias.data = self.expert_bias.data.to(torch.float32) + # Keep the QB bias in fp32 for the same reason. + if hasattr(self, 'qb_beta') and self.qb_beta is not None: + if self.qb_beta.dtype != torch.float32: + self.qb_beta.data = self.qb_beta.data.to(torch.float32) + if hasattr(self, 'qb_beta_accum') and self.qb_beta_accum is not None: + if self.qb_beta_accum.dtype != torch.float32: + self.qb_beta_accum.data = self.qb_beta_accum.data.to(torch.float32) def sinkhorn_load_balancing(self, logits: torch.Tensor): """Apply sinkhorn routing to the logits tensor. @@ -264,6 +307,81 @@ def _sinkhorn_activation(logits): scores = logits * map return scores, map + def quantile_balancing(self, logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Apply quantile-balancing (QB) routing to the logits tensor. + + Selects top-k experts per token using a dual coordinate-descent update on + a per-expert bias ``qb_beta``. Load balance is handled entirely by the bias + update; auxiliary losses must be disabled when QB is active. + + Args: + logits (torch.Tensor): The logits tensor, shape ``[num_tokens, num_experts]``. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Sparse routing probs and boolean + routing map, each shaped ``[num_tokens, num_experts]``. + """ + assert ( + not self.config.moe_router_fusion + ), "Quantile balancing routing does not support moe_router_fusion." + assert ( + self.config.moe_router_num_groups is None and self.config.moe_router_group_topk is None + ), "Quantile balancing routing does not support group-limited routing." + + local_num_tokens = logits.shape[0] + # Gather logits across TP/CP so the quantile sees a whole sequence's tokens. + # The DP reduction and qb_beta update run at the global-batch boundary in + # finalize_model_grads._update_router_qb_beta. + gather_group = self.tp_cp_group + gather_size = gather_group.size() if gather_group is not None else 1 + + should_update_beta = self.training and torch.is_grad_enabled() + + with torch.no_grad(): + logits_fp32 = logits.detach().to(dtype=torch.float32) + + if gather_size > 1: + full_logits = torch.empty( + (local_num_tokens * gather_size, self.config.num_moe_experts), + dtype=logits_fp32.dtype, + device=logits_fp32.device, + ) + torch.distributed.all_gather_into_tensor( + full_logits, logits_fp32.contiguous(), group=gather_group + ) + gather_rank = torch.distributed.get_rank(group=gather_group) + else: + full_logits = logits_fp32 + gather_rank = 0 + + # Route with the previous batch's qb_beta; in training, accumulate this + # microbatch's quantile for the next update. + full_indices, beta_local = qb_dual_update( + full_logits, self.topk, self.qb_beta, update_beta=should_update_beta + ) + if should_update_beta: + self.qb_beta_accum.add_(beta_local) + self.qb_beta_count.add_(1) + + # Take this rank's rows (all_gather orders rows by rank). + if gather_size > 1: + indices = full_indices[ + gather_rank * local_num_tokens : (gather_rank + 1) * local_num_tokens + ].contiguous() + else: + indices = full_indices + + # QB only picks the experts; reuse the shared score function for the probs. + return topk_routing_with_score_function( + logits, + self.topk, + use_pre_softmax=self.config.moe_router_pre_softmax, + scaling_factor=self.config.moe_router_topk_scaling_factor, + score_function=self.score_function, + fused=self.config.moe_router_fusion, + precomputed_indices=indices, + ) + def get_aux_loss_coeff(self, aux_loss_type: str) -> float: """Return the aux loss coeff for the given auxiliary loss type. If the auxiliary loss type is not found, return 0.0. @@ -641,6 +759,11 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N # Calculate probs and routing_map for token dispatching if self.routing_type == "sinkhorn": probs, routing_map = self.sinkhorn_load_balancing(logits) + elif self.routing_type == "quantile_balancing": + assert ( + padding_mask is None + ), "Quantile balancing routing does not support padding masks yet." + probs, routing_map = self.quantile_balancing(logits) else: probs, routing_map = topk_routing_with_score_function( logits, @@ -799,6 +922,7 @@ def _compiled_topk_routing( fused, router_replay, dense_output, + precomputed_indices, ): return topk_routing_with_score_function( logits, @@ -812,11 +936,17 @@ def _compiled_topk_routing( fused=fused, router_replay=router_replay, dense_output=dense_output, + precomputed_indices=precomputed_indices, ) def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): logits = self.gating(input).squeeze(1) # [num_tokens, num_experts] + # QB selects on (logits - qb_beta); at inference qb_beta is fixed, so it's per-token. + precomputed_indices = None + if self.qb_beta is not None: + precomputed_indices = (logits - self.qb_beta).topk(self.topk, dim=1).indices + probs, top_indices = self._compiled_topk_routing( logits, self.topk, @@ -829,6 +959,7 @@ def _forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = N fused=self.config.moe_router_fusion, router_replay=self.router_replay, dense_output=True, + precomputed_indices=precomputed_indices, ) return probs.squeeze(1), top_indices.squeeze(1) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index f8779d674e2..761504e614e 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -718,6 +718,9 @@ class TransformerConfig(ModelParallelConfig): for each individual sample. - "global_aux_loss": Load balancing loss calculated at global batch level. - "sinkhorn": Balancing algorithm used in S-BASE. + - "quantile_balancing": Dual coordinate-descent quantile balancing (QB). Load balance is + handled entirely by an internal per-expert bias update; auxiliary losses must be disabled + (`moe_aux_loss_coeff` = 0) when QB is selected. - "none": No load balancing. A list of strings can be provided to combine multiple aux-loss load balancing types. The default is "aux_loss". @@ -790,6 +793,12 @@ class TransformerConfig(ModelParallelConfig): and decreased for the experts with more assigned tokens. The default value 1e-3 is same as that used in DeepSeekV3.""" + moe_router_quantile_balancing_ema: float = 0.0 + """EMA coefficient for the quantile-balancing per-expert bias (`qb_beta`), used only when + `moe_router_load_balancing_type` is "quantile_balancing". At each global batch the bias is + updated as `qb_beta = ema * qb_beta + (1 - ema) * local_quantile`. The default 0.0 means + no memory: the bias is replaced by the latest global-batch quantile estimate each step.""" + moe_router_force_load_balancing: bool = False """[Experimental] Force load balancing with random logits for MoE router, supports naive topk and group-limited topk. This is an experimental feature and only for benchmark.""" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 4f3c0af7b30..2c3cd1cd531 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3205,9 +3205,9 @@ def _add_moe_args(parser): 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.') # Router arguments group.add_argument('--moe-router-load-balancing-type', nargs='+', type=str, - choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'], + choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'quantile_balancing', 'none'], default='aux_loss', - help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".') + help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE; "quantile_balancing" (QB) uses dual coordinate descent on a per-expert bias to handle load balance internally; "none" implies no load balancing. The default is "aux_loss".') group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') # Token dispatcher arguments diff --git a/tests/unit_tests/distributed/test_finalize_model_grads.py b/tests/unit_tests/distributed/test_finalize_model_grads.py index ee535c29baf..80d143a89a3 100644 --- a/tests/unit_tests/distributed/test_finalize_model_grads.py +++ b/tests/unit_tests/distributed/test_finalize_model_grads.py @@ -11,13 +11,21 @@ from megatron.core.distributed.finalize_model_grads import ( _allreduce_non_tensor_model_parallel_grads, _allreduce_word_embedding_grads, + _update_router_qb_beta, finalize_model_grads, + reset_model_temporary_tensors, +) +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_local_submodules, + get_gpt_layer_with_transformer_engine_spec, ) -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.moe.moe_layer import MoELayer +from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.initialize import _set_random_seed from tests.unit_tests.test_utilities import Utils @@ -117,6 +125,98 @@ def test_finalize_model_grads_requires_custom_group_before_grad_sync(self): assert model.finish_grad_sync_calls == 0 +class TestUpdateRouterQBBeta: + """Exercises the QB bias update in finalize_model_grads against a real MoE router.""" + + def setup_method(self, method): + os.environ.pop('NVTE_FUSED_ATTN', None) + os.environ.pop('NVTE_FLASH_ATTN', None) + os.environ.pop('NVTE_UNFUSED_ATTN', None) + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(1, 1) + _set_random_seed(seed_=123, data_parallel_random_init=False) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _build_moe_layer(self, ema): + num_experts = 8 + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + num_moe_experts=num_experts, + use_cpu_initialization=True, + moe_router_load_balancing_type="quantile_balancing", + moe_router_score_function="softmax", + moe_router_topk=2, + moe_aux_loss_coeff=0, + moe_router_quantile_balancing_ema=ema, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + ) + submodules = get_submodules( + get_gpt_layer_local_submodules(num_experts=num_experts, moe_grouped_gemm=False).mlp + ) + return config, MoELayer(config, submodules).cuda() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("ema", [0.0, 0.9]) + def test_update_router_qb_beta(self, ema): + config, moe_layer = self._build_moe_layer(ema) + router = moe_layer.router + router.train() + # Non-zero prior bias so the EMA term is actually exercised. + router.qb_beta.copy_(torch.randn_like(router.qb_beta)) + + # The real router forward populates qb_beta_accum / qb_beta_count. + hidden = torch.randn((32, 2, config.hidden_size)).cuda().bfloat16() + router(hidden) + router(hidden) + assert router.qb_beta_count.item() == 2 + assert router.qb_beta_accum.abs().sum().item() > 0 + + # Expected from the real accumulators: DP-avg(accum/count), EMA-blend, re-center. + local_avg = router.qb_beta_accum / router.qb_beta_count.clamp(min=1).to(torch.float32) + torch.distributed.all_reduce( + local_avg, op=torch.distributed.ReduceOp.AVG, group=dist.group.WORLD + ) + blended = ema * router.qb_beta + (1.0 - ema) * local_avg + expected = blended - blended.mean(dim=-1, keepdim=True) + + _update_router_qb_beta([moe_layer], config, dp_cp_group=dist.group.WORLD) + + torch.testing.assert_close(router.qb_beta, expected) + torch.testing.assert_close( + router.qb_beta.mean(), torch.zeros((), device=router.qb_beta.device) + ) + + # reset_model_temporary_tensors clears the accumulators for the next global batch. + reset_model_temporary_tensors(config, [moe_layer]) + torch.testing.assert_close(router.qb_beta_accum, torch.zeros_like(router.qb_beta_accum)) + assert router.qb_beta_count.item() == 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_update_router_qb_beta_skips_eval(self): + config, moe_layer = self._build_moe_layer(ema=0.0) + router = moe_layer.router + # Non-zero prior + non-uniform accumulator, so a broken eval guard would visibly + # change qb_beta (a uniform accumulator re-centers to zero and hides the bug). + router.qb_beta.copy_(torch.ones_like(router.qb_beta)) + router.qb_beta_accum.copy_( + torch.arange(router.qb_beta.numel(), dtype=torch.float32, device=router.qb_beta.device) + ) + router.qb_beta_count.fill_(1) + before = router.qb_beta.clone() + router.eval() + + _update_router_qb_beta([moe_layer], config, dp_cp_group=dist.group.WORLD) + + # Eval-mode modules are skipped, so qb_beta is unchanged. + torch.testing.assert_close(router.qb_beta, before) + + class TestAllReduceLNGrads: def init_model(self, share_embeddings_and_output_weights: bool = False): diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 194cb2a285b..f7dc78ce9a2 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -211,6 +211,7 @@ "moe_router_padding_for_fp8": False, "moe_router_padding_for_quantization": False, "moe_router_pre_softmax": False, + "moe_router_quantile_balancing_ema": 0.0, "moe_router_score_function": "sigmoid", "moe_router_topk": 6, "moe_router_topk_limited_devices": None, diff --git a/tests/unit_tests/transformer/moe/test_qb_routing.py b/tests/unit_tests/transformer/moe/test_qb_routing.py new file mode 100644 index 00000000000..5ad26487fb6 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_qb_routing.py @@ -0,0 +1,139 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +from typing import cast + +import pytest +import torch + +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules +from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules +from megatron.core.transformer.moe.moe_utils import qb_dual_update +from megatron.core.transformer.moe.router import Router +from megatron.core.transformer.spec_utils import get_submodules +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.initialize import _set_random_seed +from tests.unit_tests.test_utilities import Utils + + +class TestQBDualUpdate: + """Pure-tensor tests for the quantile-balancing dual update (CPU, no distributed).""" + + @pytest.mark.internal + @pytest.mark.parametrize("m,n,k", [(64, 8, 2), (40, 8, 1), (12, 4, 1)]) + def test_column_quantile_contract(self, m, n, k): + """qb_beta_local is the (col_target+1)-th largest score minus alpha per expert.""" + torch.manual_seed(123) + scores = torch.randn(m, n) + beta = torch.zeros(n) + + _, beta_local = qb_dual_update(scores, k, beta, update_beta=True) + + alpha = (scores - beta).topk(k + 1, dim=1).values[:, -1:] + adjusted = scores - alpha + col_target = m * k // n + expected = adjusted.sort(dim=0, descending=True).values[col_target] + torch.testing.assert_close(beta_local, expected) + + +class TestQuantileBalancingRouter: + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + _set_random_seed(seed_=123, data_parallel_random_init=False) + self.num_moe_experts = 8 + self.transformer_config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + num_moe_experts=self.num_moe_experts, + use_cpu_initialization=True, + moe_router_load_balancing_type="quantile_balancing", + moe_router_score_function="softmax", + moe_router_topk=2, + moe_aux_loss_coeff=0, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + ) + self.submodules = get_submodules( + get_gpt_layer_local_submodules( + num_experts=self.num_moe_experts, moe_grouped_gemm=False + ).mlp + ) + assert isinstance(self.submodules, MoESubmodules) + self.moe_layer = MoELayer(self.transformer_config, self.submodules) + self.router = cast(Router, self.moe_layer.router) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.internal + def test_non_qb_router_has_no_qb_buffers(self): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + num_moe_experts=self.num_moe_experts, + use_cpu_initialization=True, + moe_router_load_balancing_type="aux_loss", + moe_router_topk=2, + moe_aux_loss_coeff=0, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + ) + router = MoELayer(config, self.submodules).router + assert router.qb_beta is None + assert router.qb_beta_accum is None + assert router.qb_beta_count is None + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("moe_router_pre_softmax", [True, False]) + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid"]) + def test_qb_router_forward(self, score_function, moe_router_pre_softmax): + self.router = self.router.cuda() + self.router.config.moe_router_score_function = score_function + self.router.score_function = score_function + self.router.config.moe_router_pre_softmax = moe_router_pre_softmax + + num_tokens = 32 * 2 + hidden_states = torch.randn((32, 2, self.router.config.hidden_size)).cuda().bfloat16() + with torch.no_grad(): + probs, routing_map = self.router(hidden_states) + + assert probs.shape == (num_tokens, self.num_moe_experts) + assert routing_map.shape == (num_tokens, self.num_moe_experts) + # Each token selects exactly topk distinct experts. + assert routing_map.sum().item() == num_tokens * self.router.topk + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_qb_beta_accumulates_in_training(self): + self.router = self.router.cuda() + self.router.train() + hidden_states = torch.randn((32, 2, self.router.config.hidden_size)).cuda().bfloat16() + + assert self.router.qb_beta_count.item() == 0 + self.router(hidden_states) + assert self.router.qb_beta_count.item() == 1 + assert self.router.qb_beta_accum.abs().sum().item() > 0 + self.router(hidden_states) + assert self.router.qb_beta_count.item() == 2 + + # No accumulation outside the training path (eval / recompute). + accum_before = self.router.qb_beta_accum.clone() + with torch.no_grad(): + self.router(hidden_states) + assert self.router.qb_beta_count.item() == 2 + torch.testing.assert_close(self.router.qb_beta_accum, accum_before) + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_qb_router_rejects_padding_mask(self): + self.router = self.router.cuda() + hidden_states = torch.randn((32, 2, self.router.config.hidden_size)).cuda().bfloat16() + padding_mask = torch.zeros((32, 2), dtype=torch.bool, device=hidden_states.device) + padding_mask[-2:] = True + + with pytest.raises(AssertionError, match="does not support padding masks"): + self.router(hidden_states, padding_mask=padding_mask) diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index 9f33dd01920..b215b59cfa0 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -8,7 +8,11 @@ from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules -from megatron.core.transformer.moe.moe_utils import get_updated_expert_bias, router_gating_linear +from megatron.core.transformer.moe.moe_utils import ( + get_updated_expert_bias, + router_gating_linear, + topk_routing_with_score_function, +) from megatron.core.transformer.moe.router import Router from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig @@ -574,3 +578,40 @@ def test_router_gating_linear_bias(router_dtype): assert torch.allclose(inp.grad, ref_inp.grad, **tols) assert torch.allclose(weight.grad, ref_weight.grad, **tols) assert torch.allclose(bias.grad, ref_bias.grad, **tols) + + +@pytest.mark.internal +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) +@pytest.mark.parametrize("use_pre_softmax", [True, False]) +@pytest.mark.parametrize("topk", [1, 2]) +def test_topk_routing_precomputed_indices_equivalence(score_function, use_pre_softmax, topk): + """Passing precomputed_indices that match the function's own selection must reproduce + the standard output. Guards the shared post-top-k path reused by quantile balancing.""" + if score_function != "softmax" and use_pre_softmax: + pytest.skip("pre_softmax only applies to softmax scoring") + + torch.manual_seed(123) + num_tokens, num_experts = 64, 8 + logits = torch.randn(num_tokens, num_experts) + + kwargs = dict(use_pre_softmax=use_pre_softmax, score_function=score_function, fused=False) + probs_ref, map_ref = topk_routing_with_score_function(logits, topk, **kwargs) + _, top_indices = topk_routing_with_score_function(logits, topk, dense_output=True, **kwargs) + probs_pre, map_pre = topk_routing_with_score_function( + logits, topk, precomputed_indices=top_indices, **kwargs + ) + + # Natural top-k indices reproduce the standard output. + assert torch.equal(map_ref, map_pre) + torch.testing.assert_close(probs_ref, probs_pre) + + # Indices that differ from the natural top-k must route to exactly those experts. This + # catches a regression where the precomputed_indices branch is dropped and the function + # silently recomputes its own top-k instead of honoring the caller's indices. Bottom-k is + # disjoint from top-k since 2 * topk <= num_experts. + alt_indices = logits.topk(topk, dim=1, largest=False).indices + _, map_alt = topk_routing_with_score_function( + logits, topk, precomputed_indices=alt_indices, **kwargs + ) + expected_map = torch.zeros_like(logits, dtype=torch.bool).scatter(1, alt_indices, True) + assert torch.equal(map_alt, expected_map) diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py index c4429220181..8c84aae7097 100644 --- a/tests/unit_tests/transformer/moe/test_shared_experts.py +++ b/tests/unit_tests/transformer/moe/test_shared_experts.py @@ -10,12 +10,12 @@ from megatron.core.models.gpt import moe_module_specs from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules from megatron.core.parallel_state import get_tensor_model_parallel_world_size -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.moe import shared_experts as shared_experts_module from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules from megatron.core.transformer.moe.shared_experts import FusedSharedExpertMLP, SharedExpertMLP from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.initialize import _set_random_seed from tests.unit_tests.test_utilities import Utils @@ -353,13 +353,13 @@ def test_shared_expert_forward_backward(self, dispatcher_type: str, tp_size, ep_ tensor_model_parallel_size=tp_size, expert_model_parallel_size=ep_size ) # Create MoE layer with shared expert overlap enabled. - model_parallel_cuda_manual_seed(123) + _set_random_seed(seed_=123, data_parallel_random_init=False) moe_layer_overlap = self.get_moe_layer( moe_shared_expert_overlap=True, moe_token_dispatcher_type=dispatcher_type ).to(dtype=torch.bfloat16) # Create MoE layer with shared expert overlap disabled. - model_parallel_cuda_manual_seed(123) + _set_random_seed(seed_=123, data_parallel_random_init=False) moe_layer_no_overlap = self.get_moe_layer( moe_shared_expert_overlap=False, moe_token_dispatcher_type=dispatcher_type ).to(dtype=torch.bfloat16)