Add GLM5 (#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>
This commit is contained in:
+6
-2
@@ -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)`
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user