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

[WIP] cleanup

Browse source
This commit is contained in:
Datta Nimmaturi 2026-03-15 14:49:07 +00:00
commit da3c9c1e84
2 changed files with 6 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

13
unsloth/models/glm4_moe.py
View file

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

8
unsloth/models/loader.py
View file

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