diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index a1d88444e6..e3fd1b9cc4 100644 --- a/unsloth/kernels/fp8.py +++ b/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 diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 7d978cf184..43092bc71d 100644 --- a/unsloth/models/loader.py +++ b/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 diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index c5c417b0b2..fa970372ee 100644 --- a/unsloth/models/loader_utils.py +++ b/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: