[MCore] Add global-batch quantile balancing

megatron/core/distributed/finalize_model_grads.py

@@ -13,6 +13,7 @@
 except ImportError:
     HAVE_DTENSOR = False
 
+from megatron.core.extensions.transformer_engine import mark_qb_bin_bounds_validated
 from megatron.core.pipeline_parallel.utils import (
     get_pp_last_rank,
     is_pp_first_stage,
@@ -21,7 +22,10 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 
 from .. import parallel_state
-from ..transformer.moe.moe_utils import get_updated_expert_bias
+from ..transformer.moe.moe_utils import (
+    get_updated_expert_bias,
+    get_updated_expert_bias_with_quantile,
+)
 from ..transformer.transformer_config import TransformerConfig
 from ..utils import (
     get_attr_wrapped_model,
@@ -329,6 +333,8 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n
             if getattr(module, 'qb_beta_accum', None) is not None:
                 module.qb_beta_accum.zero_()
                 module.qb_beta_count.zero_()
+            if getattr(module, 'qb_histogram', None) is not None:
+                module.qb_histogram.zero_()
 
 
 def _update_router_expert_bias(
@@ -413,6 +419,47 @@ def _update_router_qb_beta(
         qb_beta.copy_(new_beta)
 
 
+def _update_router_expert_bias_with_quantile(
+    model: List[torch.nn.Module],
+    config: TransformerConfig,
+    reduction_groups: tuple[Optional[torch.distributed.ProcessGroup], ...] = (),
+):
+    """Pool K3 histograms and update every local router once per global batch."""
+    expert_bias_list = []
+    qb_histogram_list = []
+    qb_bin_bounds_list = []
+    for model_chunk in model:
+        for module in get_attr_wrapped_model(model_chunk, 'modules')():
+            if (
+                getattr(module, 'qb_histogram', None) is not None
+                and module.training
+                and not getattr(module, 'frozen_expert_bias', False)
+            ):
+                expert_bias_list.append(module.expert_bias)
+                qb_histogram_list.append(module.qb_histogram)
+                qb_bin_bounds_list.append(module.qb_bin_bounds)
+
+    if not expert_bias_list:
+        return
+
+    stacked_bias = torch.stack(expert_bias_list, dim=0)
+    stacked_histogram = torch.stack(qb_histogram_list, dim=0)
+    stacked_bin_bounds = torch.stack(qb_bin_bounds_list, dim=0)
+    for group in reduction_groups:
+        if get_pg_size(group) > 1:
+            torch.distributed.all_reduce(stacked_histogram, group=group)
+    updated_bias, updated_bin_bounds = get_updated_expert_bias_with_quantile(
+        stacked_histogram, stacked_bin_bounds, stacked_bias, config.moe_router_topk
+    )
+    for bias, bounds, next_bias, next_bounds in zip(
+        expert_bias_list, qb_bin_bounds_list, updated_bias, updated_bin_bounds
+    ):
+        bias.copy_(next_bias)
+        bounds.copy_(next_bounds)
+        if mark_qb_bin_bounds_validated is not None and bounds.is_cuda:
+            mark_qb_bin_bounds_validated(bounds)
+
+
 def _allreduce_non_tensor_model_parallel_grads(
     model: List[torch.nn.Module],
     config: TransformerConfig,
@@ -589,10 +636,14 @@ def finalize_model_grads(
             "If you don't need pos_embd_group, you need to explicitly set it to None."
         )
         assert hasattr(pg_collection, 'dp_cp')
-        if config.moe_router_enable_expert_bias:
-            assert hasattr(pg_collection, 'tp_dp_cp') and pg_collection.tp_dp_cp is not None, (
-                "pg_collection must have tp_dp_cp when " "moe_router_enable_expert_bias is enabled."
-            )
+        if config.moe_router_enable_expert_bias or (
+            config.moe_router_load_balancing_type == "quantile_balancing"
+            and config.moe_router_quantile_balancing_estimation_scope == "global_batch"
+            and config.gtp_weight_remat_size <= 1
+        ):
+            assert (
+                hasattr(pg_collection, 'tp_dp_cp') and pg_collection.tp_dp_cp is not None
+            ), "pg_collection must have tp_dp_cp when global router bias updates are enabled."
             tp_dp_cp_group = pg_collection.tp_dp_cp
         tp_group = pg_collection.tp
         pp_group = pg_collection.pp
@@ -687,7 +738,21 @@ 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)
+        if config.moe_router_quantile_balancing_estimation_scope == "micro_batch":
+            _update_router_qb_beta(model, config, dp_cp_group=dp_cp_group)
+        else:
+            if config.gtp_weight_remat_size > 1:
+                qb_reduction_groups = (tp_group, dp_cp_group)
+            else:
+                if pg_collection is None:
+                    # Legacy compatibility; modern callers provide pg_collection.tp_dp_cp above.
+                    tp_dp_cp_group = parallel_state.get_tensor_and_data_parallel_group(
+                        with_context_parallel=True
+                    )
+                qb_reduction_groups = (tp_dp_cp_group,)
+            _update_router_expert_bias_with_quantile(
+                model, config, reduction_groups=qb_reduction_groups
+            )
 
     reset_model_temporary_tensors(config, model)
 

megatron/core/extensions/transformer_engine.py

@@ -3673,10 +3673,30 @@ def te_general_gemm(
         fused_topk_with_score_function,
     )
 
+    try:
+        from transformer_engine.pytorch.router import (  # pylint: disable=unused-import
+            mark_qb_bin_bounds_validated,
+        )
+    except ImportError:
+        mark_qb_bin_bounds_validated = None
+
+    try:
+        _fused_topk_sig = inspect.signature(fused_topk_with_score_function)
+        fused_topk_with_score_function_supports_qb = {
+            "qb_histogram",
+            "qb_bin_bounds",
+            "qb_histogram_mode",
+        }.issubset(_fused_topk_sig.parameters)
+        del _fused_topk_sig
+    except (TypeError, ValueError):
+        fused_topk_with_score_function_supports_qb = False
+
 else:
     fused_topk_with_score_function = None
     fused_compute_score_for_moe_aux_loss = None
     fused_moe_aux_loss = None
+    fused_topk_with_score_function_supports_qb = False
+    mark_qb_bin_bounds_validated = None
 
 
 def set_save_original_input(module):

megatron/core/transformer/moe/README.md

@@ -297,6 +297,7 @@ Routers determine which expert(s) handle each token. A lightweight MLP scores ev
 | **seq_aux_loss** | Sequence-level auxiliary loss for balancing expert usage on each sequence| `--moe-router-load-balancing-type seq_aux_loss` |
 | **global_aux_loss** | Global auxiliary loss for balancing expert usage on a global batch across all ranks | `--moe-router-load-balancing-type global_aux_loss` |
 | **sinkhorn** | Optimal transport formulation for balancing expert usage | `--moe-router-load-balancing-type sinkhorn` |
+| **quantile_balancing** | Exact micro-batch or Kimi K3 global-batch histogram quantile bias updates | `--moe-router-load-balancing-type quantile_balancing --moe-router-score-function sigmoid --moe-aux-loss-coeff 0` |
 | **aux loss free** | Dynamic bias-based load balancing strategy without auxiliary loss | `--moe-router-enable-expert-bias --moe-router-bias-update-rate 1e-3`|
 | **none** | No load balancing | `--moe-router-load-balancing-type none` |
 
@@ -526,7 +527,9 @@ For MoE models, certain configurations may prevent CUDA Graph capture of MoE lay
 ### Router Arguments
 | Argument | Description | Default |
 |----------|-------------|---------|
-| --moe-router-load-balancing-type | Load balancing: aux_loss, sinkhorn, seq_aux_loss, none | aux_loss |
+| --moe-router-load-balancing-type | Load balancing: aux_loss, seq_aux_loss, global_aux_loss, sinkhorn, quantile_balancing, none | aux_loss |
+| --moe-router-quantile-balancing-estimation-scope | Quantile population: micro_batch or global_batch | micro_batch |
+| --moe-router-qb-num-bins | Uniform histogram bins per expert for global-batch quantile balancing | 1000 |
 | --moe-router-topk | Number of experts per token | 2 |
 | --moe-router-score-function | Score function: softmax, sigmoid | softmax |
 | --moe-router-pre-softmax | Softmax before top-k | False |

megatron/core/transformer/moe/moe_utils.py

@@ -42,6 +42,7 @@
         fused_sort_chunks_by_index,
         fused_sort_chunks_by_index_with_probs,
         fused_topk_with_score_function,
+        fused_topk_with_score_function_supports_qb,
         fused_unpermute,
         te_general_gemm,
     )
@@ -58,6 +59,7 @@
         fused_unpermute,
         te_general_gemm,
     ) = (None, None, None, None, None, None, None, None, None, None)
+    fused_topk_with_score_function_supports_qb = False
 
 
 def switch_load_balancing_loss_func(
@@ -776,6 +778,8 @@ def topk_routing_with_score_function(
     router_replay: Optional['RouterReplay'] = None,
     dense_output: bool = False,
     precomputed_indices: Optional[torch.Tensor] = None,
+    qb_histogram: Optional[torch.Tensor] = None,
+    qb_bin_bounds: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute the routing probabilities and map for top-k selection with score function.
 
@@ -804,6 +808,10 @@ def topk_routing_with_score_function(
                                        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.
+        qb_histogram (torch.Tensor, optional): Caller-owned int32 K3 Quantile Balancing histogram
+                                               with shape [num_experts, num_bins].
+        qb_bin_bounds (torch.Tensor, optional): FP32 CUDA tensor containing the lower and upper
+                                                K3 Quantile Balancing histogram bounds.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]:
@@ -822,6 +830,33 @@ 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
+    use_quantile_balancing = qb_histogram is not None or qb_bin_bounds is not None
+    if use_quantile_balancing and (qb_histogram is None or qb_bin_bounds is None):
+        raise ValueError("qb_histogram and qb_bin_bounds must be provided together.")
+    if use_quantile_balancing:
+        if expert_bias is None:
+            raise ValueError("Quantile Balancing requires an expert bias.")
+        if score_function != "sigmoid":
+            raise ValueError("Quantile Balancing currently requires score_function='sigmoid'.")
+        if use_pre_softmax:
+            raise ValueError("Quantile Balancing does not use pre-softmax routing.")
+        if topk >= num_experts:
+            raise ValueError("Quantile Balancing requires topk < num_experts.")
+        if num_groups is not None or group_topk is not None:
+            raise ValueError("Quantile Balancing does not support group-limited routing.")
+        if router_replay is not None:
+            raise ValueError("Quantile Balancing does not support router replay.")
+        if precomputed_indices is not None:
+            raise ValueError("Histogram Quantile Balancing computes its own top-(k+1).")
+        if qb_histogram.dim() != 2 or qb_histogram.shape[0] != num_experts:
+            raise ValueError(
+                "qb_histogram must have shape [num_experts, num_bins], got "
+                f"{tuple(qb_histogram.shape)}."
+            )
+        if qb_histogram.dtype != torch.int32:
+            raise ValueError("qb_histogram must have dtype torch.int32.")
+        if qb_bin_bounds.shape != (2,) or qb_bin_bounds.dtype != torch.float32:
+            raise ValueError("qb_bin_bounds must be an FP32 tensor with shape [2].")
     assert not (
         fused and precomputed_indices is not None
     ), "precomputed_indices is not supported with the fused top-k score function."
@@ -835,16 +870,29 @@ def topk_routing_with_score_function(
                 "Fused sqrtsoftplus score function requires TE >= 2.13.0. "
                 "Please upgrade Transformer Engine or disable moe_router_fusion."
             )
-        return fused_topk_with_score_function(
-            logits=logits,
-            topk=topk,
-            use_pre_softmax=use_pre_softmax,
-            num_groups=num_groups,
-            group_topk=group_topk,
-            scaling_factor=scaling_factor,
-            score_function=score_function,
-            expert_bias=expert_bias,
-        )
+        kwargs = {
+            "logits": logits,
+            "topk": topk,
+            "use_pre_softmax": use_pre_softmax,
+            "num_groups": num_groups,
+            "group_topk": group_topk,
+            "scaling_factor": scaling_factor,
+            "score_function": score_function,
+            "expert_bias": expert_bias,
+        }
+        if use_quantile_balancing:
+            if not fused_topk_with_score_function_supports_qb:
+                raise ValueError(
+                    "The installed Transformer Engine fused router does not expose Quantile "
+                    "Balancing histogram outputs. Upgrade Transformer Engine or disable "
+                    "moe_router_fusion."
+                )
+            kwargs.update(
+                qb_histogram=qb_histogram,
+                qb_bin_bounds=qb_bin_bounds,
+                qb_histogram_mode="fused_atomic",
+            )
+        return fused_topk_with_score_function(**kwargs)
 
     def _compute_topk(
         scores: torch.Tensor,
@@ -922,7 +970,27 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
             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)
+            if use_quantile_balancing:
+                topk_result = torch.topk(scores_for_routing, topk + 1, dim=1, sorted=True)
+                cutoff = topk_result.values[:, -1:]
+                top_indices = topk_result.indices[:, :-1]
+                with torch.no_grad():
+                    num_bins = qb_histogram.shape[1]
+                    lower, upper = qb_bin_bounds.unbind()
+                    bin_indices = torch.floor(
+                        (cutoff - scores.detach() - lower) * (num_bins / (upper - lower))
+                    ).to(torch.int64)
+                    bin_indices.clamp_(0, num_bins - 1)
+                    expert_offsets = (
+                        torch.arange(num_experts, device=logits.device, dtype=torch.int64)
+                        * num_bins
+                    )
+                    flat_indices = (bin_indices + expert_offsets).reshape(-1)
+                    qb_histogram.view(-1).scatter_add_(
+                        0, flat_indices, torch.ones_like(flat_indices, dtype=torch.int32)
+                    )
+            else:
+                _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)
             scores = torch.gather(scores, dim=1, index=top_indices)
         else:
             scores, top_indices = compute_topk(scores, topk, num_groups, group_topk)
@@ -1229,6 +1297,61 @@ def get_updated_expert_bias(
         return updated_expert_bias
 
 
+def get_updated_expert_bias_with_quantile(
+    histogram: torch.Tensor, bin_bounds: torch.Tensor, expert_bias: torch.Tensor, topk: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Recover Kimi K3 Quantile Balancing biases from a pooled global-batch histogram."""
+    if histogram.dim() < 2:
+        raise ValueError("QB histogram must have shape [..., num_experts, num_bins].")
+    num_experts, num_bins = histogram.shape[-2:]
+    if num_experts <= 0 or num_bins <= 0:
+        raise ValueError("QB histogram dimensions must be positive.")
+    if expert_bias.shape != histogram.shape[:-1]:
+        raise ValueError(
+            f"QB expert_bias shape {expert_bias.shape} does not match expected "
+            f"{histogram.shape[:-1]} for histogram shape {histogram.shape}."
+        )
+    if bin_bounds.shape != histogram.shape[:-2] + (2,):
+        raise ValueError(
+            f"QB bin_bounds shape {bin_bounds.shape} does not match expected "
+            f"{histogram.shape[:-2] + (2,)} for histogram shape {histogram.shape}."
+        )
+    if not 0 < topk < num_experts:
+        raise ValueError(f"QB topk must be in [1, {num_experts}), got {topk}.")
+
+    with torch.no_grad():
+        cumulative_counts = torch.cumsum(histogram, dim=-1, dtype=torch.int64)
+        tokens_per_expert = cumulative_counts[..., -1]
+        # Every token contributes one margin sample to every expert, so all expert totals match.
+        total_tokens = tokens_per_expert[..., 0]
+        target_quantile = total_tokens.to(torch.float32) * (topk / num_experts)
+        target_rank = torch.ceil(target_quantile).to(torch.int64).clamp_min_(1)
+        selected_bins = (
+            (cumulative_counts >= target_rank[..., None, None]).to(torch.int64).argmax(dim=-1)
+        )
+        selected_counts = torch.gather(histogram, -1, selected_bins.unsqueeze(-1)).squeeze(-1)
+        previous_bins = (selected_bins - 1).clamp_min(0)
+        counts_before = torch.gather(cumulative_counts, -1, previous_bins.unsqueeze(-1)).squeeze(-1)
+        counts_before = torch.where(
+            selected_bins == 0, torch.zeros_like(counts_before), counts_before
+        )
+        interpolation = (
+            (target_quantile[..., None] - counts_before.to(torch.float32))
+            / selected_counts.clamp_min(1).to(torch.float32)
+        ).clamp_(0.0, 1.0)
+        lower = bin_bounds[..., 0, None]
+        upper = bin_bounds[..., 1, None]
+        bin_width = (upper - lower) / num_bins
+        updated_expert_bias = lower + (selected_bins.to(torch.float32) + interpolation) * bin_width
+        updated_expert_bias -= updated_expert_bias.mean(dim=-1, keepdim=True)
+        has_tokens = total_tokens > 0
+        updated_expert_bias = torch.where(has_tokens[..., None], updated_expert_bias, expert_bias)
+        bias_min, bias_max = torch.aminmax(updated_expert_bias, dim=-1)
+        updated_bin_bounds = torch.stack((bias_min - 1.0, bias_max + 1.0), dim=-1)
+        updated_bin_bounds = torch.where(has_tokens[..., None], updated_bin_bounds, bin_bounds)
+        return updated_expert_bias, updated_bin_bounds
+
+
 def maybe_move_tensor_to_cpu(
     tensor: torch.Tensor, as_numpy: bool = False, record_stream: bool = False
 ) -> torch.Tensor:

megatron/core/transformer/moe/paged_stash.py

@@ -1053,6 +1053,14 @@ def _set_moe_paged_stash_all(self, value: bool) -> None:
         for c in self._configs_to_sync_moe_paged_stash:
             c.moe_paged_stash = value
 
+    def _reset_qb_histograms(self) -> None:
+        """Discard histogram observations from a failed full-iteration attempt."""
+        for mlp in self.moe_layers:
+            router = getattr(mlp, 'router', None)
+            qb_histogram = getattr(router, 'qb_histogram', None)
+            if qb_histogram is not None:
+                qb_histogram.zero_()
+
     def data_read(self, data_iterator, model, training, num_microbatches):
         """Read all microbatch inputs from Dataloader and copy to static buffers."""
         data_iterator_saved = []
@@ -1168,6 +1176,11 @@ def prepare_for_rerun(self, is_training=True):
             self.stash_manager.host_spill.zero_()
         self._set_moe_paged_stash_all(False)
 
+        # The dropless retry replays every microbatch. Preserve CUDA addresses while
+        # discarding the failed attempt so each token contributes exactly once.
+        if is_training:
+            self._reset_qb_histograms()
+
         # Set grad to zero.
         for model_chunk in self.model:
             model_chunk.zero_grad_buffer()