From dec0c184003a4969425f1d6f163a00744a237074 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sun, 15 Mar 2026 13:49:27 +0000 Subject: [PATCH 1/7] [WIP] Fp8 training for Moe --- unsloth/kernels/fp8.py | 76 +++++++++++++++++++++++++++ unsloth/models/_utils.py | 7 +++ unsloth/models/glm4_moe.py | 13 ++++- unsloth/models/loader.py | 66 ++++++++++++++++++++++- unsloth/models/loader_utils.py | 96 +++++++++++++++++++++++++++++++++- unsloth/models/qwen3.py | 15 ++++-- 6 files changed, 265 insertions(+), 8 deletions(-) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index a57f4ffb64..48824cdd17 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): @@ -617,8 +626,75 @@ 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) + + # 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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index cbaebcc7ac..9eddab9f71 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -94,6 +94,13 @@ 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/glm4_moe.py b/unsloth/models/glm4_moe.py index 5d04b2f1d0..3abcb0947e 100644 --- a/unsloth/models/glm4_moe.py +++ b/unsloth/models/glm4_moe.py @@ -123,6 +123,15 @@ except ImportError: torch_nn_functional_silu = torch.nn.functional.silu +def _glm4_supports_grouped_gemm_fp8_safe(experts_module) -> bool: + """ + GLM grouped GEMM expects high-precision expert weights. + FP8 routed expert weights must use the existing naive fallback path. + """ + gate_up_proj = getattr(experts_module, "gate_up_proj", None) + return gate_up_proj is not None and gate_up_proj.dtype != torch.float8_e4m3fn + + def Glm4MoeLiteMoE_fast_forward(self, hidden_states): """ Optimized MoE forward pass using grouped GEMM. @@ -155,7 +164,7 @@ def Glm4MoeLiteMoE_fast_forward(self, hidden_states): ) # Use grouped GEMM for expert computation - if HAS_GROUPED_GEMM: + if HAS_GROUPED_GEMM and _glm4_supports_grouped_gemm_fp8_safe(self.experts): # Cast hidden_states to match expert weights dtype # Under autocast, hidden_states may be fp32 while weights are bf16 hidden_states = hidden_states.to(self.experts.gate_up_proj.dtype) @@ -229,7 +238,7 @@ def Glm4MoeLiteNaiveMoe_fast_forward( # Cast routing weights to match hidden_states dtype (Qwen3 pattern) top_k_weights = top_k_weights.to(hidden_states.dtype) - if not HAS_GROUPED_GEMM: + if not HAS_GROUPED_GEMM or not _glm4_supports_grouped_gemm_fp8_safe(self): # Fallback to original naive implementation final_hidden_states = torch.zeros_like(hidden_states) with torch.no_grad(): diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index bd15ed5281..c87bd60ac7 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, @@ -65,6 +66,9 @@ 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 +from unsloth_zoo.temporary_patches.glm4_moe import ( + maybe_patch_glm4_moe_expert_fp8_scales, +) transformers_version = Version(transformers_version) SUPPORTS_FOURBIT = transformers_version >= Version("4.37") @@ -375,9 +379,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 +551,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 +810,12 @@ class FastLanguageModel(FastLlamaModel): if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) + maybe_patch_glm4_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 +1030,20 @@ 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, ) + if ( + load_in_fp8 != False + and not fast_inference + and new_model_name == old_model_name + ): + if _has_prequantized_fp8_config(model_name): + 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 +1553,12 @@ class FastModel(FastBaseModel): if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) + maybe_patch_glm4_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..113506e73c 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,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,11 @@ 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 diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index b93dddb186..43068ad481 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -300,10 +300,17 @@ def Qwen3Attention_fast_forward_inference( # Need to do it prior 2 steps before hitting full on short KV cache # or else error - self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2) - cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) - cos = cos[position_ids].unsqueeze(1) - sin = sin[position_ids].unsqueeze(1) + if position_ids.dim() == 1: + position_ids = position_ids[:, None] + position_ids = position_ids.to(Qn.device) + if position_ids.shape[-1] != Qn.shape[-2]: + position_ids = position_ids[:, -Qn.shape[-2]:] + + rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1) + self.rotary_emb.extend_rope_embedding(Vn, rotary_seq_len + 1) # +1 slack + cos, sin = self.rotary_emb.get_cached(rotary_seq_len, Qn.device.index or 0) + cos = cos[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype) + sin = sin[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype) h = self.half_head_dim RH_Q = self.RH_Q From da3c9c1e84b90cecb115d7a043848d69818af69f Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sun, 15 Mar 2026 14:49:07 +0000 Subject: [PATCH 2/7] [WIP] cleanup --- unsloth/models/glm4_moe.py | 13 ++----------- unsloth/models/loader.py | 8 ++++---- 2 files changed, 6 insertions(+), 15 deletions(-) diff --git a/unsloth/models/glm4_moe.py b/unsloth/models/glm4_moe.py index 3abcb0947e..5d04b2f1d0 100644 --- a/unsloth/models/glm4_moe.py +++ b/unsloth/models/glm4_moe.py @@ -123,15 +123,6 @@ except ImportError: torch_nn_functional_silu = torch.nn.functional.silu -def _glm4_supports_grouped_gemm_fp8_safe(experts_module) -> bool: - """ - GLM grouped GEMM expects high-precision expert weights. - FP8 routed expert weights must use the existing naive fallback path. - """ - gate_up_proj = getattr(experts_module, "gate_up_proj", None) - return gate_up_proj is not None and gate_up_proj.dtype != torch.float8_e4m3fn - - def Glm4MoeLiteMoE_fast_forward(self, hidden_states): """ Optimized MoE forward pass using grouped GEMM. @@ -164,7 +155,7 @@ def Glm4MoeLiteMoE_fast_forward(self, hidden_states): ) # Use grouped GEMM for expert computation - if HAS_GROUPED_GEMM and _glm4_supports_grouped_gemm_fp8_safe(self.experts): + if HAS_GROUPED_GEMM: # Cast hidden_states to match expert weights dtype # Under autocast, hidden_states may be fp32 while weights are bf16 hidden_states = hidden_states.to(self.experts.gate_up_proj.dtype) @@ -238,7 +229,7 @@ def Glm4MoeLiteNaiveMoe_fast_forward( # Cast routing weights to match hidden_states dtype (Qwen3 pattern) top_k_weights = top_k_weights.to(hidden_states.dtype) - if not HAS_GROUPED_GEMM or not _glm4_supports_grouped_gemm_fp8_safe(self): + if not HAS_GROUPED_GEMM: # Fallback to original naive implementation final_hidden_states = torch.zeros_like(hidden_states) with torch.no_grad(): diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index c87bd60ac7..29bf95308d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -66,8 +66,8 @@ 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 -from unsloth_zoo.temporary_patches.glm4_moe import ( - maybe_patch_glm4_moe_expert_fp8_scales, +from unsloth_zoo.temporary_patches.moe_utils_fp8 import ( + maybe_patch_stacked_moe_expert_fp8_scales, ) transformers_version = Version(transformers_version) @@ -810,7 +810,7 @@ class FastLanguageModel(FastLlamaModel): if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) - maybe_patch_glm4_moe_expert_fp8_scales( + maybe_patch_stacked_moe_expert_fp8_scales( model, model_name = model_name, token = token, @@ -1553,7 +1553,7 @@ class FastModel(FastBaseModel): if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) - maybe_patch_glm4_moe_expert_fp8_scales( + maybe_patch_stacked_moe_expert_fp8_scales( model, model_name = model_name, token = token, From c6f38aad8b2323d877999c754e470156a1045266 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Mon, 16 Mar 2026 05:51:42 +0000 Subject: [PATCH 3/7] patch for fp8 moe to use unsloth kerenls --- unsloth/kernels/fp8.py | 81 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 48824cdd17..42f5eea045 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -691,6 +691,86 @@ def compressed_linear_forward_patch(self, input): 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 + + # Pre-quantized FP8 MoE checkpoints replace `.experts` modules with + # transformers.integrations.finegrained_fp8.FP8Experts. Route those + # implementations to Unsloth's MoE backend so we avoid the optional + # Hugging Face `kernels` package at training time. + 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") @@ -698,3 +778,4 @@ 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() From 5a3b1ed73befd54ee9c7fbcb49261a543b7b40aa Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Mon, 16 Mar 2026 07:10:10 +0000 Subject: [PATCH 4/7] Undo qwen3 rope changes --- unsloth/kernels/fp8.py | 4 ---- unsloth/models/qwen3.py | 15 ++++----------- 2 files changed, 4 insertions(+), 15 deletions(-) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 42f5eea045..b7fb43adb1 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -761,10 +761,6 @@ def _patch_fp8_moe_experts(): if experts_interface is None: return - # Pre-quantized FP8 MoE checkpoints replace `.experts` modules with - # transformers.integrations.finegrained_fp8.FP8Experts. Route those - # implementations to Unsloth's MoE backend so we avoid the optional - # Hugging Face `kernels` package at training time. experts_interface["grouped_mm"] = forward_moe_backend experts_interface["batched_mm"] = forward_native_moe_loop if hasattr(finegrained_fp8, "FP8Experts"): diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index 43068ad481..b93dddb186 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -300,17 +300,10 @@ def Qwen3Attention_fast_forward_inference( # Need to do it prior 2 steps before hitting full on short KV cache # or else error - if position_ids.dim() == 1: - position_ids = position_ids[:, None] - position_ids = position_ids.to(Qn.device) - if position_ids.shape[-1] != Qn.shape[-2]: - position_ids = position_ids[:, -Qn.shape[-2]:] - - rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1) - self.rotary_emb.extend_rope_embedding(Vn, rotary_seq_len + 1) # +1 slack - cos, sin = self.rotary_emb.get_cached(rotary_seq_len, Qn.device.index or 0) - cos = cos[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype) - sin = sin[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype) + self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2) + cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) + cos = cos[position_ids].unsqueeze(1) + sin = sin[position_ids].unsqueeze(1) h = self.half_head_dim RH_Q = self.RH_Q From daa13cf3ade875cd0b88cc3f45555bda827aeaa2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 07:10:47 +0000 Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/kernels/fp8.py | 25 ++++++++++++++++--------- unsloth/models/_utils.py | 5 +++++ unsloth/models/loader_utils.py | 5 +---- 3 files changed, 22 insertions(+), 13 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_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 From dde18ddd28a14c842cbb27a4e5c0be026d6d7e63 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 14:32:34 +0000 Subject: [PATCH 6/7] Fix FP8 MoE loader: guard import, pass auth to FastModel 1. Wrap moe_utils_fp8 import in try/except so unsloth does not crash when paired with an older unsloth_zoo that lacks the module. Falls back to a no-op stub. 2. Pass token and trust_remote_code to get_model_name() and _has_prequantized_fp8_config() in the FastModel path, matching the FastLanguageModel path. Without this, private/gated repos fail silently and fall through to incorrect offline quantization. --- unsloth/models/loader.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 29bf95308d..ea7cb60295 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -66,9 +66,13 @@ 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 -from unsloth_zoo.temporary_patches.moe_utils_fp8 import ( - maybe_patch_stacked_moe_expert_fp8_scales, -) +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") @@ -1034,13 +1038,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 From 418b43dca2da56f3f7a89d048dbe91ddbeb9dd8b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 14:33:53 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/loader.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index ea7cb60295..7d978cf184 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -66,14 +66,19 @@ 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 ImportError: - def maybe_patch_stacked_moe_expert_fp8_scales(model, model_name=None, token=None, revision=None): + + 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")