diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index a57f4ffb64..044e5a9292 100644 --- a/unsloth/kernels/fp8.py +++ b/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): @@ -432,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 @@ -596,6 +605,11 @@ 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 @@ -604,9 +618,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 @@ -617,8 +640,161 @@ 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] + 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] + 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_fp8 import ( + forward_moe_backend_fp8, + ) + except Exception: + return + + experts_interface = getattr(finegrained_fp8, "ALL_FP8_EXPERTS_FUNCTIONS", None) + if experts_interface is not None: + experts_interface["grouped_mm"] = forward_moe_backend_fp8 + experts_interface["batched_mm"] = forward_moe_backend_fp8 + + 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() diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index cbaebcc7ac..37d3e9a2bc 100644 --- a/unsloth/models/_utils.py +++ b/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 ( diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index bd15ed5281..38811257ff 100644 --- a/unsloth/models/loader.py +++ b/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,15 @@ class FastLanguageModel(FastLlamaModel): if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) + 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 @@ -984,8 +1042,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, @@ -1286,7 +1362,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( @@ -1495,6 +1578,15 @@ class FastModel(FastBaseModel): if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) + 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 cf5af983a6..fa970372ee 100644 --- a/unsloth/models/loader_utils.py +++ b/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,48 @@ 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) + 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: """ Quantizes the model to fp8 using torchao and saving the quantized model to a @@ -306,6 +355,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 +387,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 +472,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