logs

src/axolotl/kernels/moe/torch_grouped.py

@@ -88,25 +88,21 @@ def _call_grouped_mm(
 def moe_ffn_forward_grouped(
     hidden_states, gate_linear, experts_module, top_k: int
 ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
-    """
-    Attempt a grouped GEMM fast path using PyTorch 2.8+.
-    If unavailable or fails, returns (None, None) so caller can fallback.
-    """
-    try:
+    """Attempt grouped GEMM fast path using PyTorch 2.8+."""
+    global LAST_ERROR
+    LAST_ERROR = None
     bsz, seqlen, hdim = hidden_states.shape
     x = hidden_states.view(-1, hdim)
     router_logits = gate_linear(x)
 
-        # topk routing in torch (keep simple to avoid dependency cycles)
+    # top-k routing executed in torch to avoid extra dependencies
     routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
     topk_weight, topk_idx = torch.topk(routing_weights, top_k, dim=-1, sorted=False)
     topk_weight = (topk_weight / topk_weight.sum(dim=-1, keepdim=True)).to(x.dtype)
 
-        # Build per-expert input lists
     flat_idx = topk_idx.view(-1)
     x_rep = x.repeat_interleave(top_k, dim=0)
 
-        # Cache stacked weights on experts (support Mixtral and Qwen-style layouts)
     E = experts_module.num_experts
     dev, dt = x.device, x.dtype
     first = experts_module[0]
@@ -133,6 +129,7 @@ def moe_ffn_forward_grouped(
                 "torch_grouped: unsupported expert layout; falling back to naive"
             )
             experts_module._ax_grouped_logged_fail = True
+        LAST_ERROR = "unsupported expert layout"
         return None, None
 
     def _resolve_expert(idx: int):
@@ -141,9 +138,7 @@ def moe_ffn_forward_grouped(
             return expert
         nested_mod = getattr(expert, nested_attr, None)
         if nested_mod is None:
-                raise AttributeError(
-                    f"expert {idx} missing nested module '{nested_attr}'"
-                )
+            raise AttributeError(f"expert {idx} missing nested module '{nested_attr}'")
         return nested_mod
 
     try:
@@ -178,7 +173,6 @@ def moe_ffn_forward_grouped(
             W13 = experts_module._stacked_w13
             W2 = experts_module._stacked_w2
         else:
-                # Qwen-style MoE: either gate_up_proj (2I x H) or (up_proj + gate_proj), down_proj (H x I)
             if (
                 not hasattr(experts_module, "_stacked_w13")
                 or experts_module._stacked_w13.device != dev
@@ -188,11 +182,9 @@ def moe_ffn_forward_grouped(
                 w2 = []
                 for i in range(E):
                     mod = _resolve_expert(i)
-                        # prefer fused gate_up_proj if present
                     if hasattr(mod, "gate_up_proj"):
                         w13.append(mod.gate_up_proj.weight.t())
                     elif hasattr(mod, "up_proj") and hasattr(mod, "gate_proj"):
-                            # concatenate [up | gate] along N
                         w13.append(
                             torch.cat(
                                 [mod.up_proj.weight.t(), mod.gate_proj.weight.t()],
@@ -231,10 +223,9 @@ def moe_ffn_forward_grouped(
             experts_module._ax_grouped_logged_fail = True
         return None, None
 
-        # Grouped GEMM for up+gate
     As: List[torch.Tensor] = []
     Bs: List[torch.Tensor] = []
-        expert_slices = []
+    expert_slices: List[Tuple[int, torch.Tensor]] = []
     for i in range(E):
         sel = flat_idx == i
         if sel.any():
@@ -244,7 +235,6 @@ def moe_ffn_forward_grouped(
             expert_slices.append((i, sel))
 
     if not As:
-            # no tokens routed — edge case
         out = torch.zeros_like(x)
         return out.view(bsz, seqlen, hdim), router_logits
 
@@ -257,13 +247,10 @@ def moe_ffn_forward_grouped(
             experts_module._ax_grouped_logged_fail = True
         return None, None
 
-        # SwiGLU on each expert block and prepare for down projection
     As2: List[torch.Tensor] = []
     Bs2: List[torch.Tensor] = []
     y_buf = torch.empty_like(x_rep)
-
-        # split Y into (I, I)
-        for Yi in Y_list:
+    for (i, _sel), Yi in zip(expert_slices, Y_list, strict=False):
         I2 = Yi.shape[-1] // 2
         Yi_hidden = F.silu(Yi[:, :I2]) * Yi[:, I2:]
         As2.append(Yi_hidden)
@@ -278,8 +265,7 @@ def moe_ffn_forward_grouped(
             experts_module._ax_grouped_logged_fail = True
         return None, None
 
-        # Write back, apply per-token weighting, and reduce over top_k
-        for (_, sel), Out_i in zip(expert_slices, Y2_list, strict=False):
+    for (_i, sel), Out_i in zip(expert_slices, Y2_list, strict=False):
         y_buf[sel] = Out_i
     y = (y_buf.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
     if not getattr(experts_module, "_ax_grouped_logged_ok", False):
@@ -288,5 +274,3 @@ def moe_ffn_forward_grouped(
         )
         experts_module._ax_grouped_logged_ok = True
     return y.view(bsz, seqlen, hdim), router_logits
-    except Exception:
-        return None, None

src/axolotl/monkeypatch/moe_grouped.py

@@ -82,8 +82,17 @@ def apply_grouped_to_moe_blocks(cfg=None) -> None:
             # One-time log per block instance indicating whether grouped engaged or fallback occurred
             if not getattr(self, "_ax_grouped_wrapper_logged", False):
                 if y is None:
+                    reason = getattr(_tg, "LAST_ERROR", None)
+                    if reason:
                         _LOG.warning(
-                        f"Grouped wrapper active but fell back to naive for {self.__class__.__name__}"
+                            "Grouped wrapper fell back to naive for %s (reason=%s)",
+                            self.__class__.__name__,
+                            reason,
+                        )
+                    else:
+                        _LOG.warning(
+                            "Grouped wrapper active but fell back to naive for %s",
+                            self.__class__.__name__,
                         )
                 else:
                     _LOG.info(