simplify

src/axolotl/kernels/moe/torch_grouped.py

@@ -23,7 +23,7 @@ def available() -> bool:
         major, minor = torch.cuda.get_device_capability()
         if major < 9:
             return False
-        return hasattr(torch.ops, "aten") and hasattr(torch.ops.aten, "_grouped_mm")
+        return hasattr(torch.ops, "_grouped_mm")
     except Exception:
         return False
 
@@ -57,7 +57,7 @@ def _call_grouped_mm(
     As: List[torch.Tensor], Bs: List[torch.Tensor]
 ) -> Optional[List[torch.Tensor]]:
     """
-    Call grouped mm using aten._grouped_mm with packed representation.
+    Call grouped mm with packed representation
 
     - A_cat: concat As along rows -> [sum_i Mi, K]
     - B_stk: stack Bs per group -> [G, K, N]
@@ -65,35 +65,23 @@ def _call_grouped_mm(
     Returns list of per-group outputs split from concatenated result.
     """
     global LAST_ERROR
-    try:
+    # Ensure 2D contiguous inputs
-        # Ensure 2D contiguous inputs
+    As2 = [a.contiguous().view(a.shape[0], a.shape[1]) for a in As]
-        As2 = [a.contiguous().view(a.shape[0], a.shape[1]) for a in As]
+    Bs2 = [b.contiguous().view(b.shape[0], b.shape[1]) for b in Bs]
-        Bs2 = [b.contiguous().view(b.shape[0], b.shape[1]) for b in Bs]
 
