Add an opt-in route_norm flag for the legacy MoE gate

top1gating leaves the chosen expert's softmax probability in the combine
weight. top2gating and topkgating rescale the kept weights so they sum to 1.
Making those agree changes the MoE output for existing checkpoints, so
route_norm defaults to None and each gate keeps its historical rule.
True or False overrides all three gates. Auxiliary loss is unchanged.

Signed-off-by: ARAVINDHAN T <arvindhant01@gmail.com>

Author: Aravind-11

deepspeed/moe/layer.py

@@ -33,6 +33,12 @@ class MoE(nn.Module):
         use_tutel (bool, optional): default=False, whether to use Tutel optimizations (if installed).
         enable_expert_tensor_parallelism (bool, optional): default=False, whether to use tensor parallelism for experts
         top2_2nd_expert_sampling (bool, optional): default=True, whether to perform sampling for 2nd expert
+        route_norm (bool, optional): default=None. None keeps the historical combine weights: k=1 leaves the chosen
+            expert's softmax probability, and k>=2 rescales the kept weights so they sum to 1. True rescales for every k.
+            False leaves the raw softmax probabilities for every k. True or False changes the MoE output relative to a
+            checkpoint trained with the other setting. The auxiliary load-balancing loss is computed before this step
+            and does not change. This is the legacy gate. It is not expert_parallel.route_norm, which configures the
+            AutoEP router.
     """
 
     def __init__(self,
@@ -50,7 +56,8 @@ def __init__(self,
                  use_rts: bool = True,
                  use_tutel: bool = False,
                  enable_expert_tensor_parallelism: bool = False,
-                 top2_2nd_expert_sampling: bool = True) -> None:
+                 top2_2nd_expert_sampling: bool = True,
+                 route_norm: Optional[bool] = None) -> None:
 
         super(MoE, self).__init__()
 
@@ -72,7 +79,7 @@ def __init__(self,
         experts = Experts(expert, self.num_local_experts, self.expert_group_name)
         self.deepspeed_moe = MOELayer(TopKGate(hidden_size, num_experts, k, capacity_factor, eval_capacity_factor,
                                                min_capacity, noisy_gate_policy, drop_tokens, use_rts, None,
-                                               top2_2nd_expert_sampling),
+                                               top2_2nd_expert_sampling, route_norm),
                                       experts,
                                       self.expert_group_name,
                                       self.ep_size,

deepspeed/moe/sharded_moe.py

@@ -232,6 +232,13 @@ def _sparse_decode(expert_output: Tensor, slots: Tensor, gates: Tensor, num_toke
     return combined.to(expert_output.dtype)
 
 
+def _route_norm_enabled(route_norm: Optional[bool], legacy_default: bool) -> bool:
+    """None keeps that gate's historical combine-weight rule. True or False overrides every gate."""
+    if route_norm is None:
+        return legacy_default
+    return route_norm
+
+
 def top1gating(logits: Tensor,
                capacity_factor: float,
                min_capacity: int,
@@ -241,8 +248,13 @@ def top1gating(logits: Tensor,
                use_rts: bool = True,
                ep_group: Union[torch.distributed.ProcessGroup, None] = None,
                sparse_routes: bool = False,
-               use_tutel: bool = False) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    """Implements Top1Gating on logits."""
+               use_tutel: bool = False,
+               route_norm: Optional[bool] = None) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """Implements Top1Gating on logits.
+
+    route_norm=None keeps the historical weights: the chosen expert's softmax probability, with no
+    second normalization. The auxiliary loss is computed from those probabilities either way.
+    """
     if noisy_gate_policy == 'RSample':
         logits_w_noise = logits + gumbel_rsample(logits.shape, device=logits.device)
     # everything is in fp32 in this function
