BlockSparseMLP: Add DS3/Dots routing

Author: turboderp

exllamav3/modules/block_sparse_mlp.py

@@ -23,6 +23,11 @@ class RoutingCFG:
     router_logits_bsz1: torch.Tensor
     routing_weights_bsz1: torch.Tensor
     selected_experts_bsz1: torch.Tensor
+    e_score_correction_bias: torch.Tensor | None
+    routed_scaling_factor: float | None
+    n_group: int | None
+    topk_group: int | None
+
 
 def routing(bsz, cfg, y, params):
     activate_all_experts = params.get("activate_all_experts")
@@ -50,6 +55,75 @@ def routing(bsz, cfg, y, params):
         return selected_experts, routing_weights
 
 
+# TODO: Optimize (for DS3)
+def routing_ds3(bsz, cfg, y, params):
+    activate_all_experts = params.get("activate_all_experts")
+    router_logits = torch.matmul(y, cfg.gate_tensor)
+
+    scores = router_logits.sigmoid()
+    scores_for_choice = scores.view(-1, cfg.num_experts) + cfg.e_score_correction_bias.unsqueeze(0)
+    group_scores = (
+        scores_for_choice.view(-1, cfg.n_group, cfg.num_experts // cfg.n_group)
+        .topk(2, dim = -1)[0]
+        .sum(dim = -1)
+    )
+    group_idx = torch.topk(group_scores, k = cfg.topk_group, dim = -1, sorted = False)[1]
+    group_mask = torch.zeros_like(group_scores)
+    group_mask.scatter_(1, group_idx, 1)
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(-1, cfg.n_group, cfg.num_experts // cfg.n_group)
+        .reshape(-1, cfg.num_experts)
+    )
+    scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)
+
+    topk_indices = torch.topk(
+        scores_for_choice,
+        k = cfg.num_experts if activate_all_experts else cfg.num_experts_per_tok,
+        dim = -1,
+        sorted = False
+    )[1]
+    topk_weights = scores.gather(1, topk_indices)
+    denominator = topk_weights.sum(dim = -1, keepdim = True) + 1e-20
+    topk_weights /= denominator
+    topk_weights = topk_weights * cfg.routed_scaling_factor
+    return topk_indices, topk_weights
+
+
+def routing_dots(bsz, cfg, y, params):
+    activate_all_experts = params.get("activate_all_experts")
+
+    if bsz == 1 and not activate_all_experts:
+        torch.matmul(y, cfg.gate_tensor, out = cfg.router_logits_bsz1)
+        cfg.router_logits_bsz1 += cfg.e_score_correction_bias
+        torch.topk(
+            cfg.router_logits_bsz1,
+            cfg.num_experts_per_tok,
+            dim = -1,
+            out = (cfg.routing_weights_bsz1, cfg.selected_experts_bsz1),
+            sorted = False
+        )
+        # TODO: Custom kernel for sigmoid normalization
+        cfg.routing_weights_bsz1.sigmoid_()
+        factor = cfg.routed_scaling_factor / (cfg.routing_weights_bsz1.sum(dim = -1, keepdim = True) + 1e-20)
+        cfg.routing_weights_bsz1 *= factor
+        return cfg.selected_experts_bsz1, cfg.routing_weights_bsz1
+
+    else:
+        router_logits = torch.matmul(y, cfg.gate_tensor)
+        router_logits += cfg.e_score_correction_bias
+        routing_weights, selected_experts = torch.topk(
+            router_logits,
+            cfg.num_experts if activate_all_experts else cfg.num_experts_per_tok,
+            dim = -1
+        )
+        # TODO: Custom kernel for sigmoid normalization
+        routing_weights.sigmoid_()
+        factor = cfg.routed_scaling_factor / (routing_weights.sum(dim = -1, keepdim = True) + 1e-20)
+        routing_weights *= factor
+        return selected_experts, routing_weights
+
+
 @dataclass
 class ExpertsCFG:
     yh: torch.Tensor
