Skip to content

Commit 1974376

Browse files
Goekdeniz-GuelmezkernelpoolAwni Hannun
authored
Add GLM5 (ml-explore#867)
* 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>
1 parent 7e67225 commit 1974376

3 files changed

Lines changed: 71 additions & 9 deletions

File tree

ACKNOWLEDGMENTS.md

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ MLX LM was developed with contributions from the following individuals:
1010
- Shunta Saito: Added support for PLaMo models.
1111
- Gökdeniz Gülmez: Added support for the following architectures:
1212
OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and
13-
`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`,
13+
`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5 (GLM MoE DSA)`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`,
1414
inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`,
1515
Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`,
1616
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:
2626
MoonshotAI's `Kimi-Linear`, LiquidAI's `LFM2` and `LFM2 MoE`,
2727
Google DeepMind's `Gemma 3`, TII's `Falcon H1` and InterLM's `InternLM 2.5`.
2828
- Ivan Fioravanti: Added support for the following architectures:
29-
ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
29+
ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
30+
- Tarjei Mandt: Added support for the following architectures: `Step 3.5 Flash`,
31+
MoonshotAI's `Kimi K2.5`, Upstage's `Solar Open`, LG AI Research's `K-Exaone MoE`,
32+
Meituan's `LongCat Flash Lite` Helped add support for the following model architectures:
33+
Z.ai & THUKEG's `GLM5 (GLM MoE DSA)`

mlx_lm/models/deepseek_v32.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def __init__(self, args: ModelArgs):
7171
self.rope = initialize_rope(
7272
dims=args.qk_rope_head_dim,
7373
base=args.rope_theta,
74-
traditional=False,
74+
traditional=True,
7575
max_position_embeddings=args.max_position_embeddings,
7676
scaling_config=args.rope_scaling,
7777
)
@@ -495,6 +495,16 @@ def __call__(
495495
return self.lm_head(out)
496496

497497
def sanitize(self, weights):
498+
# Remove multi-token prediction layers
499+
mpt_layer = self.args.num_hidden_layers
500+
new_weights = {}
501+
for k, v in weights.items():
502+
parts = k.split(".")
503+
if len(parts) >= 3 and parts[1] == "layers" and int(parts[2]) >= mpt_layer:
504+
continue
505+
new_weights[k] = v
506+
weights = new_weights
507+
498508
def dequant(weight, scale_inv):
499509
dtype = mx.bfloat16
500510
weight = mx.from_fp8(weight, dtype=mx.bfloat16)
@@ -572,12 +582,7 @@ def dequant(weight, scale_inv):
572582
weights[f"{prefix}.embed_q.weight"] = wk
573583
weights[f"{prefix}.unembed_out.weight"] = wv
574584

575-
# Remove multi-token prediction layer and any unused precomputed rotary freqs
576-
return {
577-
k: v
578-
for k, v in weights.items()
579-
if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
580-
}
585+
return weights
581586

582587
def shard(self, group: Optional[mx.distributed.Group] = None):
583588
group = group or mx.distributed.init()

mlx_lm/models/glm_moe_dsa.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Copyright © 2025 Apple Inc.
2+
3+
from dataclasses import dataclass
4+
from typing import Any, Dict, Optional
5+
6+
from .base import BaseModelArgs
7+
from .deepseek_v32 import Model as DSV32Model
8+
9+
10+
@dataclass
11+
class ModelArgs(BaseModelArgs):
12+
model_type: str
13+
vocab_size: int
14+
hidden_size: int
15+
index_head_dim: int
16+
index_n_heads: int
17+
index_topk: int
18+
intermediate_size: int
19+
moe_intermediate_size: int
20+
num_hidden_layers: int
21+
num_attention_heads: int
22+
num_key_value_heads: int
23+
n_shared_experts: Optional[int]
24+
n_routed_experts: Optional[int]
25+
routed_scaling_factor: float
26+
kv_lora_rank: int
27+
q_lora_rank: int
28+
qk_rope_head_dim: int
29+
v_head_dim: int
30+
qk_nope_head_dim: int
31+
topk_method: str
32+
scoring_func: str
33+
norm_topk_prob: bool
34+
n_group: int
35+
topk_group: int
36+
num_experts_per_tok: int
37+
moe_layer_freq: int
38+
first_k_dense_replace: int
39+
max_position_embeddings: int
40+
rms_norm_eps: float
41+
rope_parameters: Dict
42+
attention_bias: bool
43+
rope_scaling: Dict = None
44+
rope_theta: Optional[float] = None
45+
46+
def __post_init__(self):
47+
self.rope_scaling = self.rope_parameters
48+
self.rope_theta = self.rope_parameters["rope_theta"]
49+
50+
51+
class Model(DSV32Model):
52+
def __init__(self, config: ModelArgs):
53+
super().__init__(config)

0 commit comments

Comments
 (0)