@@ -319,16 +331,25 @@ def top1gating(logits: Tensor,
 
     if sparse_routes:
         gates1_s = (gates * mask1).sum(dim=1)
+        # Historical top-1 does not renormalize. An explicit True makes a kept token's weight 1.
+        # A dropped token sums to 0, and 0 / eps stays 0.
+        if _route_norm_enabled(route_norm, False):
+            denom_s = torch.clamp(gates1_s, min=torch.finfo(gates1_s.dtype).eps)
+            gates1_s = gates1_s / denom_s
         locations1_s = torch.sum(locations1 * mask1, dim=1)
         return (l_aux, capacity, num_experts, indices1_s.to(torch.int32).unsqueeze(0),
                 locations1_s.to(torch.int32).unsqueeze(0), gates1_s.unsqueeze(0), exp_counts)
 
     # Store the capacity location for each token
     locations1_s = torch.sum(locations1 * mask1, dim=1)
 
-    # Normalize gate probabilities
+    # Keep the assigned expert's softmax probability. route_norm rescales a kept row to 1.
     mask1_float = mask1.float()
     gates = gates * mask1_float
+    if _route_norm_enabled(route_norm, False):
+        denom_s = torch.sum(gates, dim=1, keepdim=True)
+        denom_s = torch.clamp(denom_s, min=torch.finfo(denom_s.dtype).eps)
+        gates = gates / denom_s
 
     locations1_sc = _one_hot_to_float(locations1_s, capacity)
     combine_weights = einsum("se,sc->sec", gates, locations1_sc)
@@ -344,8 +365,13 @@ def top2gating(logits: Tensor,
                drop_tokens: bool = True,
                ep_group: Union[torch.distributed.ProcessGroup, None] = None,
                top2_2nd_expert_sampling: bool = True,
-               sparse_routes: bool = False) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    """Implements Top2Gating on logits."""
+               sparse_routes: bool = False,
+               route_norm: Optional[bool] = None) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """Implements Top2Gating on logits.
+
+    route_norm=None keeps the historical weights: the two kept probabilities are rescaled to sum to 1.
+    The auxiliary loss is computed before that rescaling.
+    """
     # everything is in fp32 in this function
     gates = F.softmax(logits, dim=1)
 
