From 70f2b090f20289438855a83a78f5b2f079359987 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?= <60228478+Goekdeniz-Guelmez@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:32:22 +0100 Subject: [PATCH 1/9] Add GLM4 MoE DSA model implementation with configurable parameters --- mlx_lm/models/glm_moe_dsa.py | 225 +++++++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 mlx_lm/models/glm_moe_dsa.py diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py new file mode 100644 index 000000000..5d24dae74 --- /dev/null +++ b/mlx_lm/models/glm_moe_dsa.py @@ -0,0 +1,225 @@ +# Copyright © 2026 Apple Inc. + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.distributed import shard_inplace, shard_linear + +from .base import BaseModelArgs, create_attention_mask +from .pipeline import PipelineMixin + +from .deepseek_v32 import DeepseekV32Attention +from .glm4_moe_lite import Glm4MoeLiteMoE, Glm4MoeLiteMLP + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str = "glm4_moe_dsa" + vocab_size: int = 154880 + hidden_size: int = 2048 + intermediate_size: int = 10240 + moe_intermediate_size: int = 1536 + num_hidden_layers: int = 47 + num_attention_heads: int = 20 + num_key_value_heads: int = 20 + n_shared_experts: Optional[int] = 1 + n_routed_experts: Optional[int] = 64 + routed_scaling_factor: float = 1.8 + kv_lora_rank: int = 512 + q_lora_rank: int = 768 + qk_rope_head_dim: int = 64 + qk_nope_head_dim: int = 192 + v_head_dim: int = 256 + topk_method: str = "noaux_tc" + scoring_func: str = "sigmoid" + norm_topk_prob: bool = True + n_group: int = 1 + topk_group: int = 1 + num_experts_per_tok: int = 4 + moe_layer_freq: int = 1 + first_k_dense_replace: int = 1 + max_position_embeddings: int = 202752 + rms_norm_eps: float = 1e-5 + rope_theta: float = 1_000_000.0 + rope_scaling: Optional[Dict] = None + attention_bias: bool = False + attention_dropout: float = 0.0 + partial_rotary_factor: float = 1.0 + tie_word_embeddings: bool = False + num_nextn_predict_layers: int = 1 + quantization: Optional[Dict[str, Any]] = None + + +class Glm4MoeDSADecoderLayer(nn.Module): + def __init__(self, config: ModelArgs, layer_idx: int): + super().__init__() + self.self_attn = DeepseekV32Attention(config) + use_moe = ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + self.mlp = Glm4MoeLiteMoE(config) if use_moe else Glm4MoeLiteMLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + return h + r + + +class Glm4MoeDSAModel(PipelineMixin, nn.Module): + def __init__(self, config: ModelArgs): + super().__init__() + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = [ + Glm4MoeDSADecoderLayer(config, idx) for idx in range(config.num_hidden_layers) + ] + self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def __call__( + self, + x: mx.array, + cache: Optional[Any] = None, + ) -> mx.array: + h = self.embed_tokens(x) + + pipeline_rank = self.pipeline_rank + pipeline_size = self.pipeline_size + + if cache is None: + cache = [None] * len(self.pipeline_layers) + mask = create_attention_mask(h, cache[0]) + + if pipeline_rank < pipeline_size - 1: + h = mx.distributed.recv_like(h, (pipeline_rank + 1)) + + for l, c in zip(self.pipeline_layers, cache): + h = l(h, mask, cache=c) + + if pipeline_rank != 0: + h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) + if cache[-1] is not None: + cache[-1].keys = mx.depends(cache[-1].keys, h) + + if pipeline_size > 1: + h = mx.distributed.all_gather(h)[: h.shape[0]] + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, config: ModelArgs): + super().__init__() + self.args = config + self.model_type = config.model_type + self.model = Glm4MoeDSAModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + ) -> mx.array: + out = self.model(inputs, cache) + return self.lm_head(out) + + def sanitize(self, weights): + mpt_layer = self.args.num_hidden_layers + + # Stack experts + for l in range(self.args.num_hidden_layers): + prefix = f"model.layers.{l}" + for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]: + for k in ["weight", "scales", "biases"]: + if f"{prefix}.mlp.experts.0.{m}.{k}" in weights: + to_join = [ + weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}") + for e in range(self.args.n_routed_experts) + ] + weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join) + + # Remove multi-token prediction layer + return { + k: v + for k, v in weights.items() + if not k.startswith(f"model.layers.{mpt_layer}") + } + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + for layer in self.model.layers: + # Shard the self attention + layer.self_attn.q_proj = shard_linear( + layer.self_attn.q_proj, "all-to-sharded", group=group + ) + layer.self_attn.k_proj = shard_linear( + layer.self_attn.k_proj, "all-to-sharded", group=group + ) + layer.self_attn.v_proj = shard_linear( + layer.self_attn.v_proj, "all-to-sharded", group=group + ) + layer.self_attn.o_proj = shard_linear( + layer.self_attn.o_proj, "sharded-to-all", group=group + ) + layer.self_attn.n_heads //= N + layer.self_attn.n_kv_heads //= N + + # Shard the MLP + if isinstance(layer.mlp, Glm4MoeLiteMLP): + layer.mlp.gate_proj = shard_linear( + layer.mlp.gate_proj, "all-to-sharded", group=group + ) + layer.mlp.down_proj = shard_linear( + layer.mlp.down_proj, "sharded-to-all", group=group + ) + layer.mlp.up_proj = shard_linear( + layer.mlp.up_proj, "all-to-sharded", group=group + ) + + # Shard the MoE. Shard in place since the MoE should be responsible + # for aggregating the results. + else: + layer.mlp.sharding_group = group + shard_inplace( + layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group + ) + shard_inplace( + layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group + ) + shard_inplace( + layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group + ) + shard_inplace( + layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group + ) + shard_inplace( + layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group + ) + shard_inplace( + layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group + ) + + @property + def layers(self): + return self.model.pipeline_layers + + @property + def cast_predicate(self): + def predicate(k): + return "e_score_correction_bias" not in k + + return predicate From 7bf605373656e29bf420338a2f7a083efed94732 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?= <60228478+Goekdeniz-Guelmez@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:33:45 +0100 Subject: [PATCH 2/9] Update Acknowledgments to include GLM4 MoE DSA support --- ACKNOWLEDGMENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index b452d5b71..1fb4ead71 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -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`, `GLM4 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`, From a47e59bb748e1ee0ac99c6b3ee1827c44e2d2e34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?= <60228478+Goekdeniz-Guelmez@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:34:15 +0100 Subject: [PATCH 3/9] format --- mlx_lm/models/glm_moe_dsa.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index 5d24dae74..ca46603aa 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/mlx_lm/models/glm_moe_dsa.py @@ -8,10 +8,9 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear from .base import BaseModelArgs, create_attention_mask -from .pipeline import PipelineMixin - from .deepseek_v32 import DeepseekV32Attention -from .glm4_moe_lite import Glm4MoeLiteMoE, Glm4MoeLiteMLP +from .glm4_moe_lite import Glm4MoeLiteMLP, Glm4MoeLiteMoE +from .pipeline import PipelineMixin @dataclass @@ -85,7 +84,8 @@ def __init__(self, config: ModelArgs): self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = [ - Glm4MoeDSADecoderLayer(config, idx) for idx in range(config.num_hidden_layers) + Glm4MoeDSADecoderLayer(config, idx) + for idx in range(config.num_hidden_layers) ] self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) From aea6efb43cc45f308524334c3549ad7e3e9367bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?= <60228478+Goekdeniz-Guelmez@users.noreply.github.com> Date: Wed, 11 Feb 2026 21:37:26 +0100 Subject: [PATCH 4/9] update ackn. --- ACKNOWLEDGMENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 1fb4ead71..f457e7b66 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -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`, `GLM4 MoE DSA`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, +`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5`, 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`, From 0352b8dd5e0f60225ceefbb025b481aa288dfb28 Mon Sep 17 00:00:00 2001 From: Tarjei Mandt Date: Thu, 12 Feb 2026 14:03:20 +1100 Subject: [PATCH 5/9] Fixes --- mlx_lm/models/glm_moe_dsa.py | 89 ++++++++++++++++++++++++++++-------- 1 file changed, 70 insertions(+), 19 deletions(-) diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index ca46603aa..ef6813de1 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/mlx_lm/models/glm_moe_dsa.py @@ -8,6 +8,7 @@ from mlx.nn.layers.distributed import shard_inplace, shard_linear from .base import BaseModelArgs, create_attention_mask +from .cache import CacheList, KVCache from .deepseek_v32 import DeepseekV32Attention from .glm4_moe_lite import Glm4MoeLiteMLP, Glm4MoeLiteMoE from .pipeline import PipelineMixin @@ -48,6 +49,9 @@ class ModelArgs(BaseModelArgs): partial_rotary_factor: float = 1.0 tie_word_embeddings: bool = False num_nextn_predict_layers: int = 1 + index_head_dim: int = 128 + index_n_heads: int = 32 + index_topk: int = 2048 quantization: Optional[Dict[str, Any]] = None @@ -101,7 +105,9 @@ def __call__( if cache is None: cache = [None] * len(self.pipeline_layers) - mask = create_attention_mask(h, cache[0]) + mask = create_attention_mask( + h, cache[0][0] if cache[0] else None, return_array=True + ) if pipeline_rank < pipeline_size - 1: h = mx.distributed.recv_like(h, (pipeline_rank + 1)) @@ -112,7 +118,7 @@ def __call__( if pipeline_rank != 0: h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) if cache[-1] is not None: - cache[-1].keys = mx.depends(cache[-1].keys, h) + cache[-1][0].keys = mx.depends(cache[-1][0].keys, h) if pipeline_size > 1: h = mx.distributed.all_gather(h)[: h.shape[0]] @@ -137,9 +143,17 @@ def __call__( 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 - # Stack experts + # Stack experts and absorb MLA weights for l in range(self.args.num_hidden_layers): prefix = f"model.layers.{l}" for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]: @@ -151,32 +165,66 @@ def sanitize(self, weights): ] weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join) - # Remove multi-token prediction layer - return { - k: v - for k, v in weights.items() - if not k.startswith(f"model.layers.{mpt_layer}") - } + # MLA absorption: split kv_b_proj into embed_q and unembed_out + attn_prefix = f"{prefix}.self_attn" + if f"{attn_prefix}.kv_b_proj.weight" in weights: + quantized = f"{attn_prefix}.kv_b_proj.scales" in weights + v = weights.pop(f"{attn_prefix}.kv_b_proj.weight") + head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim + + if quantized: + dims = self.args.kv_lora_rank + scales = weights.pop(f"{attn_prefix}.kv_b_proj.scales") + biases = weights.pop(f"{attn_prefix}.kv_b_proj.biases") + bits = (v.shape[-1] * 32) // dims + group_size = dims // scales.shape[-1] + v = mx.dequantize( + v, scales, biases, bits=bits, group_size=group_size + ) + num_heads = self.args.num_attention_heads + v = v.reshape(num_heads, head_dim, -1) + wk = mx.contiguous( + v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2) + ) + wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :]) + if quantized: + wk, wk_scales, wk_biases = mx.quantize( + wk, bits=bits, group_size=group_size + ) + wv, wv_scales, wv_biases = mx.quantize( + wv, bits=bits, group_size=group_size + ) + weights[f"{attn_prefix}.embed_q.scales"] = wk_scales + weights[f"{attn_prefix}.unembed_out.scales"] = wv_scales + weights[f"{attn_prefix}.embed_q.biases"] = wk_biases + weights[f"{attn_prefix}.unembed_out.biases"] = wv_biases + weights[f"{attn_prefix}.embed_q.weight"] = wk + weights[f"{attn_prefix}.unembed_out.weight"] = wv + + return weights def shard(self, group: Optional[mx.distributed.Group] = None): group = group or mx.distributed.init() N = group.size() + rank = group.rank() for layer in self.model.layers: # Shard the self attention - layer.self_attn.q_proj = shard_linear( - layer.self_attn.q_proj, "all-to-sharded", group=group - ) - layer.self_attn.k_proj = shard_linear( - layer.self_attn.k_proj, "all-to-sharded", group=group - ) - layer.self_attn.v_proj = shard_linear( - layer.self_attn.v_proj, "all-to-sharded", group=group + layer.self_attn.q_b_proj = shard_linear( + layer.self_attn.q_b_proj, "all-to-sharded", group=group ) layer.self_attn.o_proj = shard_linear( layer.self_attn.o_proj, "sharded-to-all", group=group ) - layer.self_attn.n_heads //= N - layer.self_attn.n_kv_heads //= N + layer.self_attn.num_heads //= N + num_heads = layer.self_attn.num_heads + sh = rank * num_heads + eh = sh + num_heads + + def shard_heads(w): + return w[sh:eh] + + layer.self_attn.embed_q.apply(shard_heads) + layer.self_attn.unembed_out.apply(shard_heads) # Shard the MLP if isinstance(layer.mlp, Glm4MoeLiteMLP): @@ -223,3 +271,6 @@ def predicate(k): return "e_score_correction_bias" not in k return predicate + + def make_cache(self): + return [CacheList(KVCache(), KVCache()) for _ in self.layers] From d3f54e8429aa514de0ed4c57755ef69799fe16c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?= <60228478+Goekdeniz-Guelmez@users.noreply.github.com> Date: Thu, 12 Feb 2026 11:00:12 +0100 Subject: [PATCH 6/9] Update acknowledgments to include contributions for GLM MoE DSA and additional architectures --- ACKNOWLEDGMENTS.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index f457e7b66..964053048 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -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`, `GLM5`, 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`. \ No newline at end of file + 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)` \ No newline at end of file From f792f6a4895fbeb01bdcc637022808deed54fa50 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 12 Feb 2026 07:13:24 -0800 Subject: [PATCH 7/9] use dsv32 for glm5 --- mlx_lm/models/deepseek_v32.py | 20 ++- mlx_lm/models/glm_moe_dsa.py | 276 ---------------------------------- mlx_lm/utils.py | 1 + 3 files changed, 14 insertions(+), 283 deletions(-) delete mode 100644 mlx_lm/models/glm_moe_dsa.py diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index edffef74e..0fe98302b 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -50,6 +50,7 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Dict = None attention_bias: bool = False + indexer_rope_interleave: bool = False class Indexer(nn.Module): @@ -71,7 +72,7 @@ def __init__(self, args: ModelArgs): self.rope = initialize_rope( dims=args.qk_rope_head_dim, base=args.rope_theta, - traditional=False, + traditional=self.indexer_rope_interleave, max_position_embeddings=args.max_position_embeddings, scaling_config=args.rope_scaling, ) @@ -495,6 +496,16 @@ def __call__( 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 +583,7 @@ def dequant(weight, scale_inv): 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() diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py deleted file mode 100644 index ef6813de1..000000000 --- a/mlx_lm/models/glm_moe_dsa.py +++ /dev/null @@ -1,276 +0,0 @@ -# Copyright © 2026 Apple Inc. - -from dataclasses import dataclass -from typing import Any, Dict, Optional - -import mlx.core as mx -import mlx.nn as nn -from mlx.nn.layers.distributed import shard_inplace, shard_linear - -from .base import BaseModelArgs, create_attention_mask -from .cache import CacheList, KVCache -from .deepseek_v32 import DeepseekV32Attention -from .glm4_moe_lite import Glm4MoeLiteMLP, Glm4MoeLiteMoE -from .pipeline import PipelineMixin - - -@dataclass -class ModelArgs(BaseModelArgs): - model_type: str = "glm4_moe_dsa" - vocab_size: int = 154880 - hidden_size: int = 2048 - intermediate_size: int = 10240 - moe_intermediate_size: int = 1536 - num_hidden_layers: int = 47 - num_attention_heads: int = 20 - num_key_value_heads: int = 20 - n_shared_experts: Optional[int] = 1 - n_routed_experts: Optional[int] = 64 - routed_scaling_factor: float = 1.8 - kv_lora_rank: int = 512 - q_lora_rank: int = 768 - qk_rope_head_dim: int = 64 - qk_nope_head_dim: int = 192 - v_head_dim: int = 256 - topk_method: str = "noaux_tc" - scoring_func: str = "sigmoid" - norm_topk_prob: bool = True - n_group: int = 1 - topk_group: int = 1 - num_experts_per_tok: int = 4 - moe_layer_freq: int = 1 - first_k_dense_replace: int = 1 - max_position_embeddings: int = 202752 - rms_norm_eps: float = 1e-5 - rope_theta: float = 1_000_000.0 - rope_scaling: Optional[Dict] = None - attention_bias: bool = False - attention_dropout: float = 0.0 - partial_rotary_factor: float = 1.0 - tie_word_embeddings: bool = False - num_nextn_predict_layers: int = 1 - index_head_dim: int = 128 - index_n_heads: int = 32 - index_topk: int = 2048 - quantization: Optional[Dict[str, Any]] = None - - -class Glm4MoeDSADecoderLayer(nn.Module): - def __init__(self, config: ModelArgs, layer_idx: int): - super().__init__() - self.self_attn = DeepseekV32Attention(config) - use_moe = ( - config.n_routed_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % config.moe_layer_freq == 0 - ) - self.mlp = Glm4MoeLiteMoE(config) if use_moe else Glm4MoeLiteMLP(config) - self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = nn.RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def __call__( - self, - x: mx.array, - mask: Optional[mx.array] = None, - cache: Optional[Any] = None, - ) -> mx.array: - r = self.self_attn(self.input_layernorm(x), mask, cache) - h = x + r - r = self.mlp(self.post_attention_layernorm(h)) - return h + r - - -class Glm4MoeDSAModel(PipelineMixin, nn.Module): - def __init__(self, config: ModelArgs): - super().__init__() - self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) - self.layers = [ - Glm4MoeDSADecoderLayer(config, idx) - for idx in range(config.num_hidden_layers) - ] - self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def __call__( - self, - x: mx.array, - cache: Optional[Any] = None, - ) -> mx.array: - h = self.embed_tokens(x) - - pipeline_rank = self.pipeline_rank - pipeline_size = self.pipeline_size - - if cache is None: - cache = [None] * len(self.pipeline_layers) - mask = create_attention_mask( - h, cache[0][0] if cache[0] else None, return_array=True - ) - - if pipeline_rank < pipeline_size - 1: - h = mx.distributed.recv_like(h, (pipeline_rank + 1)) - - for l, c in zip(self.pipeline_layers, cache): - h = l(h, mask, cache=c) - - if pipeline_rank != 0: - h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) - if cache[-1] is not None: - cache[-1][0].keys = mx.depends(cache[-1][0].keys, h) - - if pipeline_size > 1: - h = mx.distributed.all_gather(h)[: h.shape[0]] - - return self.norm(h) - - -class Model(nn.Module): - def __init__(self, config: ModelArgs): - super().__init__() - self.args = config - self.model_type = config.model_type - self.model = Glm4MoeDSAModel(config) - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - - def __call__( - self, - inputs: mx.array, - cache: Optional[Any] = None, - ) -> mx.array: - out = self.model(inputs, cache) - 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 - - # Stack experts and absorb MLA weights - for l in range(self.args.num_hidden_layers): - prefix = f"model.layers.{l}" - for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]: - for k in ["weight", "scales", "biases"]: - if f"{prefix}.mlp.experts.0.{m}.{k}" in weights: - to_join = [ - weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}") - for e in range(self.args.n_routed_experts) - ] - weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join) - - # MLA absorption: split kv_b_proj into embed_q and unembed_out - attn_prefix = f"{prefix}.self_attn" - if f"{attn_prefix}.kv_b_proj.weight" in weights: - quantized = f"{attn_prefix}.kv_b_proj.scales" in weights - v = weights.pop(f"{attn_prefix}.kv_b_proj.weight") - head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim - - if quantized: - dims = self.args.kv_lora_rank - scales = weights.pop(f"{attn_prefix}.kv_b_proj.scales") - biases = weights.pop(f"{attn_prefix}.kv_b_proj.biases") - bits = (v.shape[-1] * 32) // dims - group_size = dims // scales.shape[-1] - v = mx.dequantize( - v, scales, biases, bits=bits, group_size=group_size - ) - num_heads = self.args.num_attention_heads - v = v.reshape(num_heads, head_dim, -1) - wk = mx.contiguous( - v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2) - ) - wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :]) - if quantized: - wk, wk_scales, wk_biases = mx.quantize( - wk, bits=bits, group_size=group_size - ) - wv, wv_scales, wv_biases = mx.quantize( - wv, bits=bits, group_size=group_size - ) - weights[f"{attn_prefix}.embed_q.scales"] = wk_scales - weights[f"{attn_prefix}.unembed_out.scales"] = wv_scales - weights[f"{attn_prefix}.embed_q.biases"] = wk_biases - weights[f"{attn_prefix}.unembed_out.biases"] = wv_biases - weights[f"{attn_prefix}.embed_q.weight"] = wk - weights[f"{attn_prefix}.unembed_out.weight"] = wv - - return weights - - def shard(self, group: Optional[mx.distributed.Group] = None): - group = group or mx.distributed.init() - N = group.size() - rank = group.rank() - for layer in self.model.layers: - # Shard the self attention - layer.self_attn.q_b_proj = shard_linear( - layer.self_attn.q_b_proj, "all-to-sharded", group=group - ) - layer.self_attn.o_proj = shard_linear( - layer.self_attn.o_proj, "sharded-to-all", group=group - ) - layer.self_attn.num_heads //= N - num_heads = layer.self_attn.num_heads - sh = rank * num_heads - eh = sh + num_heads - - def shard_heads(w): - return w[sh:eh] - - layer.self_attn.embed_q.apply(shard_heads) - layer.self_attn.unembed_out.apply(shard_heads) - - # Shard the MLP - if isinstance(layer.mlp, Glm4MoeLiteMLP): - layer.mlp.gate_proj = shard_linear( - layer.mlp.gate_proj, "all-to-sharded", group=group - ) - layer.mlp.down_proj = shard_linear( - layer.mlp.down_proj, "sharded-to-all", group=group - ) - layer.mlp.up_proj = shard_linear( - layer.mlp.up_proj, "all-to-sharded", group=group - ) - - # Shard the MoE. Shard in place since the MoE should be responsible - # for aggregating the results. - else: - layer.mlp.sharding_group = group - shard_inplace( - layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group - ) - shard_inplace( - layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group - ) - shard_inplace( - layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group - ) - shard_inplace( - layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group - ) - shard_inplace( - layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group - ) - shard_inplace( - layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group - ) - - @property - def layers(self): - return self.model.pipeline_layers - - @property - def cast_predicate(self): - def predicate(k): - return "e_score_correction_bias" not in k - - return predicate - - def make_cache(self): - return [CacheList(KVCache(), KVCache()) for _ in self.layers] diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index f51fe95f5..2f9971d35 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -51,6 +51,7 @@ "qwen2_5_vl": "qwen2_vl", "minimax_m2": "minimax", "iquestcoder": "llama", + "glm_moe_dsa": "deepseek_v32", } MAX_FILE_SIZE_GB = 5 From 13e67d06d80bb0d255b388cb1f9e5c3350ac2493 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 12 Feb 2026 08:55:08 -0800 Subject: [PATCH 8/9] fix --- mlx_lm/models/deepseek_v32.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index 0fe98302b..97147257a 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -72,7 +72,7 @@ def __init__(self, args: ModelArgs): self.rope = initialize_rope( dims=args.qk_rope_head_dim, base=args.rope_theta, - traditional=self.indexer_rope_interleave, + traditional=args.indexer_rope_interleave, max_position_embeddings=args.max_position_embeddings, scaling_config=args.rope_scaling, ) From f921946dc630698a3f0cafd171f036c9661b9bde Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 12 Feb 2026 09:43:15 -0800 Subject: [PATCH 9/9] Fix rope theta --- mlx_lm/models/deepseek_v32.py | 3 +- mlx_lm/models/glm_moe_dsa.py | 53 +++++++++++++++++++++++++++++++++++ mlx_lm/utils.py | 1 - 3 files changed, 54 insertions(+), 3 deletions(-) create mode 100644 mlx_lm/models/glm_moe_dsa.py diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index 97147257a..e40e52950 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -50,7 +50,6 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Dict = None attention_bias: bool = False - indexer_rope_interleave: bool = False class Indexer(nn.Module): @@ -72,7 +71,7 @@ def __init__(self, args: ModelArgs): self.rope = initialize_rope( dims=args.qk_rope_head_dim, base=args.rope_theta, - traditional=args.indexer_rope_interleave, + traditional=True, max_position_embeddings=args.max_position_embeddings, scaling_config=args.rope_scaling, ) diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py new file mode 100644 index 000000000..14e96e365 --- /dev/null +++ b/mlx_lm/models/glm_moe_dsa.py @@ -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) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 2f9971d35..f51fe95f5 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -51,7 +51,6 @@ "qwen2_5_vl": "qwen2_vl", "minimax_m2": "minimax", "iquestcoder": "llama", - "glm_moe_dsa": "deepseek_v32", } MAX_FILE_SIZE_GB = 5