Fix bugs in FP8 MoE support (#4312)

- B5: Add LoRA rank divisibility validation before integer division
- B15: Add missing bias after block-quantized FP8 linear path
- B21/B22: Fix _patch_fp8_moe_experts to use forward_moe_backend_fp8 for
  both grouped_mm and batched_mm slots (prevents F.linear(bf16, fp8) crash
  during decode)
- B4: Pass load_in_fp8, fast_inference, token, trust_remote_code to
  get_model_name in PEFT path
- B16: Gate maybe_patch_stacked_moe_expert_fp8_scales behind FP8 check
- B3: Narrow _has_prequantized_fp8_config to check FP8-specific indicators
  within compressed-tensors config
- Pass trust_remote_code through _has_prequantized_fp8_config call sites

unsloth/kernels/fp8.py

@@ -441,7 +441,7 @@ class FbgemmFp8Linear_matmul(torch.autograd.Function):
         elif (
             weight.shape[0] != weight_scale.shape[0]
             and weight.shape[1] == weight_scale.shape[0]
-        ) or (weight.shape[0] // 8 != 0 or weight.shape[1] // 8 != 0):
+        ) or (weight.shape[0] % 8 != 0 or weight.shape[1] % 8 != 0):
             # Either the weight/scale is transposed or its shape is not divisible by 8. Both cases, dequantizing is the preferred way.
             # The transpose case is generally noticed in backward pass when we do dY@W instead of @W.T as we do for forward.
             # The shape case, I noticed to happen in MLP of Qwen 2.5 VL 7B where the gate proj is of shape (3420, 1280) and 3420/8=427.5
@@ -605,6 +605,12 @@ except:
     pass
 
 
+HAS_FBGEMM_FP8_OPS = (
+    hasattr(torch.ops, "fbgemm")
+    and hasattr(torch.ops.fbgemm, "quantize_fp8_per_row")
+)
+
+
 @torch_compile
 def fp8_linear(X, weight, weight_scale, bias = None):
     # Per-tensor quantization: single scalar scale for entire weight
@@ -613,9 +619,18 @@ def fp8_linear(X, weight, weight_scale, bias = None):
         weight_scale.ndim == 2 and weight_scale.shape[1] > 1
     ):
         out = fp8_block_quant_linear(X, weight, weight_scale)
+        if bias is not None:
+            out = out + bias
     # Row/channel quantized FP8: 2D scale with shape (n, 1)
-    else:
+    elif HAS_FBGEMM_FP8_OPS:
         out = fbgemm_fp8_linear(X, weight, weight_scale, bias)
+    else:
+        # Fallback: dequantize FP8 weight and use standard matmul
+        W_deq = weight_dequant(weight, weight_scale).T
+        out = torch_matmul(X, W_deq)
+        if bias is not None:
+            out = out + bias
+        del W_deq
     return out
 
 
@@ -694,6 +709,10 @@ def compressed_linear_forward_patch(self, input):
 
 def _fp8_moe_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts):
     total_rank = weight_A.shape[0]
+    if num_experts == 0 or total_rank % num_experts != 0:
+        raise ValueError(
+            f"LoRA total_rank ({total_rank}) must be divisible by num_experts ({num_experts})"
+        )
     rank_per_expert = total_rank // num_experts
 
     dim_A = weight_A.shape[1]
@@ -755,19 +774,17 @@ def _fp8_moe_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts):
 def _patch_fp8_moe_experts():
     try:
         from transformers.integrations import finegrained_fp8
-        from unsloth_zoo.temporary_patches.moe_utils import (
-            forward_moe_backend,
-            forward_native_moe_loop,
+        from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+            forward_moe_backend_fp8,
         )
     except Exception:
         return
 
     experts_interface = getattr(finegrained_fp8, "ALL_FP8_EXPERTS_FUNCTIONS", None)
-    if experts_interface is None:
-        return
+    if experts_interface is not None:
+        experts_interface["grouped_mm"] = forward_moe_backend_fp8
+        experts_interface["batched_mm"] = forward_moe_backend_fp8
 
-    experts_interface["grouped_mm"] = forward_moe_backend
-    experts_interface["batched_mm"] = forward_native_moe_loop
     if hasattr(finegrained_fp8, "FP8Experts"):
         finegrained_fp8.FP8Experts._unsloth_lora_extractor_fn = staticmethod(
             _fp8_moe_lora_extractor

unsloth/models/loader.py

@@ -819,12 +819,13 @@ class FastLanguageModel(FastLlamaModel):
 
         if load_in_fp8 != False:
             _tag_model_with_fp8_torchao_config(model, fp8_mode)
-        maybe_patch_stacked_moe_expert_fp8_scales(
-            model,
-            model_name = model_name,
-            token = token,
-            revision = revision if not is_peft else None,
-        )
+        if load_in_fp8 != False or _has_prequantized_fp8_config(model_name, token=token, trust_remote_code=trust_remote_code):
+            maybe_patch_stacked_moe_expert_fp8_scales(
+                model,
+                model_name = model_name,
+                token = token,
+                revision = revision if not is_peft else None,
+            )
 
         if is_peft:
             # From https://github.com/huggingface/peft/issues/184
@@ -1359,7 +1360,14 @@ class FastModel(FastBaseModel):
             # Check base model again for PEFT
             model_name = peft_config.base_model_name_or_path
             if not use_exact_model_name:
-                model_name = get_model_name(model_name, load_in_4bit)
+                model_name = get_model_name(
+                    model_name,
+                    load_in_4bit,
+                    load_in_fp8=load_in_fp8,
+                    fast_inference=fast_inference,
+                    token=token,
+                    trust_remote_code=trust_remote_code,
+                )
             # Check if pre-quantized models are allowed
             # AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 64)
             if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
@@ -1568,12 +1576,13 @@ class FastModel(FastBaseModel):
 
         if load_in_fp8 != False:
             _tag_model_with_fp8_torchao_config(model, fp8_mode)
-        maybe_patch_stacked_moe_expert_fp8_scales(
-            model,
-            model_name = model_name,
-            token = token,
-            revision = revision if not is_peft else None,
-        )
+        if load_in_fp8 != False or _has_prequantized_fp8_config(model_name, token=token, trust_remote_code=trust_remote_code):
+            maybe_patch_stacked_moe_expert_fp8_scales(
+                model,
+                model_name = model_name,
+                token = token,
+                revision = revision if not is_peft else None,
+            )
 
         if is_peft:
             # From https://github.com/huggingface/peft/issues/184

unsloth/models/loader_utils.py

@@ -302,7 +302,25 @@ def _has_prequantized_fp8_config(
         return False
 
     quant_method = quantization_config.get("quant_method", None)
-    return quant_method in ("compressed-tensors", "fbgemm_fp8", "fp8")
+    if quant_method in ("fbgemm_fp8", "fp8"):
+        return True
+    if quant_method == "compressed-tensors":
+        # Check for FP8-specific config within compressed-tensors
+        config_groups = quantization_config.get("config_groups", {})
+        for group in config_groups.values():
+            if isinstance(group, dict):
+                weights = group.get("weights", {})
+                if isinstance(weights, dict):
+                    wtype = weights.get("type", "")
+                    num_bits = weights.get("num_bits", 0)
+                    if wtype == "float" and num_bits == 8:
+                        return True
+        # Also check top-level quantization type
+        quant_type = quantization_config.get("quantization_type", "")
+        if "fp8" in str(quant_type).lower() or "float8" in str(quant_type).lower():
+            return True
+        return False
+    return False
 
 
 def _offline_quantize_to_fp8(model_name: str, fp8_mode: str) -> str: