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

Add Gemma 4 26B MoE support (MLX) (#4844)

Browse source 
* removed bidirectional check for 31b

* Change gemma4_text for moe

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Manan Shah <mananshah@Manans-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Manan Shah 2026-04-04 02:31:02 -05:00 • committed by GitHub
commit 324bdf429d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 135 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

144
unsloth/models/gemma4_text.py
View file

@ -46,6 +46,9 @@ class ModelArgs(BaseModelArgs):
num_kv_shared_layers: int = 0
use_double_wide_mlp: bool = False
enable_moe_block: bool = False
num_experts: Optional[int] = None
top_k_experts: Optional[int] = None
moe_intermediate_size: Optional[int] = None
hidden_activation: str = "gelu_pytorch_tanh"
tie_word_embeddings: bool = True
final_logit_softcapping: Optional[float] = None
@ -123,6 +126,77 @@ class MLP(nn.Module):
return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
class Router(nn.Module):
def __init__(self, args: ModelArgs):
super().__init__()
self.args = args
self.norm = Gemma4RMSNorm(
args.hidden_size, eps = args.rms_norm_eps, with_scale = False
)
self.proj = nn.Linear(args.hidden_size, args.num_experts, bias = False)
self.scale = mx.ones((args.hidden_size,))
self.per_expert_scale = mx.ones((args.num_experts,))
self._root_size = args.hidden_size**-0.5
def __call__(self, x: mx.array):
x = self.norm(x)
x = x * self._root_size
x = x * self.scale
expert_scores = self.proj(x)
router_probs = mx.softmax(expert_scores, axis = -1)
top_k_indices = mx.argpartition(
-expert_scores, kth = self.args.top_k_experts - 1, axis = -1
)[..., : self.args.top_k_experts]
top_k_weights = mx.take_along_axis(router_probs, top_k_indices, axis = -1)
top_k_weights = top_k_weights / mx.sum(top_k_weights, axis = -1, keepdims = True)
top_k_weights = top_k_weights * self.per_expert_scale[top_k_indices]
return top_k_indices, top_k_weights
class GeGLU(nn.Module):
def __call__(self, x, gate):
return nn.gelu_approx(gate) * x
class Experts(nn.Module):
def __init__(self, args: ModelArgs):
super().__init__()
try:
from .switch_layers import SwitchGLU
except ImportError:
raise ImportError(
"Gemma4 MoE requires mlx-lm >= 0.31. Please upgrade: pip install -U mlx-lm"
)
self.switch_glu = SwitchGLU(
input_dims = args.hidden_size,
hidden_dims = args.moe_intermediate_size,
num_experts = args.num_experts,
activation = GeGLU(),
bias = False,
)
def __call__(
self,
x: mx.array,
top_k_indices: mx.array,
top_k_weights: mx.array,
) -> mx.array:
B, S, H = x.shape
K = top_k_indices.shape[-1]
x_flat = x.reshape(B * S, H)
indices_flat = top_k_indices.reshape(B * S, K)
expert_out = self.switch_glu(x_flat, indices_flat)
weights = top_k_weights.reshape(B * S, K)[..., None]
return (expert_out * weights).sum(axis = -2).reshape(B, S, H)
def build_rope(args: ModelArgs, layer_type: str, head_dim: int):
rope_config = args.rope_parameters[layer_type]
rope_type = rope_config.get("rope_type", "default")
@ -317,6 +391,21 @@ class TransformerBlock(nn.Module):
)
self.layer_scalar = mx.ones((1,))
# MoE
self.enable_moe = args.enable_moe_block
if self.enable_moe:
self.router = Router(args)
self.experts = Experts(args)
self.post_feedforward_layernorm_1 = Gemma4RMSNorm(
args.hidden_size, eps = args.rms_norm_eps
)
self.post_feedforward_layernorm_2 = Gemma4RMSNorm(
args.hidden_size, eps = args.rms_norm_eps
)
self.pre_feedforward_layernorm_2 = Gemma4RMSNorm(
args.hidden_size, eps = args.rms_norm_eps
)
if self.hidden_size_per_layer_input:
self.act = ACT2FN[args.hidden_activation]
self.per_layer_input_gate = nn.Linear(
@ -344,10 +433,24 @@ class TransformerBlock(nn.Module):
h = residual + h
residual = h
ff = self.pre_feedforward_layernorm(h)
ff = self.mlp(ff)
ff = self.post_feedforward_layernorm(ff)
h = residual + ff
if self.enable_moe:
h1 = self.pre_feedforward_layernorm(h)
h1 = self.mlp(h1)
h1 = self.post_feedforward_layernorm_1(h1)
top_k_indices, top_k_weights = self.router(h)
h2 = self.pre_feedforward_layernorm_2(h)
h2 = self.experts(h2, top_k_indices, top_k_weights)
h2 = self.post_feedforward_layernorm_2(h2)
h = h1 + h2
else:
h = self.pre_feedforward_layernorm(h)
h = self.mlp(h)
h = self.post_feedforward_layernorm(h)
h = residual + h
if self.hidden_size_per_layer_input:
residual = h
@ -371,10 +474,6 @@ def logit_softcap(softcap, x):
class Gemma4Model(nn.Module):
def __init__(self, args: ModelArgs):
super().__init__()
if args.enable_moe_block:
raise NotImplementedError(
"Gemma4 MoE layers are not implemented in mlx-lm."
)
self.args = args
self.vocab_size = args.vocab_size
@ -610,7 +709,34 @@ class Model(nn.Module):
if "lm_head.weight" not in weights:
self.tie_word_embeddings = True
self.pop("lm_head")
return weights
sanitized = {}
for k, v in weights.items():
if "rotary_emb" in k:
continue
if k.endswith(".experts.down_proj"):
k = k.replace(
".experts.down_proj", ".experts.switch_glu.down_proj.weight"
)
sanitized[k] = v
continue
if k.endswith(".experts.gate_up_proj"):
gate_key = k.replace(
".experts.gate_up_proj", ".experts.switch_glu.gate_proj.weight"
)
up_key = k.replace(
".experts.gate_up_proj", ".experts.switch_glu.up_proj.weight"
)
v = v.swapaxes(-1, -2)
mid_dim = v.shape[-1] // 2
sanitized[gate_key] = v[..., :mid_dim].swapaxes(-1, -2)
sanitized[up_key] = v[..., mid_dim:].swapaxes(-1, -2)
continue
sanitized[k] = v
return sanitized
@property
def layers(self):