ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

cleanup

Browse source
This commit is contained in:
Datta Nimmaturi 2026-03-16 15:21:59 +00:00
commit 923cac06de
4 changed files with 48 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

25
unsloth/kernels/fp8.py
View file

@ -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)

5
unsloth/models/_utils.py
View file

@ -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 (

67
unsloth/models/loader.py
View file

@ -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,

5
unsloth/models/loader_utils.py
View file

@ -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