@@ -77,6 +151,10 @@ def __init__(
         out_dtype: torch.dtype = None,
         activation_fn: str = "silu",
         interm_dtype: torch.dtype = None,
+        deepseekv3_routing: bool = False,
+        routed_scaling_factor: float | None = None,
+        n_group: int | None = None,
+        topk_group: int | None = None,
         shared_experts: MLP | GatedMLP | None = None
     ):
         super().__init__(config, key, None)
@@ -89,6 +167,11 @@ def __init__(
         self.num_experts_per_tok = num_experts_per_tok
         self.hidden_size = hidden_size
 
+        self.deepseekv3_routing = deepseekv3_routing
+        self.routed_scaling_factor = routed_scaling_factor
+        self.n_group = n_group
+        self.topk_group = topk_group
+
         self.routing_gate = Linear(
             config = config,
             key = f"{key}.{key_routing_gate}",
@@ -152,6 +235,8 @@ def __init__(
         self.routing_cfg = None
         self.experts_cfg = None
 
+        self.e_score_correction_bias = None
+
         self.shared_experts = shared_experts
         if shared_experts is not None:
             self.register_submodule(shared_experts)
@@ -161,6 +246,9 @@ def __init__(
     def load(self, device: torch.Device, **kwargs):
         super().load(device, **kwargs)
 
+        self.e_score_correction_bias = \
+            self.config.stc.get_tensor(self.key + ".gate.e_score_correction_bias", self.device, optional = True)
+
         # Test if experts can be fused
         num_exl3_tensors = 0
         num_nonexl3_tensors = 0
@@ -189,7 +277,11 @@ def load(self, device: torch.Device, **kwargs):
             num_experts_per_tok = self.num_experts_per_tok,
             router_logits_bsz1 = router_logits_bsz1,
             routing_weights_bsz1 = routing_weights_bsz1,
-            selected_experts_bsz1 = selected_experts_bsz1
+            selected_experts_bsz1 = selected_experts_bsz1,
+            e_score_correction_bias = self.e_score_correction_bias,
+            routed_scaling_factor = self.routed_scaling_factor,
+            n_group = self.n_group,
+            topk_group = self.topk_group,
         )
 
         yh = torch.empty(
@@ -231,6 +323,7 @@ def unload(self):
             self.multi_down = None
         self.routing_cfg = None
         self.experts_cfg = None
+        self.e_score_correction_bias = None
         super().unload()
 
 
@@ -245,12 +338,18 @@ def forward(
         y = x.view(-1, self.hidden_size)
         bsz = y.shape[0]
 
-        # selected_experts, routing_weights = routing(bsz, self.routing_cfg, y, params)
-        selected_experts, routing_weights = ext.blocksparse_mlp_routing(bsz, self.routing_cfg, y, params)
+        if self.deepseekv3_routing:
+            if self.n_group == 1 and self.topk_group == 1:
+                selected_experts, routing_weights = routing_dots(bsz, self.routing_cfg, y, params)
+            # else:
+            #     selected_experts, routing_weights = routing_ds3(bsz, self.routing_cfg, y, params)
+        else:
+            # selected_experts, routing_weights = routing(bsz, self.routing_cfg, y, params)
+            selected_experts, routing_weights = ext.blocksparse_mlp_routing(bsz, self.routing_cfg, y, params, False)
 
         # Torch path
         if bsz > 1 or not self.is_quantized:
-            final_hidden_states = torch.zeros_like(y)
+            final_hidden_states = torch.zeros_like(y, dtype = self.out_dtype)
 
             expert_mask = torch.nn.functional.one_hot(
                 selected_experts,
@@ -338,8 +437,14 @@ def mlp(exp_i, xc):
             )
 
             final_hidden_states = cfg.out_d[:1, ...]
-            return final_hidden_states.view(x.shape)
             final_hidden_states = final_hidden_states.view(x.shape)
             if self.shared_experts:
                 final_hidden_states += self.shared_experts.forward(x, params)
-            return final_hidden_states
+            return final_hidden_states
+
+    @override
+    def get_tensors(self):
+        t = super().get_tensors()
+        if self.e_score_correction_bias is not None:
+            t[f"{self.key}.gate.e_score_correction_bias"] = self.e_score_correction_bias.contiguous()
+        return t