diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index b452d5b..9640530 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`, `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 diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index edffef7..e40e529 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -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() diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py new file mode 100644 index 0000000..14e96e3 --- /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)