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)`. diff --git a/mlx_lm/models/Klear.py b/mlx_lm/models/Klear.py index fb7a56029..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 +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): diff --git a/mlx_lm/models/afm7.py b/mlx_lm/models/afm7.py index 8a87b4509..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 @@ -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__( diff --git a/mlx_lm/models/afmoe.py b/mlx_lm/models/afmoe.py index e23dc237a..0fe6d838c 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,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): diff --git a/mlx_lm/models/apertus.py b/mlx_lm/models/apertus.py index 4d8afa28d..4ab530334 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,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): diff --git a/mlx_lm/models/baichuan_m1.py b/mlx_lm/models/baichuan_m1.py index 1a52b4f36..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 +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): diff --git a/mlx_lm/models/bailing_moe.py b/mlx_lm/models/bailing_moe.py index ff7d908ea..f0680b780 100644 --- a/mlx_lm/models/bailing_moe.py +++ b/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) diff --git a/mlx_lm/models/bailing_moe_linear.py b/mlx_lm/models/bailing_moe_linear.py index 360f99e1b..b8a1008d0 100644 --- a/mlx_lm/models/bailing_moe_linear.py +++ b/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( 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 diff --git a/mlx_lm/models/bitnet.py b/mlx_lm/models/bitnet.py index 72794e564..173518e1e 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,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): diff --git a/mlx_lm/models/cohere.py b/mlx_lm/models/cohere.py index c9b88885a..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 +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): diff --git a/mlx_lm/models/cohere2.py b/mlx_lm/models/cohere2.py index c93e238de..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 +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): diff --git a/mlx_lm/models/deepseek.py b/mlx_lm/models/deepseek.py index 2876f7b65..3c31df108 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,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): diff --git a/mlx_lm/models/deepseek_v2.py b/mlx_lm/models/deepseek_v2.py index a263d850a..ae45e6d5d 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,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( diff --git a/mlx_lm/models/deepseek_v3.py b/mlx_lm/models/deepseek_v3.py index a2766e59e..376032e49 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,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): diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index e40e52950..12379b6e0 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,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): diff --git a/mlx_lm/models/dots1.py b/mlx_lm/models/dots1.py index c568f39db..252a1dcad 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,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): diff --git a/mlx_lm/models/ernie4_5.py b/mlx_lm/models/ernie4_5.py index d6551f821..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 +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): diff --git a/mlx_lm/models/ernie4_5_moe.py b/mlx_lm/models/ernie4_5_moe.py index 597e9c467..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 +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): diff --git a/mlx_lm/models/exaone.py b/mlx_lm/models/exaone.py index 3cd223d29..9ace043ba 100644 --- a/mlx_lm/models/exaone.py +++ b/mlx_lm/models/exaone.py @@ -29,6 +29,13 @@ 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", 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 c0e15f5e2..fe9d0ae8b 100644 --- a/mlx_lm/models/exaone4.py +++ b/mlx_lm/models/exaone4.py @@ -29,6 +29,12 @@ 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", 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 a6da99be8..cec3401d5 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, Dict, List, Optional import mlx.core as mx import mlx.nn as nn @@ -43,10 +43,12 @@ 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", 100000.0) + self.rope_scaling = self.rope_parameters @mx.compile 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 diff --git a/mlx_lm/models/gemma.py b/mlx_lm/models/gemma.py index 9c540525e..d920a4919 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, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -22,6 +22,12 @@ 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", 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 23a98ac1a..5defceae5 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, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -25,6 +25,12 @@ 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", 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 5cca40e0a..4dbb9d567 100644 --- a/mlx_lm/models/gemma3_text.py +++ b/mlx_lm/models/gemma3_text.py @@ -23,13 +23,22 @@ 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", 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 9b380cf72..6138ff452 100644 --- a/mlx_lm/models/gemma3n.py +++ b/mlx_lm/models/gemma3n.py @@ -40,6 +40,15 @@ 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", 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/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): diff --git a/mlx_lm/models/glm4.py b/mlx_lm/models/glm4.py index 611175d29..177a8b6fb 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, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -26,6 +26,12 @@ 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", 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 bcf7cc9f4..73b11eb9f 100644 --- a/mlx_lm/models/glm4_moe.py +++ b/mlx_lm/models/glm4_moe.py @@ -44,6 +44,12 @@ 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", 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 dc108afdd..39d170939 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,12 @@ 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", 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 14e96e365..661ba5023 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): - self.rope_scaling = 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) + 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 e0fda5e27..5daf6eaa3 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, Dict, Optional 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,12 @@ 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", 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 6892445d7..2000a993b 100644 --- a/mlx_lm/models/granite.py +++ b/mlx_lm/models/granite.py @@ -31,6 +31,12 @@ 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", 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 e77cc6601..42fc6a3c1 100644 --- a/mlx_lm/models/granitemoe.py +++ b/mlx_lm/models/granitemoe.py @@ -32,6 +32,12 @@ 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", 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 a95c887e2..586edf096 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, Dict, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn @@ -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,16 +47,17 @@ 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) - # Mode flags - inferred from num_local_experts + 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) + @property def use_moe(self) -> bool: return bool(self.num_local_experts) diff --git a/mlx_lm/models/helium.py b/mlx_lm/models/helium.py index 238f205fa..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 +from typing import Any, Dict, Optional 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): diff --git a/mlx_lm/models/hunyuan.py b/mlx_lm/models/hunyuan.py index 0ee3a3a07..27069121c 100644 --- a/mlx_lm/models/hunyuan.py +++ b/mlx_lm/models/hunyuan.py @@ -33,8 +33,12 @@ 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", 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 d94edeaaf..4f4e9484e 100644 --- a/mlx_lm/models/hunyuan_v1_dense.py +++ b/mlx_lm/models/hunyuan_v1_dense.py @@ -27,6 +27,12 @@ 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", 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 8fd2301dd..d0e447d1e 100644 --- a/mlx_lm/models/internlm2.py +++ b/mlx_lm/models/internlm2.py @@ -26,8 +26,14 @@ 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) + 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 @@ -38,7 +44,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'" ) diff --git a/mlx_lm/models/internlm3.py b/mlx_lm/models/internlm3.py index 766fe977f..d77ae9ed0 100644 --- a/mlx_lm/models/internlm3.py +++ b/mlx_lm/models/internlm3.py @@ -27,8 +27,14 @@ 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) + 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 5fc387f2c..75154436a 100644 --- a/mlx_lm/models/iquestloopcoder.py +++ b/mlx_lm/models/iquestloopcoder.py @@ -51,6 +51,12 @@ 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", 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 bc27540fd..849b859a4 100644 --- a/mlx_lm/models/kimi_vl.py +++ b/mlx_lm/models/kimi_vl.py @@ -40,6 +40,12 @@ 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", 100000.0) + self.rope_scaling = self.rope_parameters @dataclass diff --git a/mlx_lm/models/lfm2.py b/mlx_lm/models/lfm2.py index c7f742af9..e2dedee09 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, Dict, List, Optional import mlx.core as mx import mlx.nn as nn @@ -36,10 +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 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.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/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 diff --git a/mlx_lm/models/lille-130m.py b/mlx_lm/models/lille-130m.py index b79057216..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 +from typing import Any, Dict, Optional 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): diff --git a/mlx_lm/models/llama.py b/mlx_lm/models/llama.py index 826bc9fc0..1663f6cf1 100644 --- a/mlx_lm/models/llama.py +++ b/mlx_lm/models/llama.py @@ -33,8 +33,14 @@ 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", 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 e4e284d71..d652f2642 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, Dict, Optional, Union import mlx.core as mx import mlx.nn as nn @@ -37,6 +37,12 @@ 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) + self.rope_scaling = self.rope_parameters @dataclass diff --git a/mlx_lm/models/llama4_text.py b/mlx_lm/models/llama4_text.py index 6f45a8557..6acdea08c 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, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -26,6 +26,12 @@ 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) + 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 fb3050a32..2865c8bfd 100644 --- a/mlx_lm/models/longcat_flash.py +++ b/mlx_lm/models/longcat_flash.py @@ -43,6 +43,12 @@ 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) + 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 9c0f2a8f7..c5bfda4bc 100644 --- a/mlx_lm/models/longcat_flash_ngram.py +++ b/mlx_lm/models/longcat_flash_ngram.py @@ -43,6 +43,12 @@ 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) + 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 d87a2d94f..395584e05 100644 --- a/mlx_lm/models/mimo.py +++ b/mlx_lm/models/mimo.py @@ -27,6 +27,13 @@ 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", 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/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): diff --git a/mlx_lm/models/minicpm.py b/mlx_lm/models/minicpm.py index 660db9051..95573cbf6 100644 --- a/mlx_lm/models/minicpm.py +++ b/mlx_lm/models/minicpm.py @@ -29,6 +29,13 @@ 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", 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 977f00eae..22eeb781a 100644 --- a/mlx_lm/models/minicpm3.py +++ b/mlx_lm/models/minicpm3.py @@ -34,6 +34,13 @@ 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", 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/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 diff --git a/mlx_lm/models/mixtral.py b/mlx_lm/models/mixtral.py index 18d16644c..3f025bbe8 100644 --- a/mlx_lm/models/mixtral.py +++ b/mlx_lm/models/mixtral.py @@ -26,8 +26,13 @@ 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) + 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/nanochat.py b/mlx_lm/models/nanochat.py index 1229e6ff9..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 +from typing import Any, Dict, Optional 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): diff --git a/mlx_lm/models/nemotron-nas.py b/mlx_lm/models/nemotron-nas.py index 09087c8e5..20fa17005 100644 --- a/mlx_lm/models/nemotron-nas.py +++ b/mlx_lm/models/nemotron-nas.py @@ -114,8 +114,13 @@ 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", 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): self.block_configs = [ diff --git a/mlx_lm/models/nemotron.py b/mlx_lm/models/nemotron.py index 5d612d477..429bb1cb9 100644 --- a/mlx_lm/models/nemotron.py +++ b/mlx_lm/models/nemotron.py @@ -30,8 +30,13 @@ 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) + 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/nemotron_h.py b/mlx_lm/models/nemotron_h.py index 0c391e42c..e3c67d8af 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 diff --git a/mlx_lm/models/olmo.py b/mlx_lm/models/olmo.py index bb7592d24..c257cdd66 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, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -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) @@ -30,8 +30,13 @@ 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", 100000.0) + self.rope_traditional = self.rope_parameters.get("rope_traditional", False) + self.mlp_hidden_size = ( self.mlp_hidden_size if self.mlp_hidden_size is not None diff --git a/mlx_lm/models/olmo2.py b/mlx_lm/models/olmo2.py index 4591e00dc..01e1c79f4 100644 --- a/mlx_lm/models/olmo2.py +++ b/mlx_lm/models/olmo2.py @@ -29,8 +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) 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 f87a818d8..e2f817fcc 100644 --- a/mlx_lm/models/olmo3.py +++ b/mlx_lm/models/olmo3.py @@ -30,8 +30,12 @@ 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", 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 c971c2247..b6479d434 100644 --- a/mlx_lm/models/olmoe.py +++ b/mlx_lm/models/olmoe.py @@ -32,8 +32,14 @@ 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) + if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads 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 diff --git a/mlx_lm/models/phi3.py b/mlx_lm/models/phi3.py index b3579eae8..6f35f7b28 100644 --- a/mlx_lm/models/phi3.py +++ b/mlx_lm/models/phi3.py @@ -28,8 +28,14 @@ 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", 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 c192d603b..07388e948 100644 --- a/mlx_lm/models/phimoe.py +++ b/mlx_lm/models/phimoe.py @@ -27,6 +27,12 @@ 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", 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 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): diff --git a/mlx_lm/models/plamo.py b/mlx_lm/models/plamo.py index 762a3c588..82facaf5d 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, Dict, Optional import mlx.core as mx import mlx.nn as nn @@ -23,6 +23,12 @@ 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) + 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 b835de0a5..48b13317e 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, 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 96e1a9718..daeb91bba 100644 --- a/mlx_lm/models/qwen.py +++ b/mlx_lm/models/qwen.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. from dataclasses import dataclass +from typing import Dict, Optional import mlx.core as mx import mlx.nn as nn diff --git a/mlx_lm/models/qwen2.py b/mlx_lm/models/qwen2.py index 45df85c7b..d48a24636 100644 --- a/mlx_lm/models/qwen2.py +++ b/mlx_lm/models/qwen2.py @@ -27,6 +27,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): diff --git a/mlx_lm/models/qwen2_moe.py b/mlx_lm/models/qwen2_moe.py index 8d6022d2e..97188a4a4 100644 --- a/mlx_lm/models/qwen2_moe.py +++ b/mlx_lm/models/qwen2_moe.py @@ -29,8 +29,14 @@ 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) + 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 a59343b68..f6d9e1b5f 100644 --- a/mlx_lm/models/qwen3.py +++ b/mlx_lm/models/qwen3.py @@ -27,6 +27,12 @@ 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) + 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 52dc50f9b..660b5a135 100644 --- a/mlx_lm/models/qwen3_moe.py +++ b/mlx_lm/models/qwen3_moe.py @@ -32,6 +32,12 @@ 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) + 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 39128a445..0f3dd53ad 100644 --- a/mlx_lm/models/qwen3_next.py +++ b/mlx_lm/models/qwen3_next.py @@ -53,6 +53,12 @@ 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) + self.rope_scaling = self.rope_parameters @partial(mx.compile, shapeless=True) diff --git a/mlx_lm/models/recurrent_gemma.py b/mlx_lm/models/recurrent_gemma.py index 4659d6c3d..71f2d3dc8 100644 --- a/mlx_lm/models/recurrent_gemma.py +++ b/mlx_lm/models/recurrent_gemma.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +from ast import Dict 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 diff --git a/mlx_lm/models/rwkv7.py b/mlx_lm/models/rwkv7.py index d698171ed..26e57e861 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 diff --git a/mlx_lm/models/seed_oss.py b/mlx_lm/models/seed_oss.py index d23c9f30b..a07bf5dca 100644 --- a/mlx_lm/models/seed_oss.py +++ b/mlx_lm/models/seed_oss.py @@ -30,6 +30,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): diff --git a/mlx_lm/models/smollm3.py b/mlx_lm/models/smollm3.py index 030924c7e..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 +from typing import Dict, Optional 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) diff --git a/mlx_lm/models/solar_open.py b/mlx_lm/models/solar_open.py index a0f8b8b5d..a11c8c3f4 100644 --- a/mlx_lm/models/solar_open.py +++ b/mlx_lm/models/solar_open.py @@ -36,3 +36,9 @@ 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) + self.rope_scaling = self.rope_parameters 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): diff --git a/mlx_lm/models/starcoder2.py b/mlx_lm/models/starcoder2.py index 4cbed554f..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 +from typing import Any, Dict, Optional 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): diff --git a/mlx_lm/models/step3p5.py b/mlx_lm/models/step3p5.py index 852823ea8..b3e7aced6 100644 --- a/mlx_lm/models/step3p5.py +++ b/mlx_lm/models/step3p5.py @@ -61,6 +61,12 @@ 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", 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 06a2d653a..77681fd38 100644 --- a/mlx_lm/models/telechat3.py +++ b/mlx_lm/models/telechat3.py @@ -28,6 +28,12 @@ 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", 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 9d8c653b1..44d4c4c21 100644 --- a/mlx_lm/models/youtu_llm.py +++ b/mlx_lm/models/youtu_llm.py @@ -33,6 +33,12 @@ 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: + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters class YoutuLLMAttention(nn.Module): diff --git a/tests/test_models.py b/tests/test_models.py index 0b2a963dc..f74909f69 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -425,7 +425,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