[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

unsloth/kernels/fp8.py

@@ -60,6 +60,15 @@ except:
         "Unsloth: Could not find torchao.prototype.blockwise_fp8_inference.blockwise_quantization.blockwise_fp8_gemm"
     )
 
+try:
+    from compressed_tensors.linear.compressed_linear import CompressedLinear
+    from compressed_tensors.quantization.quant_args import QuantizationStrategy
+    from compressed_tensors.quantization.quant_config import QuantizationStatus
+except:
+    CompressedLinear = None
+    QuantizationStrategy = None
+    QuantizationStatus = None
+
 
 @triton.jit
 def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
@@ -617,8 +626,159 @@ def module_forward_patch(forward_function, scale_attr = "weight_scale"):
     return patched_forward
 
 
+def _compressed_linear_supports_unsloth_fp8(self):
+    if (
+        CompressedLinear is None
+        or QuantizationStrategy is None
+        or QuantizationStatus is None
+    ):
+        return False
+
+    weight = getattr(self, "weight", None)
+    weight_scale = getattr(self, "weight_scale", None)
+    quantization_scheme = getattr(self, "quantization_scheme", None)
+    quantization_args = getattr(quantization_scheme, "weights", None)
+    if (
+        weight is None
+        or weight_scale is None
+        or quantization_args is None
+        or weight.dtype != torch.float8_e4m3fn
+    ):
+        return False
+
+    if getattr(quantization_args, "type", None) != "float":
+        return False
+    if getattr(quantization_args, "num_bits", None) != 8:
+        return False
+    if getattr(quantization_args, "symmetric", None) is not True:
+        return False
+    if getattr(self, "weight_zero_point", None) is not None:
+        return False
+
+    strategy = getattr(quantization_args, "strategy", None)
+    if strategy not in (
+        QuantizationStrategy.TENSOR,
+        QuantizationStrategy.CHANNEL,
+        QuantizationStrategy.BLOCK,
+        "tensor",
+        "channel",
+        "block",
+    ):
+        return False
+
+    if not torch.is_tensor(weight_scale):
+        return False
+    if weight_scale.numel() == 0:
+        return False
+
+    return True
+
+
+def _compressed_linear_forward_fallback(self, input):
+    if self.quantization_status == QuantizationStatus.COMPRESSED:
+        weight_data = self.compressor.decompress_module(self)
+        param = nn.Parameter(weight_data, requires_grad = False)
+        from compressed_tensors.utils import register_offload_parameter
+
+        register_offload_parameter(self, "weight", param)
+        self.quantization_status = QuantizationStatus.FROZEN
+
+    return F.linear(input, self.weight, self.bias)
+
+
+def compressed_linear_forward_patch(self, input):
+    if _compressed_linear_supports_unsloth_fp8(self):
+        return fp8_linear(input, self.weight, self.weight_scale, self.bias)
+    return _compressed_linear_forward_fallback(self, input)
+
+
+def _fp8_moe_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts):
+    total_rank = weight_A.shape[0]
+    rank_per_expert = total_rank // num_experts
+
+    dim_A = weight_A.shape[1]
+    dim_B = weight_B.shape[0]
+
+    hidden_dim = None
+    intermediate_dim = None
+    current = wrapper
+    while hasattr(current, "base_layer"):
+        current = current.base_layer
+        if hasattr(current, "hidden_dim"):
+            hidden_dim = current.hidden_dim
+        if hasattr(current, "intermediate_dim"):
+            intermediate_dim = current.intermediate_dim
+        if hasattr(current, "gate_up_proj") and hasattr(current.gate_up_proj, "shape"):
+            shape = current.gate_up_proj.shape
+            if len(shape) == 3:
+                hidden_dim = shape[2]
+                intermediate_dim = shape[1] // 2
+
+    param_name = getattr(wrapper, "parameter_name", None)
+
+    if (
+        param_name == "down_proj"
+        and intermediate_dim is not None
+        and hidden_dim is not None
+    ):
+        first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+        first_weight = first_weight.permute(1, 0, 2).contiguous()
+        second_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
+        return first_weight, second_weight, scaling, num_experts
+
+    elif param_name == "gate_up_proj" and hidden_dim is not None:
+        first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+        first_weight = first_weight.permute(1, 0, 2).contiguous()
+        second_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
+        return first_weight, second_weight, scaling, num_experts
+
+    if hidden_dim is not None:
+        if dim_B == hidden_dim:
+            first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+            first_weight = first_weight.permute(1, 0, 2).contiguous()
+            second_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
+            return first_weight, second_weight, scaling, num_experts
+        elif dim_A == hidden_dim:
+            first_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
+            first_weight = first_weight.permute(0, 2, 1).contiguous()
+            second_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+            second_weight = second_weight.permute(1, 2, 0).contiguous()
+            return first_weight, second_weight, scaling, num_experts
+
+    first_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
+    first_weight = first_weight.permute(0, 2, 1).contiguous()
+    second_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
+    second_weight = second_weight.permute(1, 2, 0).contiguous()
+    return first_weight, second_weight, 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,
+        )
+    except Exception:
+        return
+
+    experts_interface = getattr(finegrained_fp8, "ALL_FP8_EXPERTS_FUNCTIONS", None)
+    if experts_interface is None:
+        return
+
+    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
+        )
+
+
 # Patch the forward functions of the layers (for compiled models)
 if FbgemmFp8Linear is not None:
     FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, "weight_scale")
 if FP8Linear is not None:
     FP8Linear.forward = module_forward_patch(fp8_block_quant_linear, "weight_scale_inv")
+if CompressedLinear is not None:
+    CompressedLinear.forward = compressed_linear_forward_patch
+_patch_fp8_moe_experts()

unsloth/models/_utils.py

@@ -94,6 +94,18 @@ import functools
 import textwrap
 import logging
 import warnings, subprocess, inspect, psutil, os, math
+
+try:
+    from transformers.utils import auto_docstring
+except:
+
+    def auto_docstring(*args, **kwargs):
+        def decorator(obj):
+            return obj
+
+        return decorator
+
+
 from unsloth_zoo.utils import Version, get_quant_type
 from importlib.metadata import version as importlib_version
 from ..device_type import (

unsloth/models/loader.py

@@ -34,6 +34,7 @@ from transformers import __version__ as transformers_version
 from peft import PeftConfig, PeftModel
 from .loader_utils import (
     _get_fp8_mode_and_check_settings,
+    _has_prequantized_fp8_config,
     _offline_quantize_to_fp8,
     _tag_model_with_fp8_torchao_config,
     get_model_name,
@@ -66,6 +67,18 @@ from unsloth_zoo.utils import Version, _get_dtype
 from unsloth_zoo.hf_utils import dtype_from_config
 from unsloth_zoo.tiled_mlp import patch_tiled_mlp
 
+try:
+    from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
+        maybe_patch_stacked_moe_expert_fp8_scales,
+    )
+except ImportError:
+
+    def maybe_patch_stacked_moe_expert_fp8_scales(
+        model, model_name = None, token = None, revision = None
+    ):
+        return False
+
+
 transformers_version = Version(transformers_version)
 SUPPORTS_FOURBIT = transformers_version >= Version("4.37")
 SUPPORTS_GEMMA = transformers_version >= Version("4.38")
@@ -375,9 +388,23 @@ class FastLanguageModel(FastLlamaModel):
                 model_name,
                 load_in_4bit = load_in_4bit,
                 load_in_fp8 = load_in_fp8,
+                fast_inference = fast_inference,
                 token = token,
                 trust_remote_code = trust_remote_code,
             )
+            if (
+                load_in_fp8 != False
+                and not fast_inference
+                and new_model_name == old_model_name
+            ):
+                if _has_prequantized_fp8_config(
+                    model_name,
+                    token = token,
+                    trust_remote_code = trust_remote_code,
+                ):
+                    load_in_fp8 = False
+                else:
+                    new_model_name = None
             if new_model_name is None and load_in_fp8 != False:
                 fp8_mode = _get_fp8_mode_and_check_settings(
                     load_in_fp8,
@@ -533,9 +560,31 @@ class FastLanguageModel(FastLlamaModel):
                     model_name,
                     load_in_4bit = load_in_4bit,
                     load_in_fp8 = load_in_fp8,
+                    fast_inference = fast_inference,
                     token = token,
                     trust_remote_code = trust_remote_code,
                 )
+                if (
+                    load_in_fp8 != False
+                    and not fast_inference
+                    and model_name == peft_config.base_model_name_or_path
+                ):
+                    if _has_prequantized_fp8_config(
+                        model_name,
+                        token = token,
+                        trust_remote_code = trust_remote_code,
+                    ):
+                        load_in_fp8 = False
+                    else:
+                        fp8_mode = _get_fp8_mode_and_check_settings(
+                            load_in_fp8,
+                            fast_inference,
+                            full_finetuning,
+                            load_in_4bit,
+                            load_in_8bit,
+                            load_in_16bit,
+                        )
+                        model_name = _offline_quantize_to_fp8(model_name, fp8_mode)
             # 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(
@@ -770,6 +819,12 @@ 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 is_peft:
             # From https://github.com/huggingface/peft/issues/184
@@ -984,8 +1039,26 @@ class FastModel(FastBaseModel):
         fp8_mode = None
         if not use_exact_model_name:
             new_model_name = get_model_name(
-                model_name, load_in_4bit = load_in_4bit, load_in_fp8 = load_in_fp8
+                model_name,
+                load_in_4bit = load_in_4bit,
+                load_in_fp8 = load_in_fp8,
+                fast_inference = fast_inference,
+                token = token,
+                trust_remote_code = trust_remote_code,
             )
+            if (
+                load_in_fp8 != False
+                and not fast_inference
+                and new_model_name == old_model_name
+            ):
+                if _has_prequantized_fp8_config(
+                    model_name,
+                    token = token,
+                    trust_remote_code = trust_remote_code,
+                ):
+                    load_in_fp8 = False
+                else:
+                    new_model_name = None
             if new_model_name is None and load_in_fp8 != False:
                 fp8_mode = _get_fp8_mode_and_check_settings(
                     load_in_fp8,
@@ -1495,6 +1568,12 @@ 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 is_peft:
             # From https://github.com/huggingface/peft/issues/184

unsloth/models/loader_utils.py

@@ -106,6 +106,7 @@ def __get_model_name(
     FLOAT_TO_INT_MAPPER = None,
     MAP_TO_UNSLOTH_16bit = None,
     load_in_fp8 = False,
+    fast_inference = False,
     FLOAT_TO_FP8_BLOCK_MAPPER = None,
     FLOAT_TO_FP8_ROW_MAPPER = None,
 ):
@@ -128,7 +129,7 @@ def __get_model_name(
         # For vllm >= 0.12.0, we can quantize the model to FP8 on the fly,
         # so just return the original model name. Older vllm versions will
         # fall through to offline quantization via _offline_quantize_to_fp8.
-        if importlib.util.find_spec("vllm") is not None:
+        if fast_inference and importlib.util.find_spec("vllm") is not None:
             import vllm
 
             if Version(vllm.__version__) >= Version("0.12.0"):
@@ -202,6 +203,7 @@ def _resolve_with_mappers(
     model_name,
     load_in_4bit,
     load_in_fp8,
+    fast_inference,
     int_to_float,
     float_to_int,
     map_to_unsloth_16bit,
@@ -213,6 +215,7 @@ def _resolve_with_mappers(
         FLOAT_TO_INT_MAPPER = float_to_int,
         MAP_TO_UNSLOTH_16bit = map_to_unsloth_16bit,
         load_in_fp8 = load_in_fp8,
+        fast_inference = fast_inference,
         FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER,
         FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER,
     )
@@ -222,6 +225,7 @@ def get_model_name(
     model_name,
     load_in_4bit = True,
     load_in_fp8 = False,
+    fast_inference = False,
     token = None,
     trust_remote_code = False,
 ):
@@ -230,6 +234,7 @@ def get_model_name(
         model_name = model_name,
         load_in_4bit = load_in_4bit,
         load_in_fp8 = load_in_fp8,
+        fast_inference = fast_inference,
         int_to_float = INT_TO_FLOAT_MAPPER,
         float_to_int = FLOAT_TO_INT_MAPPER,
         map_to_unsloth_16bit = MAP_TO_UNSLOTH_16bit,
@@ -245,6 +250,7 @@ def get_model_name(
 
     if (
         new_model_name is None
+        and not (load_in_fp8 != False and not fast_inference)
         and model_name.count("/") == 1
         and model_name[0].isalnum()
     ):
@@ -256,6 +262,7 @@ def get_model_name(
             model_name = model_name,
             load_in_4bit = load_in_4bit,
             load_in_fp8 = load_in_fp8,
+            fast_inference = fast_inference,
             int_to_float = NEW_INT_TO_FLOAT_MAPPER,
             float_to_int = NEW_FLOAT_TO_INT_MAPPER,
             map_to_unsloth_16bit = NEW_MAP_TO_UNSLOTH_16bit,
@@ -274,6 +281,30 @@ def get_model_name(
     return new_model_name
 
 
+def _has_prequantized_fp8_config(
+    model_name,
+    token = None,
+    trust_remote_code = False,
+) -> bool:
+    try:
+        from transformers import AutoConfig
+
+        config = AutoConfig.from_pretrained(
+            model_name,
+            token = token,
+            trust_remote_code = trust_remote_code,
+        )
+    except Exception:
+        return False
+
+    quantization_config = getattr(config, "quantization_config", None)
+    if not isinstance(quantization_config, dict):
+        return False
+
+    quant_method = quantization_config.get("quant_method", None)
+    return quant_method in ("compressed-tensors", "fbgemm_fp8", "fp8")
+
+
 def _offline_quantize_to_fp8(model_name: str, fp8_mode: str) -> str:
     """
     Quantizes the model to fp8 using torchao and saving the quantized model to a
@@ -306,6 +337,15 @@ def _offline_quantize_to_fp8(model_name: str, fp8_mode: str) -> str:
         qconfig = _get_torchao_fp8_config(fp8_mode)
         qconfig = TorchAoConfig(qconfig)
         config = AutoConfig.from_pretrained(model_name)
+        if fp8_mode == "block":
+            incompatible_shapes = _get_block_fp8_incompatible_shapes(config)
+            if incompatible_shapes:
+                raise ValueError(
+                    "Unsloth: Native block FP8 requires weight dimensions compatible "
+                    "with 128x128 block quantization, but this model exposes "
+                    f"incompatible dimensions: {', '.join(incompatible_shapes)}"
+                )
+            _check_block_fp8_deserialize_support()
         is_vlm = any(
             x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
             for x in config.architectures
@@ -329,6 +369,55 @@ def _offline_quantize_to_fp8(model_name: str, fp8_mode: str) -> str:
     return new_model_name
 
 
+def _get_block_fp8_incompatible_shapes(config) -> list[str]:
+    issues = []
+
+    def add_issue(name: str, value):
+        if isinstance(value, int) and value > 0 and value % 128 != 0:
+            issues.append(f"{name}={value}")
+
+    for attr in (
+        "hidden_size",
+        "intermediate_size",
+        "moe_intermediate_size",
+        "head_dim",
+        "v_head_dim",
+        "q_lora_rank",
+        "kv_lora_rank",
+    ):
+        add_issue(attr, getattr(config, attr, None))
+
+    qk_nope_head_dim = getattr(config, "qk_nope_head_dim", None)
+    qk_rope_head_dim = getattr(config, "qk_rope_head_dim", None)
+    kv_lora_rank = getattr(config, "kv_lora_rank", None)
+    if isinstance(qk_nope_head_dim, int) and isinstance(qk_rope_head_dim, int):
+        add_issue(
+            "qk_nope_head_dim + qk_rope_head_dim",
+            qk_nope_head_dim + qk_rope_head_dim,
+        )
+    if isinstance(kv_lora_rank, int) and isinstance(qk_rope_head_dim, int):
+        add_issue(
+            "kv_lora_rank + qk_rope_head_dim",
+            kv_lora_rank + qk_rope_head_dim,
+        )
+
+    return issues
+
+
+def _check_block_fp8_deserialize_support() -> None:
+    try:
+        from torchao.prototype.safetensors import safetensors_utils
+
+        if "PerBlock" not in getattr(safetensors_utils, "ALLOWED_CLASSES", {}):
+            raise ValueError(
+                "Unsloth: This torchao build cannot deserialize block FP8 safetensors "
+                "because `PerBlock` is missing from torchao.prototype.safetensors "
+                "allowed classes. Native block FP8 loading will fail on reload."
+            )
+    except ImportError:
+        return
+
+
 def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
     """
     Tag a model with a `TorchAOConfig` so downstream callers will know what to do with it.
@@ -365,6 +454,8 @@ def _get_fp8_mode_and_check_settings(
     assert load_in_fp8 is not False
     if load_in_fp8 is True:
         fp8_mode = "row"  # default
+        if not fast_inference and os.environ.get("UNSLOTH_HAS_FBGEMM", "0") != "1":
+            fp8_mode = "block"
     else:
         fp8_mode = load_in_fp8