Reformat models for the rope parameter args

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)`.

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, Dict, List, Optional
 
 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):

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
@@ -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__(

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,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 class Attention(nn.Module):

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,13 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
+            self.rope_traditional = self.rope_parameters.get("rope_traditional", False)
 
 
 class ApertusMLP(nn.Module):

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, Dict, List, Optional
 
 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):

mlx_lm/models/bailing_moe.py

@@ -48,6 +48,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 @partial(mx.compile, shapeless=True)

mlx_lm/models/bailing_moe_linear.py

@@ -58,6 +58,13 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
+            self.rope_traditional = self.rope_parameters.get("rope_traditional", False)
 
 
 def recurrent_gla(

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

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,13 @@ 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)
+            self.rope_scaling = self.rope_parameters
+            self.rope_traditional = self.rope_parameters.get("rope_traditional", False)
 
 
 class Attention(nn.Module):

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, Dict, Optional
 
 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):

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 Dict, Optional, Tuple
 
 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):

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,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 class DeepseekAttention(nn.Module):

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,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 def yarn_find_correction_dim(

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,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 class DeepseekV3Attention(nn.Module):

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,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 class Indexer(nn.Module):

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,12 @@ 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", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 class Dots1Attention(nn.Module):

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, Dict, Optional
 
 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):

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, Dict, Optional
 
 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):