@@ -399,16 +425,17 @@ def top2gating(logits: Tensor,
     locations1_s = torch.sum(locations1 * mask1, dim=1)
     locations2_s = torch.sum(locations2 * mask2, dim=1)
 
-    # Normalize gate probabilities
+    # Normalize gate probabilities. route_norm=None keeps this rescaling; False leaves the raw softmax values.
     mask1_float = mask1.float()
     mask2_float = mask2.float()
     gates1_s = einsum("se,se->s", gates, mask1_float)
     gates2_s = einsum("se,se->s", gates, mask2_float)
-    denom_s = gates1_s + gates2_s
-    # Avoid divide-by-zero
-    denom_s = torch.clamp(denom_s, min=torch.finfo(denom_s.dtype).eps)
-    gates1_s /= denom_s
-    gates2_s /= denom_s
+    if _route_norm_enabled(route_norm, True):
+        denom_s = gates1_s + gates2_s
+        # Avoid divide-by-zero. A dropped token is 0 / eps and stays 0.
+        denom_s = torch.clamp(denom_s, min=torch.finfo(denom_s.dtype).eps)
+        gates1_s /= denom_s
+        gates2_s /= denom_s
 
     if sparse_routes:
         # Routes evicted by the capacity limit are flagged with a negative expert index.
@@ -440,8 +467,13 @@ def topkgating(
     ep_group: Union[torch.distributed.ProcessGroup, None] = None,
     drop_policy: str = "probs",
     sparse_routes: bool = False,
+    route_norm: Optional[bool] = None,
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    """Implements TopKGating on logits."""
+    """Implements TopKGating on logits.
+
+    route_norm=None keeps the historical weights: the kept probabilities are rescaled to sum to 1.
+    The auxiliary loss is computed before that rescaling.
+    """
 
     # everything is in fp32 in this function
     # gating decisions
@@ -494,11 +526,12 @@ def topkgating(
         capacity = new_capacity
         locations = torch.cumsum(mask, dim=0) - 1
 
-    # normalize gates
+    # normalize gates. route_norm=None keeps this rescaling; False leaves the raw softmax values.
     gates_masked = gates * mask
-    gates_s = torch.sum(gates_masked, dim=-1, keepdim=True)
-    denom_s = torch.clamp(gates_s, min=torch.finfo(gates_masked.dtype).eps)
-    gates_masked = gates_masked / denom_s
+    if _route_norm_enabled(route_norm, True):
+        gates_s = torch.sum(gates_masked, dim=-1, keepdim=True)
+        denom_s = torch.clamp(gates_s, min=torch.finfo(gates_masked.dtype).eps)
+        gates_masked = gates_masked / denom_s
 
     if locations is None:
         raise ValueError(f"Locations is not set: {locations}")
@@ -539,6 +572,10 @@ class TopKGate(Module):
             size of model embedding dimension
         num_experts (int):
             number of experts in model
+        route_norm (bool, optional):
+            None keeps the historical combine weights (no renorm for k=1, renorm for k>=2).
+            True renormalizes the kept weights for every k. False leaves the raw softmax values.
+            The auxiliary loss does not use this flag.
     """
 
     wg: torch.nn.Linear
@@ -554,7 +591,8 @@ def __init__(self,
                  drop_tokens: bool = True,
                  use_rts: bool = True,
                  ep_group: Union[torch.distributed.ProcessGroup, None] = None,
-                 top2_2nd_expert_sampling: bool = True) -> None:
+                 top2_2nd_expert_sampling: bool = True,
+                 route_norm: Optional[bool] = None) -> None:
         super().__init__()
 
         self.wg = torch.nn.Linear(model_dim, num_experts, bias=False)
@@ -570,6 +608,8 @@ def __init__(self,
         self.drop_tokens = drop_tokens
         self.use_rts = use_rts
         self.top2_2nd_expert_sampling = top2_2nd_expert_sampling
+        # None is forwarded so each gating function keeps its own historical default.
+        self.route_norm = route_norm
 
     def _set_ep_group(self, ep_group):
         assert self.ep_group is None, 'Attempting to override an existing ep_group'
@@ -591,22 +631,36 @@ def forward(self,
         logits = torch.nn.functional.linear(input_fp32, weight=self.wg.weight.float(), bias=None)
 
         if self.k == 1:
-            gate_output = top1gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor,
-                                     self.min_capacity, used_token, self.noisy_gate_policy if self.training else None,
-                                     self.drop_tokens, self.use_rts, self.ep_group, sparse_routes, use_tutel)
+            gate_output = top1gating(logits,
+                                     self.capacity_factor if self.training else self.eval_capacity_factor,
+                                     self.min_capacity,
+                                     used_token,
+                                     self.noisy_gate_policy if self.training else None,
+                                     self.drop_tokens,
+                                     self.use_rts,
+                                     self.ep_group,
+                                     sparse_routes,
+                                     use_tutel,
+                                     route_norm=self.route_norm)
 
         elif self.k == 2:
-            gate_output = top2gating(logits, self.capacity_factor if self.training else self.eval_capacity_factor,
-                                     self.min_capacity, self.drop_tokens, self.ep_group, self.top2_2nd_expert_sampling,
-                                     sparse_routes)
+            gate_output = top2gating(logits,
+                                     self.capacity_factor if self.training else self.eval_capacity_factor,
+                                     self.min_capacity,
+                                     self.drop_tokens,
+                                     self.ep_group,
+                                     self.top2_2nd_expert_sampling,
+                                     sparse_routes,
+                                     route_norm=self.route_norm)
         else:
             gate_output = topkgating(logits,
                                      self.k,
                                      self.capacity_factor if self.training else self.eval_capacity_factor,
                                      self.min_capacity,
                                      self.drop_tokens,
                                      self.ep_group,
-                                     sparse_routes=sparse_routes)
+                                     sparse_routes=sparse_routes,
+                                     route_norm=self.route_norm)
 
         if self.wall_clock_breakdown:
             self.timers(TOPK_GATE_TIMER).stop()

docs/_tutorials/mixture-of-experts.md

@@ -63,6 +63,8 @@ Updated with MoE Layers
     self.fc4 = nn.Linear(84, 10)
 ```
 
+`MoE(..., route_norm=...)` is optional and defaults to `None`. That default keeps the historical combine weights: `k=1` uses the chosen expert's softmax probability, and `k>=2` rescales the kept weights so they sum to 1. `True` rescales for every `k`, and `False` leaves the raw softmax probabilities for every `k`. Either explicit value changes the layer output relative to a checkpoint trained with the other setting. The auxiliary load-balancing loss is unchanged. This constructor argument is the legacy gate. It is not the AutoEP config key `expert_parallel.route_norm`.
+
 ### Pyramid-Residual MoE
 
 Recently, we proposed a novel [Pyramid-Residual MoE](https://arxiv.org/abs/2201.05596) (PR-MoE) model architecture. To create such an MoE model, the users need to do two additional things: