From 923cac06de94e4ea79fabbefc4b5fe17e6540ec2 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Mon, 16 Mar 2026 15:21:59 +0000 Subject: [PATCH] cleanup --- unsloth/kernels/fp8.py | 25 ++++++++----- unsloth/models/_utils.py | 5 +++ unsloth/models/loader.py | 67 +++++++++++++--------------------- unsloth/models/loader_utils.py | 5 +-- 4 files changed, 48 insertions(+), 54 deletions(-) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index b7fb43adb1..a1d88444e6 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -628,9 +628,9 @@ def module_forward_patch(forward_function, scale_attr = "weight_scale"): def _compressed_linear_supports_unsloth_fp8(self): if ( - CompressedLinear is None or - QuantizationStrategy is None or - QuantizationStatus is None + CompressedLinear is None + or QuantizationStrategy is None + or QuantizationStatus is None ): return False @@ -639,10 +639,10 @@ def _compressed_linear_supports_unsloth_fp8(self): 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 + weight is None + or weight_scale is None + or quantization_args is None + or weight.dtype != torch.float8_e4m3fn ): return False @@ -679,6 +679,7 @@ def _compressed_linear_forward_fallback(self, input): 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 @@ -715,7 +716,11 @@ def _fp8_moe_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts): param_name = getattr(wrapper, "parameter_name", None) - if param_name == "down_proj" and intermediate_dim is not None and hidden_dim is not 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) @@ -764,7 +769,9 @@ def _patch_fp8_moe_experts(): 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) + finegrained_fp8.FP8Experts._unsloth_lora_extractor_fn = staticmethod( + _fp8_moe_lora_extractor + ) # Patch the forward functions of the layers (for compiled models) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 9eddab9f71..37d3e9a2bc 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -94,13 +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 ( diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index cddc897320..7d978cf184 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -66,11 +66,18 @@ from ..device_type import ( 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 Exception as e: - maybe_patch_stacked_moe_expert_fp8_scales = None - logger.info(f"Unsloth: Failed to import Moe FP8 patches due to {e}. Please update unsloth and unsloth_zoo using `pip install unsloth unsloth_zoo --upgrade`") + 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") @@ -225,32 +232,6 @@ def _fix_rope_inv_freq(model): return model -def _apply_post_load_fp8_patches( - model, - config, - load_in_fp8, - fp8_mode, - model_name, - token, - revision, -): - if load_in_fp8 != False: - _tag_model_with_fp8_torchao_config(model, fp8_mode) - quant_config = getattr(config, "quantization_config", None) - quant_method = getattr(quant_config, "quant_method", None) if quant_config else None - if quant_method and quant_method in ["compressed-tensors", "fp8", "fbgemm_fp8"] and maybe_patch_stacked_moe_expert_fp8_scales is None and is_moe_model(model): - raise ValueError("FP8 MoE models need unslot and unsloth_zoo to be updated. Please update via `pip install unsloth unsloth_zoo --upgrade`") - if maybe_patch_stacked_moe_expert_fp8_scales is not None: - maybe_patch_stacked_moe_expert_fp8_scales( - model, - model_name = model_name, - token = token, - revision = revision, - ) - - return model - - class FastLanguageModel(FastLlamaModel): @staticmethod def from_pretrained( @@ -836,11 +817,10 @@ class FastLanguageModel(FastLlamaModel): elif isinstance(quantization_config, dict): model.config.update({"quantization_config": quantization_config}) - model = _apply_post_load_fp8_patches( - model = model, - config = model_config, - load_in_fp8 = load_in_fp8, - fp8_mode = fp8_mode, + 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, @@ -1063,13 +1043,19 @@ class FastModel(FastBaseModel): 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): + if _has_prequantized_fp8_config( + model_name, + token = token, + trust_remote_code = trust_remote_code, + ): load_in_fp8 = False else: new_model_name = None @@ -1580,11 +1566,10 @@ class FastModel(FastBaseModel): elif isinstance(quantization_config, dict): model.config.update({"quantization_config": quantization_config}) - model = _apply_post_load_fp8_patches( - model = model, - config = model_config, - load_in_fp8 = load_in_fp8, - fp8_mode = fp8_mode, + 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, diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 113506e73c..c5c417b0b2 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -454,10 +454,7 @@ 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" - ): + if not fast_inference and os.environ.get("UNSLOTH_HAS_FBGEMM", "0") != "1": fp8_mode = "block" else: fp8_mode = load_in_fp8