From 26ff2a3a618c66227bce47f182a66244d04e518e Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:30:21 +0100 Subject: [PATCH 01/91] feat: add optional rope_parameters to ModelArgs for dynamic rope_theta configuration --- mlx_lm/models/afm7.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/afm7.py b/mlx_lm/models/afm7.py index 8a87b4509..81123ec77 100644 --- a/mlx_lm/models/afm7.py +++ b/mlx_lm/models/afm7.py @@ -25,9 +25,14 @@ class ModelArgs(BaseModelArgs): num_heads: int num_kv_heads: int hidden_dim_scale_factor: float = 3.25 + rope_parameters: Optional[Dict] = None rope_theta: float = 50000 rms_norm_eps: float = 1e-5 + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 50000.0) + class FusedLoRALinear(nn.Module): def __init__( From edd1499e71d71ab6fccbfae3ddc234f625bedee5 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:30:55 +0100 Subject: [PATCH 02/91] feat: add rope_parameters to ModelArgs for dynamic rope_theta configuration --- mlx_lm/models/youtu_llm.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/youtu_llm.py b/mlx_lm/models/youtu_llm.py index 9d8c653b1..0cf41a615 100644 --- a/mlx_lm/models/youtu_llm.py +++ b/mlx_lm/models/youtu_llm.py @@ -33,6 +33,11 @@ class ModelArgs(BaseModelArgs): attention_bias: bool = False mlp_bias: bool = False tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) class YoutuLLMAttention(nn.Module): From 47489ee8ca9b70152812c969b1557a99315682bc Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:31:10 +0100 Subject: [PATCH 03/91] feat: initialize rope_theta from rope_parameters in ModelArgs --- mlx_lm/models/telechat3.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/telechat3.py b/mlx_lm/models/telechat3.py index 06a2d653a..18e0fd7f1 100644 --- a/mlx_lm/models/telechat3.py +++ b/mlx_lm/models/telechat3.py @@ -28,6 +28,11 @@ class ModelArgs(BaseModelArgs): head_dim: Optional[int] = None rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) class Telechat3Attention(nn.Module): From 47be5ed832c3e5bec4de8a4bfde497308afc9e11 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:31:24 +0100 Subject: [PATCH 04/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/step3p5.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/step3p5.py b/mlx_lm/models/step3p5.py index 852823ea8..d5615bb2f 100644 --- a/mlx_lm/models/step3p5.py +++ b/mlx_lm/models/step3p5.py @@ -61,6 +61,11 @@ class ModelArgs(BaseModelArgs): swiglu_limits: Optional[List[float]] = None swiglu_limits_shared: Optional[List[float]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) class ZeroCenteredRMSNorm(nn.Module): From 069a859db8d05727dc18f13cdb0c1b54c1285c92 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:31:54 +0100 Subject: [PATCH 05/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/starcoder2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/starcoder2.py b/mlx_lm/models/starcoder2.py index 4cbed554f..e0c8d9bac 100644 --- a/mlx_lm/models/starcoder2.py +++ b/mlx_lm/models/starcoder2.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -21,6 +21,11 @@ class ModelArgs(BaseModelArgs): vocab_size: int = 49152 rope_theta: float = 100000 tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From 61cce1d1bd3e51fcd9c966c05ccb9439e3ad66d0 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:32:42 +0100 Subject: [PATCH 06/91] feat: update ModelArgs to set rope_theta from rope_parameters in post_init --- mlx_lm/models/stablelm.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/stablelm.py b/mlx_lm/models/stablelm.py index 0f90f721b..ab4299110 100644 --- a/mlx_lm/models/stablelm.py +++ b/mlx_lm/models/stablelm.py @@ -1,7 +1,8 @@ -# Copyright © 2023-2024 Apple Inc. +# Copyright © 2026 Apple Inc. import math from dataclasses import dataclass +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -25,6 +26,11 @@ class ModelArgs(BaseModelArgs): layer_norm_eps: float use_parallel_residual: bool = False qk_layernorm: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class LayerNormPerHead(nn.Module): From fe4c5c25a8ab30a6d8d6a12d4aef972dd5181f32 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:33:21 +0100 Subject: [PATCH 07/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/solar_open.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/solar_open.py b/mlx_lm/models/solar_open.py index a0f8b8b5d..a3ed331d3 100644 --- a/mlx_lm/models/solar_open.py +++ b/mlx_lm/models/solar_open.py @@ -36,3 +36,8 @@ class ModelArgs(BaseModelArgs): topk_group: int = 1 scoring_func: str = "sigmoid" topk_method: str = "noaux_tc" + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) From 38fcf6d04228bca4c929608fff82ef11eb2d174e Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:33:57 +0100 Subject: [PATCH 08/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/smollm3.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/smollm3.py b/mlx_lm/models/smollm3.py index 030924c7e..2e4f8f641 100644 --- a/mlx_lm/models/smollm3.py +++ b/mlx_lm/models/smollm3.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Optional +from typing import Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -14,9 +14,12 @@ class ModelArgs(llama.ModelArgs): model_type: str no_rope_layer_interval: int = 4 no_rope_layers: Optional[list[int]] = None + rope_parameters: Optional[Dict] = None def __post_init__(self): super().__post_init__() + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) if self.no_rope_layers is None: self.no_rope_layers = [ int((i + 1) % self.no_rope_layer_interval != 0) From ac2ec74356401739fe078db1940d49321e4cf6bb Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:34:24 +0100 Subject: [PATCH 09/91] feat: update ModelArgs to set rope_theta from rope_parameters in post_init --- mlx_lm/models/seed_oss.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/seed_oss.py b/mlx_lm/models/seed_oss.py index d23c9f30b..f4fe23cae 100644 --- a/mlx_lm/models/seed_oss.py +++ b/mlx_lm/models/seed_oss.py @@ -30,6 +30,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From 9aa4637cc66e3d6d7aa76a7e275407f066e8e76c Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:34:38 +0100 Subject: [PATCH 10/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/rwkv7.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/rwkv7.py b/mlx_lm/models/rwkv7.py index d698171ed..6f0ecdaea 100644 --- a/mlx_lm/models/rwkv7.py +++ b/mlx_lm/models/rwkv7.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from functools import partial -from typing import Optional +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -25,6 +25,11 @@ class ModelArgs(BaseModelArgs): gate_low_rank_dim: int decay_low_rank_dim: int tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) @partial(mx.compile, shapeless=True) From 8736c9e4a1a4e60c8bdbe88cf87fa96131fb9e7d Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:35:04 +0100 Subject: [PATCH 11/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/recurrent_gemma.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/recurrent_gemma.py b/mlx_lm/models/recurrent_gemma.py index 4659d6c3d..50ab4d330 100644 --- a/mlx_lm/models/recurrent_gemma.py +++ b/mlx_lm/models/recurrent_gemma.py @@ -1,5 +1,6 @@ # Copyright © 2023-2024 Apple Inc. +from ast import Dict import math from dataclasses import dataclass from typing import List, Optional @@ -29,8 +30,11 @@ class ModelArgs(BaseModelArgs): embeddings_scale_by_sqrt_dim: bool = True block_types: Optional[List[str]] = None _block_types: Optional[List[str]] = None + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) # For some reason these have different names in 2B and 9B if self.block_types is None: self.block_types = self._block_types From 144d58ced2d7392fe5acc7fb4daeb39a3fcf5a1f Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:35:17 +0100 Subject: [PATCH 12/91] feat: add post_init method to ModelArgs to initialize rope_theta from rope_parameters --- mlx_lm/models/qwen3.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/qwen3.py b/mlx_lm/models/qwen3.py index a59343b68..336134ed1 100644 --- a/mlx_lm/models/qwen3.py +++ b/mlx_lm/models/qwen3.py @@ -27,6 +27,11 @@ class ModelArgs(BaseModelArgs): head_dim: int tie_word_embeddings: bool rope_scaling: Optional[Dict[str, Union[float, str]]] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From e8e2dd46b26d157a93eff7488644ac4bdceae354 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:35:59 +0100 Subject: [PATCH 13/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/qwen3_next.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/qwen3_next.py b/mlx_lm/models/qwen3_next.py index 47c9971bd..56710a713 100644 --- a/mlx_lm/models/qwen3_next.py +++ b/mlx_lm/models/qwen3_next.py @@ -51,6 +51,11 @@ class ModelArgs(BaseModelArgs): attention_bias: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None full_attention_interval: int = 4 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Qwen3NextRMSNormGated(nn.Module): From 353862551a94f1cd363974dc7f2225da6ccff181 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:36:10 +0100 Subject: [PATCH 14/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/qwen3_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/qwen3_moe.py b/mlx_lm/models/qwen3_moe.py index 52dc50f9b..5202e6f72 100644 --- a/mlx_lm/models/qwen3_moe.py +++ b/mlx_lm/models/qwen3_moe.py @@ -32,6 +32,11 @@ class ModelArgs(BaseModelArgs): max_position_embeddings: int norm_topk_prob: bool rope_scaling: Optional[Dict[str, Union[float, str]]] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From 6f33002ffd7596573dbfbc74fb92c8425e980c60 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:36:41 +0100 Subject: [PATCH 15/91] feat: add rope_parameters to ModelArgs and implement post_init to initialize rope_theta --- mlx_lm/models/qwen2.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/qwen2.py b/mlx_lm/models/qwen2.py index 45df85c7b..8987274c3 100644 --- a/mlx_lm/models/qwen2.py +++ b/mlx_lm/models/qwen2.py @@ -27,6 +27,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) class Attention(nn.Module): From 438b06d7585a26b55e304bc79d72aea5efa0cdda Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:37:01 +0100 Subject: [PATCH 16/91] feat: add rope_parameters to ModelArgs and update rope_theta in post_init --- mlx_lm/models/qwen2_moe.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/qwen2_moe.py b/mlx_lm/models/qwen2_moe.py index 8d6022d2e..dd2708af4 100644 --- a/mlx_lm/models/qwen2_moe.py +++ b/mlx_lm/models/qwen2_moe.py @@ -29,8 +29,12 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From 24a82bf595e47dcd676c9c5fa49aff2d9c853184 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:37:29 +0100 Subject: [PATCH 17/91] feat: initialize rope_theta in post_init if rope_parameters are provided --- mlx_lm/models/qwen.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/qwen.py b/mlx_lm/models/qwen.py index 96e1a9718..a57566ae6 100644 --- a/mlx_lm/models/qwen.py +++ b/mlx_lm/models/qwen.py @@ -4,6 +4,7 @@ import mlx.core as mx import mlx.nn as nn +from typing import Optional, Dict from .activations import swiglu from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention @@ -22,8 +23,11 @@ class ModelArgs(BaseModelArgs): no_bias: bool = True vocab_size: int = 151936 num_key_value_heads = None + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From 50c0a190d53f9b5f1a765cddb7921a1c7dd56cc8 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:37:50 +0100 Subject: [PATCH 18/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/plamo2.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/plamo2.py b/mlx_lm/models/plamo2.py index b835de0a5..0f0d18a47 100644 --- a/mlx_lm/models/plamo2.py +++ b/mlx_lm/models/plamo2.py @@ -1,8 +1,7 @@ # Copyright © 2025 Apple Inc. -import math from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -35,6 +34,11 @@ class ModelArgs(BaseModelArgs): mamba_enabled: bool = True intermediate_size: int = 13312 vocab_size: int = 32000 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class RMSNorm(nn.Module): From 28ab8189234e77cffe281dde7b5d732187b963c1 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:38:15 +0100 Subject: [PATCH 19/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/plamo.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/plamo.py b/mlx_lm/models/plamo.py index 762a3c588..4bbd65f7c 100644 --- a/mlx_lm/models/plamo.py +++ b/mlx_lm/models/plamo.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -23,6 +23,11 @@ class ModelArgs(BaseModelArgs): n_shared_head: int = 8 rope_theta: float = 10000 rope_traditional: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From bbeeb1f7d0fdae5b027f5de242ee0b7f000155c2 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:40:25 +0100 Subject: [PATCH 20/91] refactor: update ModelArgs to inherit from BaseModelArgs and clean up unused method --- mlx_lm/models/phixtral.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/mlx_lm/models/phixtral.py b/mlx_lm/models/phixtral.py index 5477c2c06..bb4768e31 100644 --- a/mlx_lm/models/phixtral.py +++ b/mlx_lm/models/phixtral.py @@ -1,19 +1,19 @@ -# Copyright © 2023-2024 Apple Inc. +# Copyright © 2026 Apple Inc. import inspect import math from dataclasses import dataclass -from typing import Tuple +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn -from .base import create_attention_mask, scaled_dot_product_attention +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention from .switch_layers import SwitchMLP @dataclass -class ModelArgs: +class ModelArgs(BaseModelArgs): model_type: str num_vocab: int = 51200 model_dim: int = 2560 @@ -23,16 +23,6 @@ class ModelArgs: num_experts_per_tok: int = 2 num_local_experts: int = 4 - @classmethod - def from_dict(cls, params): - return cls( - **{ - k: v - for k, v in params.items() - if k in inspect.signature(cls).parameters - } - ) - class RoPEAttention(nn.Module): def __init__(self, dims: int, num_heads: int, rotary_dim: int): From 6b2c2f7ff8f0013ad7fe7d6c963e0cf647748014 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:40:40 +0100 Subject: [PATCH 21/91] feat: update ModelArgs to set rope_theta from rope_parameters in post_init --- mlx_lm/models/phimoe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/phimoe.py b/mlx_lm/models/phimoe.py index c192d603b..2ff6c6afa 100644 --- a/mlx_lm/models/phimoe.py +++ b/mlx_lm/models/phimoe.py @@ -27,6 +27,11 @@ class ModelArgs(BaseModelArgs): num_local_experts: int = 16 num_experts_per_tok: int = 2 rope_theta: float = 10000.0 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 33a8a412303833ad5dc26fd244e45afbd07ccb1d Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:41:30 +0100 Subject: [PATCH 22/91] feat: initialize rope_theta from rope_parameters in post_init if provided --- mlx_lm/models/phi3.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/phi3.py b/mlx_lm/models/phi3.py index b3579eae8..779dc3ddb 100644 --- a/mlx_lm/models/phi3.py +++ b/mlx_lm/models/phi3.py @@ -28,8 +28,12 @@ class ModelArgs(BaseModelArgs): max_position_embeddings: int = 131072 original_max_position_embeddings: int = 4096 tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From dcd28a7e0dd97bfddc87c90b203ff214783f0219 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:42:12 +0100 Subject: [PATCH 23/91] feat: update ModelArgs to allow dynamic initialization of rope_theta from rope_parameters --- mlx_lm/models/phi.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/phi.py b/mlx_lm/models/phi.py index cbf07ec51..1d711b5f0 100644 --- a/mlx_lm/models/phi.py +++ b/mlx_lm/models/phi.py @@ -1,7 +1,8 @@ -# Copyright © 2023-2024 Apple Inc. +# Copyright © 2026 Apple Inc. import math from dataclasses import dataclass +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -11,7 +12,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "phi" + model_type: str max_position_embeddings: int = 2048 vocab_size: int = 51200 hidden_size: int = 2560 @@ -22,8 +23,12 @@ class ModelArgs(BaseModelArgs): intermediate_size: int = 10240 layer_norm_eps: float = 1e-5 rope_theta: float = 10000.0 + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From 59b803fe7ede7c7b4168e80564a271f570e51d20 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:42:52 +0100 Subject: [PATCH 24/91] feat: add rope_parameters to ModelArgs for dynamic rope_theta initialization --- mlx_lm/models/olmoe.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/olmoe.py b/mlx_lm/models/olmoe.py index c971c2247..e994a8e2c 100644 --- a/mlx_lm/models/olmoe.py +++ b/mlx_lm/models/olmoe.py @@ -32,8 +32,12 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From f9531e8f89f01beaae7943438b3ad9a47904b166 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:43:11 +0100 Subject: [PATCH 25/91] feat: add rope_parameters to ModelArgs for dynamic rope_theta assignment in post_init --- mlx_lm/models/olmo3.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlx_lm/models/olmo3.py b/mlx_lm/models/olmo3.py index f87a818d8..3f5176751 100644 --- a/mlx_lm/models/olmo3.py +++ b/mlx_lm/models/olmo3.py @@ -30,8 +30,11 @@ class ModelArgs(BaseModelArgs): head_dim: Optional[int] = None rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads if self.layer_types is None: From 595278a30b37aa60987f562ac3ceb7070e8d12ec Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:43:32 +0100 Subject: [PATCH 26/91] feat: initialize rope_theta from rope_parameters in post_init if provided --- mlx_lm/models/olmo2.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlx_lm/models/olmo2.py b/mlx_lm/models/olmo2.py index 4591e00dc..002464296 100644 --- a/mlx_lm/models/olmo2.py +++ b/mlx_lm/models/olmo2.py @@ -29,8 +29,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From 3f5f6be5da0746b290a7acb8af526aebb13ca1cc Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:44:29 +0100 Subject: [PATCH 27/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/olmo.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/olmo.py b/mlx_lm/models/olmo.py index bb7592d24..8e13a06df 100644 --- a/mlx_lm/models/olmo.py +++ b/mlx_lm/models/olmo.py @@ -2,7 +2,7 @@ import sys from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -30,8 +30,12 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False mlp_ratio: int = 4 weight_tying: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.mlp_hidden_size = ( self.mlp_hidden_size if self.mlp_hidden_size is not None From 73b4f733d3afd84aa3119b9735f6177db0cd055f Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:44:46 +0100 Subject: [PATCH 28/91] feat: add rope_parameters to ModelArgs for dynamic rope_theta assignment in post_init --- mlx_lm/models/nemotron.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlx_lm/models/nemotron.py b/mlx_lm/models/nemotron.py index 5d612d477..349b97d18 100644 --- a/mlx_lm/models/nemotron.py +++ b/mlx_lm/models/nemotron.py @@ -30,8 +30,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) if self.rope_scaling: if not "factor" in self.rope_scaling: raise ValueError(f"rope_scaling must contain 'factor'") From 91a2dd2f0e57f8a758df612a5d7f5b434457ccf6 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:45:21 +0100 Subject: [PATCH 29/91] feat: add rope_parameters to ModelArgs for dynamic rope_theta assignment in post_init --- mlx_lm/models/nemotron-nas.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/nemotron-nas.py b/mlx_lm/models/nemotron-nas.py index 09087c8e5..504c5affd 100644 --- a/mlx_lm/models/nemotron-nas.py +++ b/mlx_lm/models/nemotron-nas.py @@ -114,8 +114,12 @@ class ModelArgs(BaseModelArgs): rope_scaling: Optional[Dict[str, Union[float, str]]] = None max_position_embeddings: int = 131072 tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) + # Automatically parse block_configs if they are loaded as dicts if self.block_configs and isinstance(self.block_configs[0], dict): self.block_configs = [ From ac1cd067ce701006fd5edad24ca2d1d3fecae1bc Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:46:00 +0100 Subject: [PATCH 30/91] refactor: remove unused import of partial in nemotron_h.py --- mlx_lm/models/nemotron_h.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mlx_lm/models/nemotron_h.py b/mlx_lm/models/nemotron_h.py index 5fb500eea..8dd69f98e 100644 --- a/mlx_lm/models/nemotron_h.py +++ b/mlx_lm/models/nemotron_h.py @@ -1,7 +1,6 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from functools import partial from typing import Any, List, Optional, Tuple import mlx.core as mx From 05bc8dc7799d0b2ef242c40236fb62705ba8ac32 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:46:26 +0100 Subject: [PATCH 31/91] feat: update ModelArgs to include rope_parameters and initialize rope_theta in post_init --- mlx_lm/models/nanochat.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/nanochat.py b/mlx_lm/models/nanochat.py index 1229e6ff9..8d2c1c898 100644 --- a/mlx_lm/models/nanochat.py +++ b/mlx_lm/models/nanochat.py @@ -3,7 +3,7 @@ import math from dataclasses import dataclass from functools import partial -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -13,7 +13,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "nanochat" + model_type: str hidden_size: int = 1280 num_hidden_layers: int = 20 num_attention_heads: int = 10 @@ -22,6 +22,11 @@ class ModelArgs(BaseModelArgs): max_position_embeddings: int = 2048 intermediate_size: int = 5120 # 4 * hidden_size rope_theta: float = 10000.0 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) def rms_norm(x): From ae75aa9f049ea774148d44d18bd33ccc70a65b05 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:46:52 +0100 Subject: [PATCH 32/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/mixtral.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlx_lm/models/mixtral.py b/mlx_lm/models/mixtral.py index 18d16644c..e297f57fa 100644 --- a/mlx_lm/models/mixtral.py +++ b/mlx_lm/models/mixtral.py @@ -26,8 +26,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From 141e66b9da1a26de176d68685032bb003ec4ceb2 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:47:40 +0100 Subject: [PATCH 33/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/minimax.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/minimax.py b/mlx_lm/models/minimax.py index 9bf78d9a4..fd116b218 100644 --- a/mlx_lm/models/minimax.py +++ b/mlx_lm/models/minimax.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from functools import lru_cache -from typing import Any, List, Optional +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -32,6 +32,11 @@ class ModelArgs(BaseModelArgs): scoring_func: str = "sigmoid" head_dim: Optional[int] = None use_qk_norm: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) @lru_cache From 427b7b1ce31e732d361f3f3be52363daf9526e32 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:47:57 +0100 Subject: [PATCH 34/91] feat: add post_init method to ModelArgs for dynamic rope_theta initialization from rope_parameters --- mlx_lm/models/minicpm3.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/minicpm3.py b/mlx_lm/models/minicpm3.py index 977f00eae..34141f932 100644 --- a/mlx_lm/models/minicpm3.py +++ b/mlx_lm/models/minicpm3.py @@ -34,6 +34,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[str, float]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) class Attention(nn.Module): From f6a619f8f0d98b064f64d2ff52e422dbec8895bf Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:48:13 +0100 Subject: [PATCH 35/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/minicpm.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/minicpm.py b/mlx_lm/models/minicpm.py index 660db9051..2d6d96418 100644 --- a/mlx_lm/models/minicpm.py +++ b/mlx_lm/models/minicpm.py @@ -29,6 +29,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[str, float]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) class MLP(nn.Module): From 287fa617d1a97a9030d9b88e56d7fb01ac294378 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:49:32 +0100 Subject: [PATCH 36/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/mimo.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/mimo.py b/mlx_lm/models/mimo.py index d87a2d94f..4ef9891ff 100644 --- a/mlx_lm/models/mimo.py +++ b/mlx_lm/models/mimo.py @@ -27,6 +27,11 @@ class ModelArgs(BaseModelArgs): rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False num_nextn_predict_layers: int = 2 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From bb6662844da6597a37c83e45a544745e684f7c24 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:49:55 +0100 Subject: [PATCH 37/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/mimo_v2_flash.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/mlx_lm/models/mimo_v2_flash.py b/mlx_lm/models/mimo_v2_flash.py index bec873e17..11ccc1e06 100644 --- a/mlx_lm/models/mimo_v2_flash.py +++ b/mlx_lm/models/mimo_v2_flash.py @@ -1,8 +1,6 @@ -# Copyright © 2024 Apple Inc. +# Copyright © 2026 Apple Inc. -import math from dataclasses import dataclass -from functools import partial from typing import Any, Dict, List, Optional import mlx.core as mx @@ -49,6 +47,11 @@ class ModelArgs(BaseModelArgs): swa_head_dim: int swa_v_head_dim: int partial_rotary_factor: int + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From 8d68b311b20e9b93a7e20715ebcfe69ba0f6ca14 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:50:26 +0100 Subject: [PATCH 38/91] feat: add rope_parameters to ModelArgs and implement post_init for rope_theta initialization --- mlx_lm/models/longcat_flash.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/longcat_flash.py b/mlx_lm/models/longcat_flash.py index fb3050a32..cf349058d 100644 --- a/mlx_lm/models/longcat_flash.py +++ b/mlx_lm/models/longcat_flash.py @@ -43,6 +43,11 @@ class ModelArgs(BaseModelArgs): norm_topk_prob: bool = False router_bias: bool = False rope_scaling: Optional[Dict] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class LongcatFlashMLA(nn.Module): From 160b9d11280b4cf8ccb8c29d18e5410216c1bed0 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:50:39 +0100 Subject: [PATCH 39/91] feat: add rope_parameters to ModelArgs and implement post_init for rope_theta initialization --- mlx_lm/models/longcat_flash_ngram.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/longcat_flash_ngram.py b/mlx_lm/models/longcat_flash_ngram.py index 9c0f2a8f7..4e2e94829 100644 --- a/mlx_lm/models/longcat_flash_ngram.py +++ b/mlx_lm/models/longcat_flash_ngram.py @@ -43,6 +43,11 @@ class ModelArgs(BaseModelArgs): norm_topk_prob: bool = False router_bias: bool = False rope_scaling: Optional[Dict] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class NgramEmbedding(nn.Module): From 20ed8eb8b5acccecfb1f11fb37747137c0a85f95 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:50:57 +0100 Subject: [PATCH 40/91] feat: add rope_parameters to TextArgs and implement post_init for rope_theta initialization --- mlx_lm/models/llama4.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/llama4.py b/mlx_lm/models/llama4.py index e4e284d71..f39d222de 100644 --- a/mlx_lm/models/llama4.py +++ b/mlx_lm/models/llama4.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Union +from typing import Any, Optional, Union, Dict import mlx.core as mx import mlx.nn as nn @@ -37,6 +37,11 @@ class TextArgs(BaseModelArgs): attn_temperature_tuning: int = 4 floor_scale: int = 8192 attn_scale: float = 0.1 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) @dataclass From 3701dd79b15e72ee3f693408e3e7e8774de3bca5 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:51:13 +0100 Subject: [PATCH 41/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/llama4_text.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/llama4_text.py b/mlx_lm/models/llama4_text.py index 6f45a8557..3301ea94c 100644 --- a/mlx_lm/models/llama4_text.py +++ b/mlx_lm/models/llama4_text.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -26,6 +26,11 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool no_rope_layers: list use_qk_norm: bool + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Attention(nn.Module): From faa303a2384a2a4e61c5db4cde13535259535f12 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:51:47 +0100 Subject: [PATCH 42/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/llama.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/llama.py b/mlx_lm/models/llama.py index 826bc9fc0..c12ddbe31 100644 --- a/mlx_lm/models/llama.py +++ b/mlx_lm/models/llama.py @@ -33,8 +33,12 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool = True layer_types: Optional[List[str]] = None sliding_window: Optional[int] = None + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From de952792b03c878d1d5340da2049957751e27a4d Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:52:08 +0100 Subject: [PATCH 43/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/lille-130m.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/lille-130m.py b/mlx_lm/models/lille-130m.py index b79057216..775eb59ea 100644 --- a/mlx_lm/models/lille-130m.py +++ b/mlx_lm/models/lille-130m.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -22,6 +22,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float vocab_size: int tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class Lille130mAttention(nn.Module): From 99132f110ad57b27a0617fbbbd744fe720ba8f56 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:52:31 +0100 Subject: [PATCH 44/91] feat: update ModelArgs to include rope_parameters as a Dict and initialize rope_theta in post_init --- mlx_lm/models/lfm2.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/lfm2.py b/mlx_lm/models/lfm2.py index c7f742af9..52df0e5c0 100644 --- a/mlx_lm/models/lfm2.py +++ b/mlx_lm/models/lfm2.py @@ -1,6 +1,6 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -36,8 +36,11 @@ class ModelArgs(BaseModelArgs): rope_parameters: Optional[dict] = None full_attn_idxs: Optional[List[int]] = None layer_types: Optional[List[str]] = None + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) if self.rope_parameters is not None and "rope_theta" in self.rope_parameters: self.rope_theta = self.rope_parameters["rope_theta"] if self.num_key_value_heads is None: From d1bb0a40db15be222ae8f51165540684b0666b97 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:53:18 +0100 Subject: [PATCH 45/91] feat: update rope_theta initialization to use default value if not provided in rope_parameters --- mlx_lm/models/lfm2_moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/lfm2_moe.py b/mlx_lm/models/lfm2_moe.py index 39ea433e6..dc7ea42d6 100644 --- a/mlx_lm/models/lfm2_moe.py +++ b/mlx_lm/models/lfm2_moe.py @@ -41,8 +41,8 @@ class ModelArgs(BaseModelArgs): layer_types: Optional[List[str]] = None def __post_init__(self): - if self.rope_parameters is not None and "rope_theta" in self.rope_parameters: - self.rope_theta = self.rope_parameters["rope_theta"] + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) if self.full_attn_idxs is None: self.full_attn_idxs = [ i From 58a8dd346ae5e88821116c99dd1e85e6031f9c7f Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:53:31 +0100 Subject: [PATCH 46/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/Klear.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/Klear.py b/mlx_lm/models/Klear.py index fb7a56029..b51c06d74 100644 --- a/mlx_lm/models/Klear.py +++ b/mlx_lm/models/Klear.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -31,6 +31,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float max_position_embeddings: int norm_topk_prob: bool + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class KlearAttention(nn.Module): From 4857d118effe41ef0aabc7513a8dab292a1a72c0 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:54:39 +0100 Subject: [PATCH 47/91] feat: initialize rope_theta from rope_parameters in TextArgs post_init --- mlx_lm/models/kimi_vl.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/kimi_vl.py b/mlx_lm/models/kimi_vl.py index bc27540fd..9cfa7a7fe 100644 --- a/mlx_lm/models/kimi_vl.py +++ b/mlx_lm/models/kimi_vl.py @@ -40,6 +40,11 @@ class TextArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Dict = None attention_bias: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) @dataclass From 38ff81741389e52f0aa70d55e32e7d1c85ad9cf4 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:55:29 +0100 Subject: [PATCH 48/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/iquestloopcoder.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/iquestloopcoder.py b/mlx_lm/models/iquestloopcoder.py index 5fc387f2c..6b4fc8cf0 100644 --- a/mlx_lm/models/iquestloopcoder.py +++ b/mlx_lm/models/iquestloopcoder.py @@ -51,6 +51,11 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool = False loop_num: int = 2 loop_window_size: int = 64 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) class LoopGateProjection(nn.Module): From 0f65ca42ea28b7b02d5bc4c322dfb472f639f8f5 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:55:56 +0100 Subject: [PATCH 49/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/internlm3.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx_lm/models/internlm3.py b/mlx_lm/models/internlm3.py index 766fe977f..fbb16e49a 100644 --- a/mlx_lm/models/internlm3.py +++ b/mlx_lm/models/internlm3.py @@ -27,8 +27,12 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads From fca0729e9dbac788d80a6135569acb9f1352a52e Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:56:16 +0100 Subject: [PATCH 50/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/internlm2.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/internlm2.py b/mlx_lm/models/internlm2.py index 8fd2301dd..f97b8bce2 100644 --- a/mlx_lm/models/internlm2.py +++ b/mlx_lm/models/internlm2.py @@ -26,8 +26,12 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads @@ -38,7 +42,7 @@ def __post_init__(self): if self.rope_scaling["type"] not in ["linear", "dynamic"]: raise ValueError( - "rope_scaling 'type' currently only supports 'linear' or 'dynamic" + "rope_scaling 'type' currently only supports 'linear' or 'dynamic'" ) From e8e509c250ccd08891111bf50172d347e2bbc36c Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:56:42 +0100 Subject: [PATCH 51/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/hunyuan.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlx_lm/models/hunyuan.py b/mlx_lm/models/hunyuan.py index 0ee3a3a07..ff6004c37 100644 --- a/mlx_lm/models/hunyuan.py +++ b/mlx_lm/models/hunyuan.py @@ -33,8 +33,11 @@ class ModelArgs(BaseModelArgs): moe_intermediate_size: Optional[Union[int, list]] = None rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) if self.rope_scaling: required_keys = {"factor", "type"} From b0c22bb0308382b8e81644d1aa6a6a13de9273aa Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:56:58 +0100 Subject: [PATCH 52/91] feat: add rope_parameters to ModelArgs and update rope_theta in post_init --- mlx_lm/models/hunyuan_v1_dense.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/hunyuan_v1_dense.py b/mlx_lm/models/hunyuan_v1_dense.py index d94edeaaf..547373929 100644 --- a/mlx_lm/models/hunyuan_v1_dense.py +++ b/mlx_lm/models/hunyuan_v1_dense.py @@ -27,6 +27,11 @@ class ModelArgs(BaseModelArgs): rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False head_dim: Optional[int] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) def __post_init__(self): if self.rope_scaling: From bc7234e87240293da2c7c68bb9494c3bdf1feecc Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:57:27 +0100 Subject: [PATCH 53/91] feat: update ModelArgs to include rope_parameters and adjust rope_theta in post_init --- mlx_lm/models/helium.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/helium.py b/mlx_lm/models/helium.py index 238f205fa..c60a680d1 100644 --- a/mlx_lm/models/helium.py +++ b/mlx_lm/models/helium.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -12,6 +12,7 @@ @dataclass class ModelArgs(BaseModelArgs): + model_type: str hidden_size: int num_hidden_layers: int intermediate_size: int @@ -23,9 +24,13 @@ class ModelArgs(BaseModelArgs): head_dim: int max_position_embeddings: int mlp_bias: bool - model_type: str rope_theta: float tie_word_embeddings: bool + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class HeliumAttention(nn.Module): From 0be2cc09dc2ffbce699310c36fcc556e6f433e85 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:58:07 +0100 Subject: [PATCH 54/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/granitemoehybrid.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/granitemoehybrid.py b/mlx_lm/models/granitemoehybrid.py index a95c887e2..fd55e66b8 100644 --- a/mlx_lm/models/granitemoehybrid.py +++ b/mlx_lm/models/granitemoehybrid.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional, Tuple +from typing import Any, List, Optional, Tuple, Dict import mlx.core as mx import mlx.nn as nn @@ -62,6 +62,12 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool = True time_step_limit: Tuple[float, float] = (0.001, 100.0) + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + # Mode flags - inferred from num_local_experts @property def use_moe(self) -> bool: From b552e904f246b7a6712a93ce3c60bca4b071e2da Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:58:18 +0100 Subject: [PATCH 55/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/granitemoe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/granitemoe.py b/mlx_lm/models/granitemoe.py index e77cc6601..0cb30c7d2 100644 --- a/mlx_lm/models/granitemoe.py +++ b/mlx_lm/models/granitemoe.py @@ -32,6 +32,11 @@ class ModelArgs(BaseModelArgs): num_experts_per_tok: int rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class GraniteMoeAttention(nn.Module): From 962546d56fcb631e612ed786bac66231169c5257 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:58:28 +0100 Subject: [PATCH 56/91] feat: initialize rope_theta in ModelArgs post_init based on rope_parameters --- mlx_lm/models/granite.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/granite.py b/mlx_lm/models/granite.py index 6892445d7..f0eee1305 100644 --- a/mlx_lm/models/granite.py +++ b/mlx_lm/models/granite.py @@ -31,6 +31,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 194d852a27caa4d542d57847039dd418e10aac86 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:59:15 +0100 Subject: [PATCH 57/91] feat: update ModelArgs to require model_type and add rope_parameters with post_init handling --- mlx_lm/models/gpt_oss.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/gpt_oss.py b/mlx_lm/models/gpt_oss.py index e0fda5e27..9eb5659de 100644 --- a/mlx_lm/models/gpt_oss.py +++ b/mlx_lm/models/gpt_oss.py @@ -3,7 +3,7 @@ import math from dataclasses import dataclass from functools import partial -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -17,7 +17,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "gpt_oss" + model_type: str num_hidden_layers: int = 36 num_local_experts: int = 128 num_experts_per_tok: int = 4 @@ -32,6 +32,11 @@ class ModelArgs(BaseModelArgs): rope_theta: int = 150000 rope_scaling: Any = None layer_types: list = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 150000.0) # These operators emulate particular methods in torch that don't exist in MLX natively From 951ca9086ee7ae6c100bf5d6f3d5da39199d6eff Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:59:44 +0100 Subject: [PATCH 58/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/glm4.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/glm4.py b/mlx_lm/models/glm4.py index 611175d29..3e69c279a 100644 --- a/mlx_lm/models/glm4.py +++ b/mlx_lm/models/glm4.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -26,6 +26,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float rope_traditional: bool = True max_position_embeddings: int = 32768 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Glm4MLP(nn.Module): From d5b421ed8b7272fa9d2c0da22a21570fb89abd17 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 18:59:57 +0100 Subject: [PATCH 59/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/glm4_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/glm4_moe.py b/mlx_lm/models/glm4_moe.py index bcf7cc9f4..bcf192ce9 100644 --- a/mlx_lm/models/glm4_moe.py +++ b/mlx_lm/models/glm4_moe.py @@ -44,6 +44,11 @@ class ModelArgs(BaseModelArgs): partial_rotary_factor: float scoring_func: str = "sigmoid" topk_method: str = "noaux_tc" + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From ee2b05da78e8ac9271b76f1dda70fdd35195e951 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:00:26 +0100 Subject: [PATCH 60/91] feat: update ModelArgs to make model_type required and adjust rope_theta initialization --- mlx_lm/models/glm4_moe_lite.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/glm4_moe_lite.py b/mlx_lm/models/glm4_moe_lite.py index dc108afdd..3e14c2146 100644 --- a/mlx_lm/models/glm4_moe_lite.py +++ b/mlx_lm/models/glm4_moe_lite.py @@ -18,7 +18,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "glm4_moe_lite" + model_type: str vocab_size: int = 154880 hidden_size: int = 2048 intermediate_size: int = 10240 @@ -44,7 +44,7 @@ class ModelArgs(BaseModelArgs): 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_theta: float = 1000000.0 rope_scaling: Optional[Dict] = None attention_bias: bool = False attention_dropout: float = 0.0 @@ -52,6 +52,11 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool = False num_nextn_predict_layers: int = 1 quantization: Optional[Dict[str, Any]] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) class Glm4MoeLiteAttention(nn.Module): From 4ec1c731ca71c28aba0d5a199585db5040d62ba7 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:01:22 +0100 Subject: [PATCH 61/91] feat: update ModelArgs to set rope_theta from rope_parameters in post_init --- mlx_lm/models/glm.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/mlx_lm/models/glm.py b/mlx_lm/models/glm.py index ab3853d0a..63cb8dc8b 100644 --- a/mlx_lm/models/glm.py +++ b/mlx_lm/models/glm.py @@ -1,14 +1,13 @@ -# Copyright © 2025 Apple Inc. +# Copyright © 2026 Apple Inc. from dataclasses import dataclass -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn from .activations import swiglu from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention -from .rope_utils import initialize_rope @dataclass @@ -20,12 +19,17 @@ class ModelArgs(BaseModelArgs): num_attention_heads: int rms_norm_eps: float vocab_size: int - head_dim: int num_key_value_heads: int + head_dim: Optional[int] = None max_position_embeddings: Optional[int] = None attention_bias: bool = False rope_theta: float = 10000 tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class GLMAttention(nn.Module): From 66a91f629706c701e11f271d40e0b1e48486caa7 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:01:53 +0100 Subject: [PATCH 62/91] feat: update ModelArgs to make rope_parameters optional and adjust rope_theta initialization --- mlx_lm/models/glm_moe_dsa.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index 14e96e365..3c8937197 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/mlx_lm/models/glm_moe_dsa.py @@ -38,14 +38,15 @@ class ModelArgs(BaseModelArgs): 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 + rope_parameters: Optional[Dict] = None def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) self.rope_scaling = self.rope_parameters - self.rope_theta = self.rope_parameters["rope_theta"] class Model(DSV32Model): From 9921082c68ca4fd7ba25d7acbf7c7be1cb8e79e1 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:02:12 +0100 Subject: [PATCH 63/91] feat: add rope_parameters to TextConfig and initialize rope_theta in post_init --- mlx_lm/models/gemma3n.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/gemma3n.py b/mlx_lm/models/gemma3n.py index 9b380cf72..292ee182a 100644 --- a/mlx_lm/models/gemma3n.py +++ b/mlx_lm/models/gemma3n.py @@ -40,6 +40,11 @@ class TextConfig(BaseModelArgs): altup_active_idx: int laurel_rank: int rope_scaling: Optional[Dict] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) @dataclass From 378bd010c41c7814dc10340f128965984a866e82 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:02:39 +0100 Subject: [PATCH 64/91] feat: add rope_parameters to ModelArgs and update rope_theta initialization in post_init --- mlx_lm/models/gemma3_text.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/gemma3_text.py b/mlx_lm/models/gemma3_text.py index 5cca40e0a..6a477ef27 100644 --- a/mlx_lm/models/gemma3_text.py +++ b/mlx_lm/models/gemma3_text.py @@ -23,13 +23,18 @@ class ModelArgs(BaseModelArgs): rms_norm_eps: float = 1.0e-6 vocab_size: int = 262144 num_key_value_heads: int = 1 - rope_theta: float = 1_000_000.0 + rope_theta: float = 1000000.0 rope_local_base_freq: float = 10_000.0 query_pre_attn_scalar: float = 256 sliding_window: int = 512 sliding_window_pattern: int = 6 max_position_embeddings: int = 32768 rope_scaling: Dict = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) class Attention(nn.Module): From fcecc454b7cab818e78a31c9dba6d62a7b4f317b Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:04:45 +0100 Subject: [PATCH 65/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/gemma2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/gemma2.py b/mlx_lm/models/gemma2.py index 23a98ac1a..cacd9bb4a 100644 --- a/mlx_lm/models/gemma2.py +++ b/mlx_lm/models/gemma2.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -25,6 +25,11 @@ class ModelArgs(BaseModelArgs): attn_logit_softcapping: float = 50.0 final_logit_softcapping: float = 30.0 query_pre_attn_scalar: float = 144.0 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class RMSNorm(nn.Module): From 05beddece39652a842c104bb8fa090447025df62 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:04:57 +0100 Subject: [PATCH 66/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/gemma.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/gemma.py b/mlx_lm/models/gemma.py index 9c540525e..722dba700 100644 --- a/mlx_lm/models/gemma.py +++ b/mlx_lm/models/gemma.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -22,6 +22,11 @@ class ModelArgs(BaseModelArgs): num_key_value_heads: int rope_theta: float = 10000 rope_traditional: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class RMSNorm(nn.Module): From 466bbcfc8b9c89fc10ad679db4f9bb66444648a6 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:05:36 +0100 Subject: [PATCH 67/91] feat: add model_type attribute to ModelArgs class --- mlx_lm/models/falcon_h1.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/falcon_h1.py b/mlx_lm/models/falcon_h1.py index 773893485..ae7b62eec 100644 --- a/mlx_lm/models/falcon_h1.py +++ b/mlx_lm/models/falcon_h1.py @@ -20,6 +20,7 @@ @dataclass class ModelArgs(BaseModelArgs): + model_type: str attention_bias: bool = False attention_in_multiplier: float = 1.0 attention_out_multiplier: float = 0.9375 @@ -49,7 +50,6 @@ class ModelArgs(BaseModelArgs): mlp_multipliers: List[float] = field( default_factory=lambda: [0.8838834764831844, 0.5859375] ) - model_type: str = "falcon_h1" num_attention_heads: int = 8 num_hidden_layers: int = 36 num_key_value_heads: int = 2 From b9da26ec642bf63afaec32e7bf0863307ddc0f39 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:05:50 +0100 Subject: [PATCH 68/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/exaone4.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/exaone4.py b/mlx_lm/models/exaone4.py index c0e15f5e2..0e52947c8 100644 --- a/mlx_lm/models/exaone4.py +++ b/mlx_lm/models/exaone4.py @@ -29,6 +29,11 @@ class ModelArgs(BaseModelArgs): rope_scaling: Dict[str, Union[float, str]] sliding_window: Optional[int] sliding_window_pattern: Optional[str] + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 1205638fb5f4c38f19dfc05a72731bfdc262498e Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:06:04 +0100 Subject: [PATCH 69/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/exaone.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/exaone.py b/mlx_lm/models/exaone.py index 3cd223d29..226655fed 100644 --- a/mlx_lm/models/exaone.py +++ b/mlx_lm/models/exaone.py @@ -29,6 +29,11 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool = True attention_bias: bool = False mlp_bias: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class AttentionModule(nn.Module): From 47ebffbefd91496e8773fd43a68e0e00f4c1fe7c Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:06:32 +0100 Subject: [PATCH 70/91] feat: update rope_theta initialization in ModelArgs to use default value --- mlx_lm/models/exaone_moe.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mlx_lm/models/exaone_moe.py b/mlx_lm/models/exaone_moe.py index a6da99be8..31b72b31c 100644 --- a/mlx_lm/models/exaone_moe.py +++ b/mlx_lm/models/exaone_moe.py @@ -1,7 +1,7 @@ # Copyright © 2026 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -43,10 +43,11 @@ class ModelArgs(BaseModelArgs): rope_scaling: Optional[dict] = None rope_parameters: Optional[dict] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None def __post_init__(self): - if self.rope_parameters is not None and "rope_theta" in self.rope_parameters: - self.rope_theta = self.rope_parameters["rope_theta"] + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) @mx.compile From ba47261d4c7912a75ec673caad0d3cb1107587ee Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:07:08 +0100 Subject: [PATCH 71/91] feat: initialize rope_theta from rope_parameters in ModelArgs post_init --- mlx_lm/models/ernie4_5.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/ernie4_5.py b/mlx_lm/models/ernie4_5.py index d6551f821..8c949d9da 100644 --- a/mlx_lm/models/ernie4_5.py +++ b/mlx_lm/models/ernie4_5.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -26,6 +26,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float use_bias: bool tie_word_embeddings: bool + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 1355bc4b8276c417481327fdd1bbd6f0c40a8082 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:07:37 +0100 Subject: [PATCH 72/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/ernie4_5_moe.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/ernie4_5_moe.py b/mlx_lm/models/ernie4_5_moe.py index 597e9c467..331fd3e53 100644 --- a/mlx_lm/models/ernie4_5_moe.py +++ b/mlx_lm/models/ernie4_5_moe.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass, field -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -37,6 +37,11 @@ class ModelArgs(BaseModelArgs): moe_layer_end_index: Optional[int] = None head_dim: Optional[int] = None moe_gate_act: str = "softmax" + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From ea9fe1fcd9dcf5b436b9b90dc0f54ee7aa16f998 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:07:53 +0100 Subject: [PATCH 73/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/dots1.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/dots1.py b/mlx_lm/models/dots1.py index c568f39db..69e1c81e3 100644 --- a/mlx_lm/models/dots1.py +++ b/mlx_lm/models/dots1.py @@ -1,7 +1,6 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from functools import partial from typing import Any, Dict, Optional, Union import mlx.core as mx @@ -40,6 +39,11 @@ class ModelArgs(BaseModelArgs): mlp_bias: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Dots1Attention(nn.Module): From ffdcca9accea6a53b677af52df7309c1db2991ae Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:08:13 +0100 Subject: [PATCH 74/91] feat: update ModelArgs to require model_type and add rope_parameters with post_init handling --- mlx_lm/models/deepseek.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek.py b/mlx_lm/models/deepseek.py index 2876f7b65..8bb3a0b37 100644 --- a/mlx_lm/models/deepseek.py +++ b/mlx_lm/models/deepseek.py @@ -11,7 +11,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "deepseek" + model_type: str vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 11008 @@ -29,6 +29,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Optional[Dict] = None attention_bias: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class DeepseekAttention(nn.Module): From 1316e2b71c2e52be3fb10de1cfd5ce46e30d007a Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:08:33 +0100 Subject: [PATCH 75/91] feat: update ModelArgs to require model_type and add rope_parameters with post_init handling --- mlx_lm/models/deepseek_v32.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index e40e52950..684ab0753 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -18,7 +18,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "deepseek_v32" + model_type: str vocab_size: int = 102400 hidden_size: int = 4096 index_head_dim: int = 128 @@ -50,6 +50,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Dict = None attention_bias: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Indexer(nn.Module): From 3a526161be7922efab26a57d373ef7960172c42c Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:08:50 +0100 Subject: [PATCH 76/91] feat: update ModelArgs to require model_type and add rope_parameters with post_init handling --- mlx_lm/models/deepseek_v3.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/deepseek_v3.py b/mlx_lm/models/deepseek_v3.py index a2766e59e..cd1025e5c 100644 --- a/mlx_lm/models/deepseek_v3.py +++ b/mlx_lm/models/deepseek_v3.py @@ -2,7 +2,6 @@ import math from dataclasses import dataclass -from functools import partial from typing import Any, Dict, Optional import mlx.core as mx @@ -19,7 +18,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "deepseek_v3" + model_type: str vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 11008 @@ -48,6 +47,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Dict = None attention_bias: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class DeepseekV3Attention(nn.Module): From 6c5d78b8d5abd260d44c3dbddef6b3976d311336 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:09:05 +0100 Subject: [PATCH 77/91] feat: update ModelArgs to require model_type and add rope_parameters with post_init handling --- mlx_lm/models/deepseek_v2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/deepseek_v2.py b/mlx_lm/models/deepseek_v2.py index a263d850a..b434b2d30 100644 --- a/mlx_lm/models/deepseek_v2.py +++ b/mlx_lm/models/deepseek_v2.py @@ -16,7 +16,7 @@ @dataclass class ModelArgs(BaseModelArgs): - model_type: str = "deepseek_v2" + model_type: str vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 11008 @@ -43,6 +43,11 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000.0 rope_scaling: Dict = None attention_bias: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) def yarn_find_correction_dim( From 7b4af214174b4beba9456034757bc0e944d86042 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:09:30 +0100 Subject: [PATCH 78/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/cohere2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/cohere2.py b/mlx_lm/models/cohere2.py index c93e238de..4f816e132 100644 --- a/mlx_lm/models/cohere2.py +++ b/mlx_lm/models/cohere2.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Optional, Tuple, Dict import mlx.core as mx import mlx.nn as nn @@ -28,6 +28,11 @@ class ModelArgs(BaseModelArgs): layer_norm_bias: bool = False sliding_window: int = 4096 sliding_window_pattern: int = 4 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 50000.0) class Attention(nn.Module): From ffa4b062ac39427bd4298b1b7d65635cb58a9b56 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:09:52 +0100 Subject: [PATCH 79/91] feat: update ModelArgs to initialize rope_theta from rope_parameters in post_init --- mlx_lm/models/cohere.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/cohere.py b/mlx_lm/models/cohere.py index c9b88885a..57984dba4 100644 --- a/mlx_lm/models/cohere.py +++ b/mlx_lm/models/cohere.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -25,6 +25,11 @@ class ModelArgs(BaseModelArgs): attention_bias: bool = False layer_norm_bias: bool = False use_qk_norm: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 8000000.0) class LayerNorm2D(nn.Module): From 1520a93762693bd8a40a022f857d2ca3640b78a5 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:10:09 +0100 Subject: [PATCH 80/91] feat: update ModelArgs to set rope_theta from rope_parameters in post_init --- mlx_lm/models/bitnet.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/bitnet.py b/mlx_lm/models/bitnet.py index 72794e564..427e6ef14 100644 --- a/mlx_lm/models/bitnet.py +++ b/mlx_lm/models/bitnet.py @@ -1,7 +1,6 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from functools import partial from typing import Any, Dict, Optional, Union import mlx.core as mx @@ -30,6 +29,11 @@ class ModelArgs(BaseModelArgs): rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 4fe3a47774cf2afb9296b5a8e9ecdc0517a38a1b Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:10:21 +0100 Subject: [PATCH 81/91] fix: remove unused import of Any from typing in base.py --- mlx_lm/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/base.py b/mlx_lm/models/base.py index d7c3efb28..729a484dc 100644 --- a/mlx_lm/models/base.py +++ b/mlx_lm/models/base.py @@ -2,7 +2,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Optional +from typing import Optional import mlx.core as mx from mlx.utils import tree_map From 323d65bf4759b0239de8e473e08392b6a0df51b2 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:10:34 +0100 Subject: [PATCH 82/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/bailing_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/bailing_moe.py b/mlx_lm/models/bailing_moe.py index ff7d908ea..fb5634b8b 100644 --- a/mlx_lm/models/bailing_moe.py +++ b/mlx_lm/models/bailing_moe.py @@ -48,6 +48,11 @@ class ModelArgs(BaseModelArgs): topk_group: int = 4 moe_shared_expert_intermediate_size: Optional[int] = None moe_router_enable_shared_expert: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) @partial(mx.compile, shapeless=True) From f00b4f3d5a8f777d27881a6ed88c8a73481e25d0 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:10:52 +0100 Subject: [PATCH 83/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/bailing_moe_linear.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlx_lm/models/bailing_moe_linear.py b/mlx_lm/models/bailing_moe_linear.py index 360f99e1b..6d172e2e7 100644 --- a/mlx_lm/models/bailing_moe_linear.py +++ b/mlx_lm/models/bailing_moe_linear.py @@ -58,6 +58,11 @@ class ModelArgs(BaseModelArgs): moe_shared_expert_intermediate_size: Optional[int] = None moe_router_enable_shared_expert: bool = True head_dim: Optional[int] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) def recurrent_gla( From 55d29e248fdb7ee5e04c854f3d0e08dcd2898b64 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:11:45 +0100 Subject: [PATCH 84/91] feat: add rope_parameters to ModelArgs and initialize rope_theta in post_init --- mlx_lm/models/apertus.py | 6 +++++- mlx_lm/models/baichuan_m1.py | 9 +++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/mlx_lm/models/apertus.py b/mlx_lm/models/apertus.py index 4d8afa28d..587719cf3 100644 --- a/mlx_lm/models/apertus.py +++ b/mlx_lm/models/apertus.py @@ -1,7 +1,6 @@ # Copyright © 2023-2025 Apple Inc. from dataclasses import dataclass -from functools import partial from typing import Any, Dict, Optional, Union import mlx.core as mx @@ -31,6 +30,11 @@ class ModelArgs(BaseModelArgs): tie_word_embeddings: bool rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class ApertusMLP(nn.Module): diff --git a/mlx_lm/models/baichuan_m1.py b/mlx_lm/models/baichuan_m1.py index 1a52b4f36..5f56c5ce6 100644 --- a/mlx_lm/models/baichuan_m1.py +++ b/mlx_lm/models/baichuan_m1.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Optional, Dict import mlx.core as mx import mlx.nn as nn @@ -13,6 +13,7 @@ @dataclass class ModelArgs(BaseModelArgs): + model_type: str vocab_size: int hidden_size: int intermediate_size: int @@ -24,10 +25,14 @@ class ModelArgs(BaseModelArgs): sliding_window_layers: List[int] conv_window: int rms_norm_eps: float - model_type: str = "baichuan_m1" num_swa_attention_heads: Optional[int] = None num_swa_key_value_heads: Optional[int] = None tie_word_embeddings: bool = False + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 1f7401b4b12f2df4211bc8958b9a6f6c3b9930f8 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 6 Mar 2026 19:12:05 +0100 Subject: [PATCH 85/91] feat: add rope_parameters to ModelArgs and update rope_theta in post_init --- mlx_lm/models/afmoe.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx_lm/models/afmoe.py b/mlx_lm/models/afmoe.py index e23dc237a..184512251 100644 --- a/mlx_lm/models/afmoe.py +++ b/mlx_lm/models/afmoe.py @@ -31,7 +31,6 @@ class ModelArgs(BaseModelArgs): rope_theta: float = 10000 rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False - # MoE config num_experts: int = 128 num_experts_per_tok: int = 8 num_shared_experts: int = 1 @@ -43,6 +42,11 @@ class ModelArgs(BaseModelArgs): topk_group: int = 1 sliding_window: int = 2048 mup_enabled: bool = True + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) class Attention(nn.Module): From 1245c2ec70e32b7865e3c724a95379a6b86afca8 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Sat, 7 Mar 2026 14:56:08 +0100 Subject: [PATCH 86/91] format --- mlx_lm/models/Klear.py | 2 +- mlx_lm/models/baichuan_m1.py | 2 +- mlx_lm/models/cohere.py | 2 +- mlx_lm/models/cohere2.py | 2 +- mlx_lm/models/ernie4_5.py | 2 +- mlx_lm/models/ernie4_5_moe.py | 2 +- mlx_lm/models/exaone_moe.py | 2 +- mlx_lm/models/gemma.py | 2 +- mlx_lm/models/gemma2.py | 2 +- mlx_lm/models/glm4.py | 2 +- mlx_lm/models/gpt_oss.py | 2 +- mlx_lm/models/granitemoehybrid.py | 2 +- mlx_lm/models/helium.py | 2 +- mlx_lm/models/lfm2.py | 2 +- mlx_lm/models/lille-130m.py | 2 +- mlx_lm/models/llama4.py | 2 +- mlx_lm/models/llama4_text.py | 2 +- mlx_lm/models/nanochat.py | 2 +- mlx_lm/models/olmo.py | 2 +- mlx_lm/models/plamo.py | 2 +- mlx_lm/models/plamo2.py | 2 +- mlx_lm/models/qwen.py | 2 +- mlx_lm/models/recurrent_gemma.py | 2 +- mlx_lm/models/smollm3.py | 2 +- mlx_lm/models/starcoder2.py | 2 +- 25 files changed, 25 insertions(+), 25 deletions(-) diff --git a/mlx_lm/models/Klear.py b/mlx_lm/models/Klear.py index b51c06d74..bba975864 100644 --- a/mlx_lm/models/Klear.py +++ b/mlx_lm/models/Klear.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional, Dict +from typing import Any, Dict, List, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/baichuan_m1.py b/mlx_lm/models/baichuan_m1.py index 5f56c5ce6..236f46ef7 100644 --- a/mlx_lm/models/baichuan_m1.py +++ b/mlx_lm/models/baichuan_m1.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional, Dict +from typing import Any, Dict, List, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/cohere.py b/mlx_lm/models/cohere.py index 57984dba4..7978631a2 100644 --- a/mlx_lm/models/cohere.py +++ b/mlx_lm/models/cohere.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/cohere2.py b/mlx_lm/models/cohere2.py index 4f816e132..d9ddf1528 100644 --- a/mlx_lm/models/cohere2.py +++ b/mlx_lm/models/cohere2.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Optional, Tuple, Dict +from typing import Dict, Optional, Tuple import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/ernie4_5.py b/mlx_lm/models/ernie4_5.py index 8c949d9da..503562e0a 100644 --- a/mlx_lm/models/ernie4_5.py +++ b/mlx_lm/models/ernie4_5.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/ernie4_5_moe.py b/mlx_lm/models/ernie4_5_moe.py index 331fd3e53..2c6c8119c 100644 --- a/mlx_lm/models/ernie4_5_moe.py +++ b/mlx_lm/models/ernie4_5_moe.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass, field -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/exaone_moe.py b/mlx_lm/models/exaone_moe.py index 31b72b31c..e4093d37c 100644 --- a/mlx_lm/models/exaone_moe.py +++ b/mlx_lm/models/exaone_moe.py @@ -1,7 +1,7 @@ # Copyright © 2026 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional, Dict +from typing import Any, Dict, List, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/gemma.py b/mlx_lm/models/gemma.py index 722dba700..8899b458c 100644 --- a/mlx_lm/models/gemma.py +++ b/mlx_lm/models/gemma.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/gemma2.py b/mlx_lm/models/gemma2.py index cacd9bb4a..5544701ec 100644 --- a/mlx_lm/models/gemma2.py +++ b/mlx_lm/models/gemma2.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/glm4.py b/mlx_lm/models/glm4.py index 3e69c279a..da3d80156 100644 --- a/mlx_lm/models/glm4.py +++ b/mlx_lm/models/glm4.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/gpt_oss.py b/mlx_lm/models/gpt_oss.py index 9eb5659de..ebc6ab2a3 100644 --- a/mlx_lm/models/gpt_oss.py +++ b/mlx_lm/models/gpt_oss.py @@ -3,7 +3,7 @@ import math from dataclasses import dataclass from functools import partial -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/granitemoehybrid.py b/mlx_lm/models/granitemoehybrid.py index fd55e66b8..c38bb4d23 100644 --- a/mlx_lm/models/granitemoehybrid.py +++ b/mlx_lm/models/granitemoehybrid.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional, Tuple, Dict +from typing import Any, Dict, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/helium.py b/mlx_lm/models/helium.py index c60a680d1..bcb838921 100644 --- a/mlx_lm/models/helium.py +++ b/mlx_lm/models/helium.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/lfm2.py b/mlx_lm/models/lfm2.py index 52df0e5c0..b000e2547 100644 --- a/mlx_lm/models/lfm2.py +++ b/mlx_lm/models/lfm2.py @@ -1,6 +1,6 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, List, Optional, Dict +from typing import Any, Dict, List, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/lille-130m.py b/mlx_lm/models/lille-130m.py index 775eb59ea..1299593c6 100644 --- a/mlx_lm/models/lille-130m.py +++ b/mlx_lm/models/lille-130m.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/llama4.py b/mlx_lm/models/llama4.py index f39d222de..7149c99fc 100644 --- a/mlx_lm/models/llama4.py +++ b/mlx_lm/models/llama4.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Union, Dict +from typing import Any, Dict, Optional, Union import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/llama4_text.py b/mlx_lm/models/llama4_text.py index 3301ea94c..95a4c4d53 100644 --- a/mlx_lm/models/llama4_text.py +++ b/mlx_lm/models/llama4_text.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/nanochat.py b/mlx_lm/models/nanochat.py index 8d2c1c898..cc97da6a6 100644 --- a/mlx_lm/models/nanochat.py +++ b/mlx_lm/models/nanochat.py @@ -3,7 +3,7 @@ import math from dataclasses import dataclass from functools import partial -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/olmo.py b/mlx_lm/models/olmo.py index 8e13a06df..e06a04a48 100644 --- a/mlx_lm/models/olmo.py +++ b/mlx_lm/models/olmo.py @@ -2,7 +2,7 @@ import sys from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/plamo.py b/mlx_lm/models/plamo.py index 4bbd65f7c..575d5655d 100644 --- a/mlx_lm/models/plamo.py +++ b/mlx_lm/models/plamo.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/plamo2.py b/mlx_lm/models/plamo2.py index 0f0d18a47..6c81693d1 100644 --- a/mlx_lm/models/plamo2.py +++ b/mlx_lm/models/plamo2.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/qwen.py b/mlx_lm/models/qwen.py index a57566ae6..945775b42 100644 --- a/mlx_lm/models/qwen.py +++ b/mlx_lm/models/qwen.py @@ -1,10 +1,10 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn -from typing import Optional, Dict from .activations import swiglu from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention diff --git a/mlx_lm/models/recurrent_gemma.py b/mlx_lm/models/recurrent_gemma.py index 50ab4d330..71f2d3dc8 100644 --- a/mlx_lm/models/recurrent_gemma.py +++ b/mlx_lm/models/recurrent_gemma.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. -from ast import Dict import math +from ast import Dict from dataclasses import dataclass from typing import List, Optional diff --git a/mlx_lm/models/smollm3.py b/mlx_lm/models/smollm3.py index 2e4f8f641..5c3a10b69 100644 --- a/mlx_lm/models/smollm3.py +++ b/mlx_lm/models/smollm3.py @@ -1,7 +1,7 @@ # Copyright © 2025 Apple Inc. from dataclasses import dataclass -from typing import Optional, Dict +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/starcoder2.py b/mlx_lm/models/starcoder2.py index e0c8d9bac..ebd7af23c 100644 --- a/mlx_lm/models/starcoder2.py +++ b/mlx_lm/models/starcoder2.py @@ -1,7 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn From 345c26f9187824082b0325b08b5e58d64a55397a 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 Mar 2026 09:19:53 +0100 Subject: [PATCH 87/91] feat: update ModelArgs initialization in test_phi2 with new parameters including rope_theta --- mlx_lm/models/afm7.py | 2 +- tests/test_models.py | 14 +++++++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/afm7.py b/mlx_lm/models/afm7.py index 81123ec77..f8d3b8a1b 100644 --- a/mlx_lm/models/afm7.py +++ b/mlx_lm/models/afm7.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from functools import partial from itertools import accumulate -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/tests/test_models.py b/tests/test_models.py index b384d748a..162611e65 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -397,7 +397,19 @@ def test_bitnet(self): def test_phi2(self): from mlx_lm.models import phi - args = phi.ModelArgs() + args = phi.ModelArgs( + model_type="phi", + max_position_embeddings=128, + vocab_size=10_000, + hidden_size=1024, + num_attention_heads=4, + num_hidden_layers=4, + num_key_value_heads=4, + partial_rotary_factor=0.4, + intermediate_size=2048, + layer_norm_eps=1e-5, + rope_theta=1000, + ) model = phi.Model(args) self.model_test_runner( model, args.model_type, args.vocab_size, args.num_hidden_layers From 8dc64492d6f89fcc679615de6b7be90f2db765d6 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, 17 Mar 2026 14:37:54 +0100 Subject: [PATCH 88/91] Update rope parameters across multiple models - Increased default value of `rope_theta` from 10000.0 to 100000.0 in various model argument classes. - Introduced `rope_scaling` and `rope_traditional` parameters in several models to enhance flexibility in rope parameter configurations. - Ensured consistent handling of `rope_parameters` across models including `deepseek`, `exaone`, `gemma`, `glm`, `hunyuan`, `llama`, `qwen`, and others. --- mlx_lm/models/afmoe.py | 3 ++- mlx_lm/models/apertus.py | 4 +++- mlx_lm/models/bailing_moe.py | 3 ++- mlx_lm/models/bailing_moe_linear.py | 4 +++- mlx_lm/models/bitnet.py | 4 +++- mlx_lm/models/deepseek.py | 3 ++- mlx_lm/models/deepseek_v2.py | 3 ++- mlx_lm/models/deepseek_v3.py | 3 ++- mlx_lm/models/deepseek_v32.py | 3 ++- mlx_lm/models/dots1.py | 3 ++- mlx_lm/models/exaone.py | 4 +++- mlx_lm/models/exaone4.py | 3 ++- mlx_lm/models/exaone_moe.py | 3 ++- mlx_lm/models/gemma.py | 3 ++- mlx_lm/models/gemma2.py | 3 ++- mlx_lm/models/gemma3_text.py | 5 ++++- mlx_lm/models/gemma3n.py | 4 +++- mlx_lm/models/glm4.py | 3 ++- mlx_lm/models/glm4_moe.py | 3 ++- mlx_lm/models/glm4_moe_lite.py | 3 ++- mlx_lm/models/glm_moe_dsa.py | 4 ++-- mlx_lm/models/gpt_oss.py | 3 ++- mlx_lm/models/granite.py | 3 ++- mlx_lm/models/granitemoe.py | 3 ++- mlx_lm/models/granitemoehybrid.py | 11 ----------- mlx_lm/models/hunyuan.py | 3 ++- mlx_lm/models/hunyuan_v1_dense.py | 3 ++- mlx_lm/models/internlm2.py | 4 +++- mlx_lm/models/internlm3.py | 4 +++- mlx_lm/models/iquestloopcoder.py | 3 ++- mlx_lm/models/kimi_linear.py | 6 ++++++ mlx_lm/models/kimi_vl.py | 3 ++- mlx_lm/models/lfm2.py | 4 +--- mlx_lm/models/llama.py | 4 +++- mlx_lm/models/llama4.py | 1 + mlx_lm/models/llama4_text.py | 1 + mlx_lm/models/longcat_flash.py | 1 + mlx_lm/models/longcat_flash_ngram.py | 1 + mlx_lm/models/mimo.py | 4 +++- mlx_lm/models/minicpm.py | 4 +++- mlx_lm/models/minicpm3.py | 4 +++- mlx_lm/models/mixtral.py | 2 ++ mlx_lm/models/nemotron-nas.py | 3 ++- mlx_lm/models/nemotron.py | 4 +++- mlx_lm/models/olmo.py | 3 ++- mlx_lm/models/olmo2.py | 4 +++- mlx_lm/models/olmo3.py | 3 ++- mlx_lm/models/olmoe.py | 4 +++- mlx_lm/models/phi3.py | 4 +++- mlx_lm/models/phimoe.py | 3 ++- mlx_lm/models/phixtral.py | 1 + mlx_lm/models/plamo.py | 1 + mlx_lm/models/plamo2.py | 5 ----- mlx_lm/models/qwen.py | 3 --- mlx_lm/models/qwen2.py | 4 +++- mlx_lm/models/qwen2_moe.py | 4 +++- mlx_lm/models/qwen3.py | 1 + mlx_lm/models/qwen3_moe.py | 1 + mlx_lm/models/qwen3_next.py | 1 + mlx_lm/models/rwkv7.py | 5 ----- mlx_lm/models/seed_oss.py | 2 ++ mlx_lm/models/solar_open.py | 1 + mlx_lm/models/step3p5.py | 3 ++- mlx_lm/models/telechat3.py | 3 ++- mlx_lm/models/youtu_llm.py | 5 +++-- 65 files changed, 135 insertions(+), 76 deletions(-) diff --git a/mlx_lm/models/afmoe.py b/mlx_lm/models/afmoe.py index 184512251..0fe6d838c 100644 --- a/mlx_lm/models/afmoe.py +++ b/mlx_lm/models/afmoe.py @@ -46,7 +46,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/apertus.py b/mlx_lm/models/apertus.py index 587719cf3..4ab530334 100644 --- a/mlx_lm/models/apertus.py +++ b/mlx_lm/models/apertus.py @@ -34,7 +34,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class ApertusMLP(nn.Module): diff --git a/mlx_lm/models/bailing_moe.py b/mlx_lm/models/bailing_moe.py index fb5634b8b..f0680b780 100644 --- a/mlx_lm/models/bailing_moe.py +++ b/mlx_lm/models/bailing_moe.py @@ -52,7 +52,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters @partial(mx.compile, shapeless=True) diff --git a/mlx_lm/models/bailing_moe_linear.py b/mlx_lm/models/bailing_moe_linear.py index 6d172e2e7..b8a1008d0 100644 --- a/mlx_lm/models/bailing_moe_linear.py +++ b/mlx_lm/models/bailing_moe_linear.py @@ -62,7 +62,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) def recurrent_gla( diff --git a/mlx_lm/models/bitnet.py b/mlx_lm/models/bitnet.py index 427e6ef14..173518e1e 100644 --- a/mlx_lm/models/bitnet.py +++ b/mlx_lm/models/bitnet.py @@ -33,7 +33,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Attention(nn.Module): diff --git a/mlx_lm/models/deepseek.py b/mlx_lm/models/deepseek.py index 8bb3a0b37..3c31df108 100644 --- a/mlx_lm/models/deepseek.py +++ b/mlx_lm/models/deepseek.py @@ -33,7 +33,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class DeepseekAttention(nn.Module): diff --git a/mlx_lm/models/deepseek_v2.py b/mlx_lm/models/deepseek_v2.py index b434b2d30..ae45e6d5d 100644 --- a/mlx_lm/models/deepseek_v2.py +++ b/mlx_lm/models/deepseek_v2.py @@ -47,7 +47,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters def yarn_find_correction_dim( diff --git a/mlx_lm/models/deepseek_v3.py b/mlx_lm/models/deepseek_v3.py index cd1025e5c..376032e49 100644 --- a/mlx_lm/models/deepseek_v3.py +++ b/mlx_lm/models/deepseek_v3.py @@ -51,7 +51,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class DeepseekV3Attention(nn.Module): diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index 684ab0753..12379b6e0 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -54,7 +54,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Indexer(nn.Module): diff --git a/mlx_lm/models/dots1.py b/mlx_lm/models/dots1.py index 69e1c81e3..252a1dcad 100644 --- a/mlx_lm/models/dots1.py +++ b/mlx_lm/models/dots1.py @@ -43,7 +43,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Dots1Attention(nn.Module): diff --git a/mlx_lm/models/exaone.py b/mlx_lm/models/exaone.py index 226655fed..9ace043ba 100644 --- a/mlx_lm/models/exaone.py +++ b/mlx_lm/models/exaone.py @@ -33,7 +33,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class AttentionModule(nn.Module): diff --git a/mlx_lm/models/exaone4.py b/mlx_lm/models/exaone4.py index 0e52947c8..fe9d0ae8b 100644 --- a/mlx_lm/models/exaone4.py +++ b/mlx_lm/models/exaone4.py @@ -33,7 +33,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/exaone_moe.py b/mlx_lm/models/exaone_moe.py index e4093d37c..cec3401d5 100644 --- a/mlx_lm/models/exaone_moe.py +++ b/mlx_lm/models/exaone_moe.py @@ -47,7 +47,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters @mx.compile diff --git a/mlx_lm/models/gemma.py b/mlx_lm/models/gemma.py index 8899b458c..d920a4919 100644 --- a/mlx_lm/models/gemma.py +++ b/mlx_lm/models/gemma.py @@ -26,7 +26,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class RMSNorm(nn.Module): diff --git a/mlx_lm/models/gemma2.py b/mlx_lm/models/gemma2.py index 5544701ec..5defceae5 100644 --- a/mlx_lm/models/gemma2.py +++ b/mlx_lm/models/gemma2.py @@ -29,7 +29,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class RMSNorm(nn.Module): diff --git a/mlx_lm/models/gemma3_text.py b/mlx_lm/models/gemma3_text.py index 6a477ef27..96761a911 100644 --- a/mlx_lm/models/gemma3_text.py +++ b/mlx_lm/models/gemma3_text.py @@ -34,7 +34,10 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_local_base_freq = self.rope_parameters.get("rope_local_base_freq", 10000.0) + class Attention(nn.Module): diff --git a/mlx_lm/models/gemma3n.py b/mlx_lm/models/gemma3n.py index 292ee182a..ee3a4ace4 100644 --- a/mlx_lm/models/gemma3n.py +++ b/mlx_lm/models/gemma3n.py @@ -44,7 +44,9 @@ class TextConfig(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_local_base_freq = self.rope_parameters.get("rope_local_base_freq", 10000.0) @dataclass diff --git a/mlx_lm/models/glm4.py b/mlx_lm/models/glm4.py index da3d80156..177a8b6fb 100644 --- a/mlx_lm/models/glm4.py +++ b/mlx_lm/models/glm4.py @@ -30,7 +30,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Glm4MLP(nn.Module): diff --git a/mlx_lm/models/glm4_moe.py b/mlx_lm/models/glm4_moe.py index bcf192ce9..73b11eb9f 100644 --- a/mlx_lm/models/glm4_moe.py +++ b/mlx_lm/models/glm4_moe.py @@ -48,7 +48,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/glm4_moe_lite.py b/mlx_lm/models/glm4_moe_lite.py index 3e14c2146..39d170939 100644 --- a/mlx_lm/models/glm4_moe_lite.py +++ b/mlx_lm/models/glm4_moe_lite.py @@ -56,7 +56,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Glm4MoeLiteAttention(nn.Module): diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index 3c8937197..661ba5023 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/mlx_lm/models/glm_moe_dsa.py @@ -45,8 +45,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) - self.rope_scaling = self.rope_parameters + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Model(DSV32Model): diff --git a/mlx_lm/models/gpt_oss.py b/mlx_lm/models/gpt_oss.py index ebc6ab2a3..5daf6eaa3 100644 --- a/mlx_lm/models/gpt_oss.py +++ b/mlx_lm/models/gpt_oss.py @@ -36,7 +36,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 150000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters # These operators emulate particular methods in torch that don't exist in MLX natively diff --git a/mlx_lm/models/granite.py b/mlx_lm/models/granite.py index f0eee1305..2000a993b 100644 --- a/mlx_lm/models/granite.py +++ b/mlx_lm/models/granite.py @@ -35,7 +35,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/granitemoe.py b/mlx_lm/models/granitemoe.py index 0cb30c7d2..42fc6a3c1 100644 --- a/mlx_lm/models/granitemoe.py +++ b/mlx_lm/models/granitemoe.py @@ -36,7 +36,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class GraniteMoeAttention(nn.Module): diff --git a/mlx_lm/models/granitemoehybrid.py b/mlx_lm/models/granitemoehybrid.py index c38bb4d23..586edf096 100644 --- a/mlx_lm/models/granitemoehybrid.py +++ b/mlx_lm/models/granitemoehybrid.py @@ -21,7 +21,6 @@ @dataclass class ModelArgs(BaseModelArgs): - # Required fields (no defaults) model_type: str vocab_size: int hidden_size: int @@ -38,14 +37,9 @@ class ModelArgs(BaseModelArgs): layer_types: List[str] rms_norm_eps: float rope_theta: float - - # Optional fields (with defaults) - # MoE parameters (optional for dense mode) num_local_experts: Optional[int] = None num_experts_per_tok: Optional[int] = None shared_intermediate_size: Optional[int] = None - - # Mamba parameters (optional for non-hybrid mode) mamba_n_heads: Optional[int] = None mamba_d_head: Optional[int] = None mamba_proj_bias: Optional[bool] = None @@ -53,11 +47,7 @@ class ModelArgs(BaseModelArgs): mamba_d_conv: Optional[int] = None mamba_n_groups: Optional[int] = None mamba_conv_bias: Optional[bool] = None - - # Dense MLP parameters (for non-MoE mode) mlp_bias: bool = False - - # Other optional parameters position_embedding_type: str = "rope" tie_word_embeddings: bool = True time_step_limit: Tuple[float, float] = (0.001, 100.0) @@ -68,7 +58,6 @@ def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) - # Mode flags - inferred from num_local_experts @property def use_moe(self) -> bool: return bool(self.num_local_experts) diff --git a/mlx_lm/models/hunyuan.py b/mlx_lm/models/hunyuan.py index ff6004c37..27069121c 100644 --- a/mlx_lm/models/hunyuan.py +++ b/mlx_lm/models/hunyuan.py @@ -37,7 +37,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters if self.rope_scaling: required_keys = {"factor", "type"} diff --git a/mlx_lm/models/hunyuan_v1_dense.py b/mlx_lm/models/hunyuan_v1_dense.py index 547373929..4f4e9484e 100644 --- a/mlx_lm/models/hunyuan_v1_dense.py +++ b/mlx_lm/models/hunyuan_v1_dense.py @@ -31,7 +31,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters def __post_init__(self): if self.rope_scaling: diff --git a/mlx_lm/models/internlm2.py b/mlx_lm/models/internlm2.py index f97b8bce2..d0e447d1e 100644 --- a/mlx_lm/models/internlm2.py +++ b/mlx_lm/models/internlm2.py @@ -30,7 +30,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/internlm3.py b/mlx_lm/models/internlm3.py index fbb16e49a..d77ae9ed0 100644 --- a/mlx_lm/models/internlm3.py +++ b/mlx_lm/models/internlm3.py @@ -31,7 +31,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/iquestloopcoder.py b/mlx_lm/models/iquestloopcoder.py index 6b4fc8cf0..75154436a 100644 --- a/mlx_lm/models/iquestloopcoder.py +++ b/mlx_lm/models/iquestloopcoder.py @@ -55,7 +55,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class LoopGateProjection(nn.Module): diff --git a/mlx_lm/models/kimi_linear.py b/mlx_lm/models/kimi_linear.py index b2fa6d18f..afd1c7ad8 100644 --- a/mlx_lm/models/kimi_linear.py +++ b/mlx_lm/models/kimi_linear.py @@ -52,6 +52,12 @@ class ModelArgs(BaseModelArgs): use_grouped_topk: bool = True num_expert_group: int = 1 topk_group: int = 1 + rope_parameters: Optional[Dict] = None + + def __post_init__(self): + if self.rope_parameters is not None: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class KimiMLP(nn.Module): diff --git a/mlx_lm/models/kimi_vl.py b/mlx_lm/models/kimi_vl.py index 9cfa7a7fe..849b859a4 100644 --- a/mlx_lm/models/kimi_vl.py +++ b/mlx_lm/models/kimi_vl.py @@ -44,7 +44,8 @@ class TextArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters @dataclass diff --git a/mlx_lm/models/lfm2.py b/mlx_lm/models/lfm2.py index b000e2547..e2dedee09 100644 --- a/mlx_lm/models/lfm2.py +++ b/mlx_lm/models/lfm2.py @@ -40,9 +40,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) - if self.rope_parameters is not None and "rope_theta" in self.rope_parameters: - self.rope_theta = self.rope_parameters["rope_theta"] + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads if self.full_attn_idxs is None: diff --git a/mlx_lm/models/llama.py b/mlx_lm/models/llama.py index c12ddbe31..1663f6cf1 100644 --- a/mlx_lm/models/llama.py +++ b/mlx_lm/models/llama.py @@ -37,7 +37,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/llama4.py b/mlx_lm/models/llama4.py index 7149c99fc..d652f2642 100644 --- a/mlx_lm/models/llama4.py +++ b/mlx_lm/models/llama4.py @@ -42,6 +42,7 @@ class TextArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters @dataclass diff --git a/mlx_lm/models/llama4_text.py b/mlx_lm/models/llama4_text.py index 95a4c4d53..6acdea08c 100644 --- a/mlx_lm/models/llama4_text.py +++ b/mlx_lm/models/llama4_text.py @@ -31,6 +31,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/longcat_flash.py b/mlx_lm/models/longcat_flash.py index cf349058d..2865c8bfd 100644 --- a/mlx_lm/models/longcat_flash.py +++ b/mlx_lm/models/longcat_flash.py @@ -48,6 +48,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class LongcatFlashMLA(nn.Module): diff --git a/mlx_lm/models/longcat_flash_ngram.py b/mlx_lm/models/longcat_flash_ngram.py index 4e2e94829..c5bfda4bc 100644 --- a/mlx_lm/models/longcat_flash_ngram.py +++ b/mlx_lm/models/longcat_flash_ngram.py @@ -48,6 +48,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class NgramEmbedding(nn.Module): diff --git a/mlx_lm/models/mimo.py b/mlx_lm/models/mimo.py index 4ef9891ff..395584e05 100644 --- a/mlx_lm/models/mimo.py +++ b/mlx_lm/models/mimo.py @@ -31,7 +31,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Attention(nn.Module): diff --git a/mlx_lm/models/minicpm.py b/mlx_lm/models/minicpm.py index 2d6d96418..95573cbf6 100644 --- a/mlx_lm/models/minicpm.py +++ b/mlx_lm/models/minicpm.py @@ -33,7 +33,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class MLP(nn.Module): diff --git a/mlx_lm/models/minicpm3.py b/mlx_lm/models/minicpm3.py index 34141f932..22eeb781a 100644 --- a/mlx_lm/models/minicpm3.py +++ b/mlx_lm/models/minicpm3.py @@ -38,7 +38,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Attention(nn.Module): diff --git a/mlx_lm/models/mixtral.py b/mlx_lm/models/mixtral.py index e297f57fa..3f025bbe8 100644 --- a/mlx_lm/models/mixtral.py +++ b/mlx_lm/models/mixtral.py @@ -31,6 +31,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/nemotron-nas.py b/mlx_lm/models/nemotron-nas.py index 504c5affd..20fa17005 100644 --- a/mlx_lm/models/nemotron-nas.py +++ b/mlx_lm/models/nemotron-nas.py @@ -118,7 +118,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters # Automatically parse block_configs if they are loaded as dicts if self.block_configs and isinstance(self.block_configs[0], dict): diff --git a/mlx_lm/models/nemotron.py b/mlx_lm/models/nemotron.py index 349b97d18..429bb1cb9 100644 --- a/mlx_lm/models/nemotron.py +++ b/mlx_lm/models/nemotron.py @@ -34,7 +34,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.rope_scaling: if not "factor" in self.rope_scaling: raise ValueError(f"rope_scaling must contain 'factor'") diff --git a/mlx_lm/models/olmo.py b/mlx_lm/models/olmo.py index e06a04a48..6eb79fa85 100644 --- a/mlx_lm/models/olmo.py +++ b/mlx_lm/models/olmo.py @@ -34,7 +34,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) self.mlp_hidden_size = ( self.mlp_hidden_size diff --git a/mlx_lm/models/olmo2.py b/mlx_lm/models/olmo2.py index 002464296..01e1c79f4 100644 --- a/mlx_lm/models/olmo2.py +++ b/mlx_lm/models/olmo2.py @@ -33,7 +33,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/olmo3.py b/mlx_lm/models/olmo3.py index 3f5176751..e2f817fcc 100644 --- a/mlx_lm/models/olmo3.py +++ b/mlx_lm/models/olmo3.py @@ -34,7 +34,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads if self.layer_types is None: diff --git a/mlx_lm/models/olmoe.py b/mlx_lm/models/olmoe.py index e994a8e2c..b6479d434 100644 --- a/mlx_lm/models/olmoe.py +++ b/mlx_lm/models/olmoe.py @@ -36,7 +36,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/phi3.py b/mlx_lm/models/phi3.py index 779dc3ddb..6f35f7b28 100644 --- a/mlx_lm/models/phi3.py +++ b/mlx_lm/models/phi3.py @@ -32,7 +32,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/phimoe.py b/mlx_lm/models/phimoe.py index 2ff6c6afa..07388e948 100644 --- a/mlx_lm/models/phimoe.py +++ b/mlx_lm/models/phimoe.py @@ -31,7 +31,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 10000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/phixtral.py b/mlx_lm/models/phixtral.py index bb4768e31..4bfc987a4 100644 --- a/mlx_lm/models/phixtral.py +++ b/mlx_lm/models/phixtral.py @@ -24,6 +24,7 @@ class ModelArgs(BaseModelArgs): num_local_experts: int = 4 + class RoPEAttention(nn.Module): def __init__(self, dims: int, num_heads: int, rotary_dim: int): super().__init__() diff --git a/mlx_lm/models/plamo.py b/mlx_lm/models/plamo.py index 575d5655d..82facaf5d 100644 --- a/mlx_lm/models/plamo.py +++ b/mlx_lm/models/plamo.py @@ -28,6 +28,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Attention(nn.Module): diff --git a/mlx_lm/models/plamo2.py b/mlx_lm/models/plamo2.py index 6c81693d1..48b13317e 100644 --- a/mlx_lm/models/plamo2.py +++ b/mlx_lm/models/plamo2.py @@ -34,11 +34,6 @@ class ModelArgs(BaseModelArgs): mamba_enabled: bool = True intermediate_size: int = 13312 vocab_size: int = 32000 - rope_parameters: Optional[Dict] = None - - def __post_init__(self): - if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) class RMSNorm(nn.Module): diff --git a/mlx_lm/models/qwen.py b/mlx_lm/models/qwen.py index 945775b42..daeb91bba 100644 --- a/mlx_lm/models/qwen.py +++ b/mlx_lm/models/qwen.py @@ -23,11 +23,8 @@ class ModelArgs(BaseModelArgs): no_bias: bool = True vocab_size: int = 151936 num_key_value_heads = None - rope_parameters: Optional[Dict] = None def __post_init__(self): - if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/qwen2.py b/mlx_lm/models/qwen2.py index 8987274c3..d48a24636 100644 --- a/mlx_lm/models/qwen2.py +++ b/mlx_lm/models/qwen2.py @@ -31,7 +31,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Attention(nn.Module): diff --git a/mlx_lm/models/qwen2_moe.py b/mlx_lm/models/qwen2_moe.py index dd2708af4..97188a4a4 100644 --- a/mlx_lm/models/qwen2_moe.py +++ b/mlx_lm/models/qwen2_moe.py @@ -33,7 +33,9 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 1000000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/mlx_lm/models/qwen3.py b/mlx_lm/models/qwen3.py index 336134ed1..f6d9e1b5f 100644 --- a/mlx_lm/models/qwen3.py +++ b/mlx_lm/models/qwen3.py @@ -32,6 +32,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/qwen3_moe.py b/mlx_lm/models/qwen3_moe.py index 5202e6f72..660b5a135 100644 --- a/mlx_lm/models/qwen3_moe.py +++ b/mlx_lm/models/qwen3_moe.py @@ -37,6 +37,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Attention(nn.Module): diff --git a/mlx_lm/models/qwen3_next.py b/mlx_lm/models/qwen3_next.py index 689c1b187..0f3dd53ad 100644 --- a/mlx_lm/models/qwen3_next.py +++ b/mlx_lm/models/qwen3_next.py @@ -58,6 +58,7 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters @partial(mx.compile, shapeless=True) diff --git a/mlx_lm/models/rwkv7.py b/mlx_lm/models/rwkv7.py index 6f0ecdaea..26e57e861 100644 --- a/mlx_lm/models/rwkv7.py +++ b/mlx_lm/models/rwkv7.py @@ -25,11 +25,6 @@ class ModelArgs(BaseModelArgs): gate_low_rank_dim: int decay_low_rank_dim: int tie_word_embeddings: bool = False - rope_parameters: Optional[Dict] = None - - def __post_init__(self): - if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) @partial(mx.compile, shapeless=True) diff --git a/mlx_lm/models/seed_oss.py b/mlx_lm/models/seed_oss.py index f4fe23cae..a07bf5dca 100644 --- a/mlx_lm/models/seed_oss.py +++ b/mlx_lm/models/seed_oss.py @@ -35,6 +35,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) class Attention(nn.Module): diff --git a/mlx_lm/models/solar_open.py b/mlx_lm/models/solar_open.py index a3ed331d3..a11c8c3f4 100644 --- a/mlx_lm/models/solar_open.py +++ b/mlx_lm/models/solar_open.py @@ -41,3 +41,4 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters diff --git a/mlx_lm/models/step3p5.py b/mlx_lm/models/step3p5.py index d5615bb2f..b3e7aced6 100644 --- a/mlx_lm/models/step3p5.py +++ b/mlx_lm/models/step3p5.py @@ -65,7 +65,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class ZeroCenteredRMSNorm(nn.Module): diff --git a/mlx_lm/models/telechat3.py b/mlx_lm/models/telechat3.py index 18e0fd7f1..77681fd38 100644 --- a/mlx_lm/models/telechat3.py +++ b/mlx_lm/models/telechat3.py @@ -32,7 +32,8 @@ class ModelArgs(BaseModelArgs): def __post_init__(self): if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class Telechat3Attention(nn.Module): diff --git a/mlx_lm/models/youtu_llm.py b/mlx_lm/models/youtu_llm.py index 0cf41a615..44d4c4c21 100644 --- a/mlx_lm/models/youtu_llm.py +++ b/mlx_lm/models/youtu_llm.py @@ -36,8 +36,9 @@ class ModelArgs(BaseModelArgs): rope_parameters: Optional[Dict] = None def __post_init__(self): - if self.rope_parameters is not None: - self.rope_theta = self.rope_parameters.get("rope_theta", 500000.0) + if self.rope_parameters: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class YoutuLLMAttention(nn.Module): From 31ba659aded6e6b8c53962b9decc778d254cf8f5 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, 17 Mar 2026 14:38:43 +0100 Subject: [PATCH 89/91] Update acknowledgments to include Qwen3.5 and Qwen3.5MoE support --- ACKNOWLEDGMENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 964053048..77e33566a 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -16,7 +16,7 @@ Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite Mo Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`, Alibaba Qwen's `Qwen3Next`, Tele-AI's `TeleChat3`, and Allenai's `OLMoE` and `Olmo 3`; Helped add support for the following model architectures: -Alibaba Qwen's `Qwen3 & Qwen3MoE)`; Added support for the following training algorithms: +Alibaba Qwen's `Qwen3 & Qwen3MoE), Qwen3.5 & Qwen3.5MoE`; Added support for the following training algorithms: `Full Weight Fine-Tuning`, and the `Muon` optimizer; Added support for the following other features: `Multiple Optimizers to choose for training`, and `reporting training metrics to WandB (Weights & Biases)`. From c0ef19dd6c05d6010186db3f4c064dd75a999816 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, 17 Mar 2026 14:43:09 +0100 Subject: [PATCH 90/91] Fix installation message formatting in olmo.py --- mlx_lm/models/olmo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/models/olmo.py b/mlx_lm/models/olmo.py index 6eb79fa85..c257cdd66 100644 --- a/mlx_lm/models/olmo.py +++ b/mlx_lm/models/olmo.py @@ -13,7 +13,7 @@ try: import hf_olmo except ImportError: - print("To run olmo install ai2-olmo: pip install ai2-olmo") + print("To run olmo install ai2-olmo: `pip install ai2-olmo`") sys.exit(1) From e57b467f7b244b7a0d4629d37358cd295edb6b96 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, 17 Mar 2026 14:43:22 +0100 Subject: [PATCH 91/91] format --- mlx_lm/models/gemma3_text.py | 5 +++-- mlx_lm/models/gemma3n.py | 4 +++- mlx_lm/models/phixtral.py | 1 - 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/mlx_lm/models/gemma3_text.py b/mlx_lm/models/gemma3_text.py index 96761a911..4dbb9d567 100644 --- a/mlx_lm/models/gemma3_text.py +++ b/mlx_lm/models/gemma3_text.py @@ -36,8 +36,9 @@ def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) self.rope_scaling = self.rope_parameters - self.rope_local_base_freq = self.rope_parameters.get("rope_local_base_freq", 10000.0) - + self.rope_local_base_freq = self.rope_parameters.get( + "rope_local_base_freq", 10000.0 + ) class Attention(nn.Module): diff --git a/mlx_lm/models/gemma3n.py b/mlx_lm/models/gemma3n.py index ee3a4ace4..6138ff452 100644 --- a/mlx_lm/models/gemma3n.py +++ b/mlx_lm/models/gemma3n.py @@ -46,7 +46,9 @@ def __post_init__(self): if self.rope_parameters is not None: self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) self.rope_scaling = self.rope_parameters - self.rope_local_base_freq = self.rope_parameters.get("rope_local_base_freq", 10000.0) + self.rope_local_base_freq = self.rope_parameters.get( + "rope_local_base_freq", 10000.0 + ) @dataclass diff --git a/mlx_lm/models/phixtral.py b/mlx_lm/models/phixtral.py index 4bfc987a4..bb4768e31 100644 --- a/mlx_lm/models/phixtral.py +++ b/mlx_lm/models/phixtral.py @@ -24,7 +24,6 @@ class ModelArgs(BaseModelArgs): num_local_experts: int = 4 - class RoPEAttention(nn.Module): def __init__(self, dims: int, num_heads: int, rotary_dim: int): super().__init__()