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,