-        if not As2:
+    if not As2:
-            return []
+        return []
-        device = As2[0].device
+    device = As2[0].device
-        A_cat = torch.cat(As2, dim=0)
+    A_cat = torch.cat(As2, dim=0).to(torch.bfloat16)
-        B_stk = torch.stack(Bs2, dim=0)
+    B_stk = torch.stack(Bs2, dim=0).to(torch.bfloat16)
-        offs = torch.tensor([a.shape[0] for a in As2], device=device, dtype=torch.int32)
+    offs = torch.tensor([a.shape[0] for a in As2], device=device, dtype=torch.int32)
-
+    Y_cat = torch._grouped_mm(A_cat, B_stk, offs)  # type: ignore[attr-defined]
-        if hasattr(torch.ops.aten, "_grouped_mm"):
+    outs: List[torch.Tensor] = []
-            try:
+    start = 0
-                Y_cat = torch.ops.aten._grouped_mm(A_cat, B_stk, offs)  # type: ignore[attr-defined]
+    for m in offs.tolist():
-                outs: List[torch.Tensor] = []
+        outs.append(Y_cat[start : start + m, :])
-                start = 0
+        start += m
-                for m in offs.tolist():
+    return outs
-                    outs.append(Y_cat[start : start + m, :])
-                    start += m
-                return outs
-            except Exception as e:
-                LAST_ERROR = f"_grouped_mm failed: {e}"
-                return None
-        LAST_ERROR = "aten._grouped_mm not present"
-        return None
-    except Exception as e:
-        LAST_ERROR = f"call error: {e}"
-        return None
 
 
 def moe_ffn_forward_grouped(
@@ -103,6 +91,9 @@ def moe_ffn_forward_grouped(
     global LAST_ERROR
     LAST_ERROR = None
     bsz, seqlen, hdim = hidden_states.shape
+    compute_dtype = gate_linear.weight.dtype
+    if hidden_states.dtype != compute_dtype:
+        hidden_states = hidden_states.to(dtype=compute_dtype)
     x = hidden_states.view(-1, hdim)
     router_logits = gate_linear(x)
 
@@ -114,16 +105,7 @@ def moe_ffn_forward_grouped(
     flat_idx = topk_idx.view(-1)
     x_rep = x.repeat_interleave(top_k, dim=0)
 
-    try:
+    E = _num_experts(experts_module)
-        E = _num_experts(experts_module)
-    except AttributeError as err:
-        LAST_ERROR = str(err)
-        if not getattr(experts_module, "_ax_grouped_logged_fail", False):
-            _LOGGER.warning(
-                "torch_grouped: could not determine expert count; falling back to naive"
-            )
-            experts_module._ax_grouped_logged_fail = True
-        return None, None
     dev, dt = x.device, x.dtype
     first = experts_module[0]
 
@@ -161,87 +143,76 @@ def moe_ffn_forward_grouped(
             raise AttributeError(f"expert {idx} missing nested module '{nested_attr}'")
         return nested_mod
 
-    try:
+    if is_mixtral:
-        if is_mixtral:
+        if (
-            if (
+            not hasattr(experts_module, "_stacked_w1")
-                not hasattr(experts_module, "_stacked_w1")
+            or experts_module._stacked_w1.device != dev
-                or experts_module._stacked_w1.device != dev
+            or experts_module._stacked_w1.dtype != dt
-                or experts_module._stacked_w1.dtype != dt
+        ):
-            ):
+            mods = [_resolve_expert(i) for i in range(E)]
-                mods = [_resolve_expert(i) for i in range(E)]
+            w1 = [mod.w1.weight.t() for mod in mods]
-                w1 = [mod.w1.weight.t() for mod in mods]
+            w3 = [mod.w3.weight.t() for mod in mods]
-                w3 = [mod.w3.weight.t() for mod in mods]
+            w2 = [mod.w2.weight.t() for mod in mods]
-                w2 = [mod.w2.weight.t() for mod in mods]
+            experts_module._stacked_w1 = (
-                experts_module._stacked_w1 = (
+                torch.stack(w1, dim=0)
-                    torch.stack(w1, dim=0)
+                .to(device=dev, dtype=dt, non_blocking=True)
-                    .to(device=dev, dtype=dt, non_blocking=True)
+                .contiguous()
-                    .contiguous()
-                )
-                experts_module._stacked_w3 = (
-                    torch.stack(w3, dim=0)
-                    .to(device=dev, dtype=dt, non_blocking=True)
-                    .contiguous()
-                )
-                experts_module._stacked_w2 = (
-                    torch.stack(w2, dim=0)
-                    .to(device=dev, dtype=dt, non_blocking=True)
-                    .contiguous()
-                )
-                experts_module._stacked_w13 = torch.cat(
-                    [experts_module._stacked_w1, experts_module._stacked_w3], dim=-1
-                ).contiguous()
-            W13 = experts_module._stacked_w13
-            W2 = experts_module._stacked_w2
-        else:
-            if (
-                not hasattr(experts_module, "_stacked_w13")
-                or experts_module._stacked_w13.device != dev
-                or experts_module._stacked_w13.dtype != dt
-            ):
-                w13 = []
-                w2 = []
-                for i in range(E):
-                    mod = _resolve_expert(i)
-                    if hasattr(mod, "gate_up_proj"):
-                        w13.append(mod.gate_up_proj.weight.t())
-                    elif hasattr(mod, "up_proj") and hasattr(mod, "gate_proj"):
-                        w13.append(
-                            torch.cat(
-                                [mod.up_proj.weight.t(), mod.gate_proj.weight.t()],
-                                dim=-1,
-                            )
-                        )
-                    else:
-                        LAST_ERROR = "unrecognized Qwen MoE expert weight layout"
-                        if not getattr(
-                            experts_module, "_ax_grouped_logged_fail", False
-                        ):
-                            _LOGGER.warning(
-                                "torch_grouped: could not resolve Qwen MoE expert weights; fallback to naive"
-                            )
-                            experts_module._ax_grouped_logged_fail = True
-                        return None, None
-                    w2.append(mod.down_proj.weight.t())
-                experts_module._stacked_w13 = (
-                    torch.stack(w13, dim=0)
-                    .to(device=dev, dtype=dt, non_blocking=True)
-                    .contiguous()
-                )
-                experts_module._stacked_w2 = (
-                    torch.stack(w2, dim=0)
-                    .to(device=dev, dtype=dt, non_blocking=True)
-                    .contiguous()
-                )
-            W13 = experts_module._stacked_w13
-            W2 = experts_module._stacked_w2
-    except AttributeError as err:
-        LAST_ERROR = str(err)
-        if not getattr(experts_module, "_ax_grouped_logged_fail", False):
-            _LOGGER.warning(
-                "torch_grouped: expert weights missing expected attributes; falling back to naive"
             )
-            experts_module._ax_grouped_logged_fail = True
+            experts_module._stacked_w3 = (
-        return None, None
+                torch.stack(w3, dim=0)
+                .to(device=dev, dtype=dt, non_blocking=True)
+                .contiguous()
+            )
+            experts_module._stacked_w2 = (
+                torch.stack(w2, dim=0)
+                .to(device=dev, dtype=dt, non_blocking=True)
+                .contiguous()
+            )
+            experts_module._stacked_w13 = torch.cat(
+                [experts_module._stacked_w1, experts_module._stacked_w3], dim=-1
+            ).contiguous()
+        W13 = experts_module._stacked_w13
+        W2 = experts_module._stacked_w2
+    else:
+        if (
+            not hasattr(experts_module, "_stacked_w13")
+            or experts_module._stacked_w13.device != dev
+            or experts_module._stacked_w13.dtype != dt
+        ):
+            w13 = []
+            w2 = []
+            for i in range(E):
+                mod = _resolve_expert(i)
+                if hasattr(mod, "gate_up_proj"):
+                    w13.append(mod.gate_up_proj.weight.t())
+                elif hasattr(mod, "up_proj") and hasattr(mod, "gate_proj"):
+                    w13.append(
+                        torch.cat(
+                            [mod.up_proj.weight.t(), mod.gate_proj.weight.t()],
+                            dim=-1,
+                        )
+                    )
+                else:
+                    LAST_ERROR = "unrecognized Qwen MoE expert weight layout"
+                    if not getattr(experts_module, "_ax_grouped_logged_fail", False):
+                        _LOGGER.warning(
+                            "torch_grouped: could not resolve Qwen MoE expert weights; fallback to naive"
+                        )
+                        experts_module._ax_grouped_logged_fail = True
+                    return None, None
+                w2.append(mod.down_proj.weight.t())
+            experts_module._stacked_w13 = (
+                torch.stack(w13, dim=0)
+                .to(device=dev, dtype=dt, non_blocking=True)
+                .contiguous()
+            )
+            experts_module._stacked_w2 = (
+                torch.stack(w2, dim=0)
+                .to(device=dev, dtype=dt, non_blocking=True)
+                .contiguous()
+            )
+        W13 = experts_module._stacked_w13
+        W2 = experts_module._stacked_w2
 
     As: List[torch.Tensor] = []
     Bs: List[torch.Tensor] = []
@@ -251,7 +222,7 @@ def moe_ffn_forward_grouped(
         if sel.any():
             Xi = x_rep[sel].contiguous()
             As.append(Xi)
-            Bs.append(W13[i].contiguous())
+            Bs.append(W13[i].reshape(hdim, -1).contiguous())
             expert_slices.append((i, sel))
 
     if not As:
@@ -264,19 +235,13 @@ def moe_ffn_forward_grouped(
         target_dtype: torch.dtype,
     ) -> Optional[List[torch.Tensor]]:
         global LAST_ERROR
-        try:
+        if target_dtype != dt:
-            if target_dtype != dt:
+            a_tensors = [t.to(target_dtype).contiguous() for t in a_tensors]
-                a_tensors = [t.to(target_dtype).contiguous() for t in a_tensors]
+            b_tensors = [t.to(target_dtype).contiguous() for t in b_tensors]
-                b_tensors = [t.to(target_dtype).contiguous() for t in b_tensors]
+        outputs = _call_grouped_mm(a_tensors, b_tensors)
-            outputs = _call_grouped_mm(a_tensors, b_tensors)
+        if outputs is not None and target_dtype != dt:
-            if outputs is not None and target_dtype != dt:
+            outputs = [t.to(dt).contiguous() for t in outputs]
-                outputs = [t.to(dt).contiguous() for t in outputs]
+        return outputs
-            return outputs
-        except RuntimeError as err:
-            LAST_ERROR = f"grouped_mm cast failure: {err}"  # type: ignore[assignment]
-            if torch.cuda.is_available():  # pragma: no cover - defensive
-                torch.cuda.synchronize()
-            return None
 
     def _try_grouped_mm(
         a_tensors: List[torch.Tensor], b_tensors: List[torch.Tensor]
@@ -312,7 +277,7 @@ def moe_ffn_forward_grouped(
         I2 = Yi.shape[-1] // 2
         Yi_hidden = F.silu(Yi[:, :I2]) * Yi[:, I2:]
         As2.append(Yi_hidden)
-        Bs2.append(W2[i].contiguous())
+        Bs2.append(W2[i].reshape(I2, hdim).contiguous())
 
     Y2_list, _cast_used_down = _try_grouped_mm(As2, Bs2)
     if Y2_list is None: