* Add GLM4 MoE DSA model implementation with configurable parameters

* Update Acknowledgments to include GLM4 MoE DSA support

* format

* update ackn.

* Fixes

* Update acknowledgments to include contributions for GLM MoE DSA and additional architectures

* use dsv32 for glm5

* fix

* Fix rope theta

---------

Co-authored-by: Tarjei Mandt <kernelpool@gmail.com>
Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Gökdeniz Gülmez
2026-02-12 19:11:16 +01:00
committed by GitHub
parent 7e67225e1d
commit 1974376d70
3 changed files with 71 additions and 9 deletions
+6 -2
View File
@@ -10,7 +10,7 @@ MLX LM was developed with contributions from the following individuals:
- Shunta Saito: Added support for PLaMo models.
- Gökdeniz Gülmez: Added support for the following architectures:
OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and
`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`,
`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5 (GLM MoE DSA)`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`,
inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`,
Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`,
Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`,
@@ -26,4 +26,8 @@ Added support for the following other features:
MoonshotAI's `Kimi-Linear`, LiquidAI's `LFM2` and `LFM2 MoE`,
Google DeepMind's `Gemma 3`, TII's `Falcon H1` and InterLM's `InternLM 2.5`.
- Ivan Fioravanti: Added support for the following architectures:
ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
- Tarjei Mandt: Added support for the following architectures: `Step 3.5 Flash`,
MoonshotAI's `Kimi K2.5`, Upstage's `Solar Open`, LG AI Research's `K-Exaone MoE`,
Meituan's `LongCat Flash Lite` Helped add support for the following model architectures:
Z.ai & THUKEG's `GLM5 (GLM MoE DSA)`
+12 -7
View File
@@ -71,7 +71,7 @@ class Indexer(nn.Module):
self.rope = initialize_rope(
dims=args.qk_rope_head_dim,
base=args.rope_theta,
traditional=False,
traditional=True,
max_position_embeddings=args.max_position_embeddings,
scaling_config=args.rope_scaling,
)
@@ -495,6 +495,16 @@ class Model(nn.Module):
return self.lm_head(out)
def sanitize(self, weights):
# Remove multi-token prediction layers
mpt_layer = self.args.num_hidden_layers
new_weights = {}
for k, v in weights.items():
parts = k.split(".")
if len(parts) >= 3 and parts[1] == "layers" and int(parts[2]) >= mpt_layer:
continue
new_weights[k] = v
weights = new_weights
def dequant(weight, scale_inv):
dtype = mx.bfloat16
weight = mx.from_fp8(weight, dtype=mx.bfloat16)
@@ -572,12 +582,7 @@ class Model(nn.Module):
weights[f"{prefix}.embed_q.weight"] = wk
weights[f"{prefix}.unembed_out.weight"] = wv
# Remove multi-token prediction layer and any unused precomputed rotary freqs
return {
k: v
for k, v in weights.items()
if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
}
return weights
def shard(self, group: Optional[mx.distributed.Group] = None):
group = group or mx.distributed.init()
+53
View File
@@ -0,0 +1,53 @@
# Copyright © 2025 Apple Inc.
from dataclasses import dataclass
from typing import Any, Dict, Optional
from .base import BaseModelArgs
from .deepseek_v32 import Model as DSV32Model
@dataclass
class ModelArgs(BaseModelArgs):
model_type: str
vocab_size: int
hidden_size: int
index_head_dim: int
index_n_heads: int
index_topk: int
intermediate_size: int
moe_intermediate_size: int
num_hidden_layers: int
num_attention_heads: int
num_key_value_heads: int
n_shared_experts: Optional[int]
n_routed_experts: Optional[int]
routed_scaling_factor: float
kv_lora_rank: int
q_lora_rank: int
qk_rope_head_dim: int
v_head_dim: int
qk_nope_head_dim: int
topk_method: str
scoring_func: str
norm_topk_prob: bool
n_group: int
topk_group: int
num_experts_per_tok: int
moe_layer_freq: int
first_k_dense_replace: int
max_position_embeddings: int
rms_norm_eps: float
rope_parameters: Dict
attention_bias: bool
rope_scaling: Dict = None
rope_theta: Optional[float] = None
def __post_init__(self):
self.rope_scaling = self.rope_parameters
self.rope_theta = self.rope_parameters["rope_theta"]
class Model(DSV32Model):
def __init__(self, config: ModelArgs):
super().__init__(config)