Fix Nemotron H loading error (#426)

* fix

* format

ACKNOWLEDGMENTS.md

@@ -8,5 +8,5 @@ with a short description of your contribution(s) below. For example:
 MLX LM was developed with contributions from the following individuals:
 
 - Shunta Saito: Added support for PLaMo models.
-- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, Rednote `dots.llm1`, Baisu's `Ernie4.5 MoE`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`, and `reporting training metrics to WandB (Weights & Biases)`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, Rednote `dots.llm1`, Baisu's `Ernie4.5 MoE`, inclusionAI's `Bailing MoE e.g. Ling-family`, IBM's `Granite MoE`, Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, and Allenai's `OLMoE`; 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)`.
 - Prince Canuma: Helped add support for the following model architectures: HuggingFace's `Starcoder2`, Cohere's `Cohere (1 and 2)`, Alibaba Qwen's `Qwen (2, 3 and MoE)`, Microsoft's `Phi (3 and 3.5 MoE)`, `BitNet1.58`, Meta's `Llama (3 and 4)`, Google DeepMind's `Gemma 3`, and InterLM's `InternLM 2.5`.

README.md

@@ -12,6 +12,7 @@ Some key features include:
   fine-tuning](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md)
   with support for quantized models.
 * Distributed inference and fine-tuning with `mx.distributed`
+* Supports experiment tracking using SwanLab and W&B.
 
 The easiest way to get started is to install the `mlx-lm` package:
 

mlx_lm/LEARNED_QUANTS.md

@@ -22,7 +22,7 @@ quantized model can be further refined with DWQ.
 To get started, first install the requirements:
 
 ```
-pip install mlx-lm[quant]
+pip install "mlx-lm[train]"
 ```
 
 ### DWQ

mlx_lm/LORA.md

@@ -26,6 +26,12 @@ LoRA (QLoRA).[^qlora] LoRA fine-tuning works with the following model families:
 
 ## Run
 
+First, make sure you have the training dependenices installed:
+
+```shell
+pip install "mlx-lm[train]"
+```
+
 The main command is `mlx_lm.lora`. To see a full list of command-line options run:
 
 ```shell
@@ -78,8 +84,14 @@ You can resume fine-tuning with an existing adapter with
 
 #### Logging
 
-You can log training metrics to Weights & Biases by passing a project name with
-the `--wandb` flag. Make sure to install wandb with `pip install wandb`.
+You can log training metrics to Weights & Biases using `--report-to wandb`, or
+to SwanLab using `--report-to swanlab`. Make sure to install the required
+packages beforehand: `pip install wandb` or `pip install swanlab`. You can
+enable both tracking tools simultaneously by separating them with a comma, for
+example: `--report-to wandb,swanlab`.
+
+To specify a project name for the logging tracker, use `--project-name <YOUR
+PROJECT NAME>`.
 
 #### Prompt Masking
 

mlx_lm/__main__.py

@@ -9,6 +9,7 @@ if __name__ == "__main__":
         "quant.dwq",
         "quant.dynamic_quant",
         "quant.gptq",
+        "benchmark",
         "cache_prompt",
         "chat",
         "convert",
@@ -16,6 +17,7 @@ if __name__ == "__main__":
         "fuse",
         "generate",
         "lora",
+        "perplexity",
         "server",
         "manage",
         "upload",

mlx_lm/_version.py

@@ -1,3 +1,3 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.26.0"
+__version__ = "0.27.1"

mlx_lm/benchmark.py

@@ -0,0 +1,106 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+
+from mlx_lm import stream_generate
+from mlx_lm.generate import DEFAULT_MODEL
+from mlx_lm.tokenizer_utils import load_tokenizer
+from mlx_lm.utils import (
+    fetch_from_hub,
+    get_model_path,
+)
+
+
+def setup_arg_parser():
+    """Set up and return the argument parser."""
+    parser = argparse.ArgumentParser(description="LLM benchmarking script")
+    parser.add_argument(
+        "--model",
+        type=str,
+        help=(
+            "The path to the local model directory or Hugging Face repo. "
+            f"If no model is specified, then {DEFAULT_MODEL} is used."
+        ),
+        default=None,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
+    parser.add_argument(
+        "--prompt-tokens",
+        "-p",
+        default=512,
+        help="Length of prompt",
+        type=int,
+    )
+    parser.add_argument(
+        "--generation-tokens",
+        "-g",
+        default=1024,
+        help="Length of completion",
+        type=int,
+    )
+    parser.add_argument(
+        "--num-trials",
+        "-n",
+        default=5,
+        help="Number of timing trials",
+        type=int,
+    )
+    return parser
+
+
+def main():
+    parser = setup_arg_parser()
+    args = parser.parse_args()
+    mx.random.seed(0)
+
+    model_path = args.model or DEFAULT_MODEL
+
+    model_path, _ = get_model_path(model_path, revision=None)
+    model, config, _ = fetch_from_hub(model_path, trust_remote_code=True)
+    tokenizer = load_tokenizer(
+        model_path,
+        eos_token_ids=[],  # Empty to avoid early stopping
+        tokenizer_config_extra={"trust_remote_code": True},
+    )
+
+    prompt_tokens = args.prompt_tokens
+    generation_tokens = args.generation_tokens
+    prompt = mx.random.randint(0, config["vocab_size"], (prompt_tokens,))
+
+    def _bench():
+        for response in stream_generate(
+            model, tokenizer, prompt, max_tokens=generation_tokens
+        ):
+            pass
+        return response
+
+    print("Running warmup..")
+    _bench()
+
+    report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
+    print(f"Timing with {prompt_tokens=} and {generation_tokens=}.")
+    responses = []
+    for i in range(args.num_trials):
+        response = _bench()
+        responses.append(response)
+        results = [(k, getattr(response, k)) for k in report_keys]
+        results = [f"{k}={v:.3f}" for k, v in results]
+        print(f"Trial {i+1}:  " + ", ".join(results))
+
+    def avg(k):
+        vals = (getattr(response, k) for response in responses)
+        return sum(vals) / args.num_trials
+
+    results = [(k, avg(k)) for k in report_keys]
+    results = [f"{k}={v:.3f}" for k, v in results]
+    print(f"Averages: " + ", ".join(results))
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/chat.py

@@ -27,6 +27,11 @@ def setup_arg_parser():
         help="The path to the local model directory or Hugging Face repo.",
         default=DEFAULT_MODEL,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -69,6 +74,11 @@ def setup_arg_parser():
         default=DEFAULT_MAX_TOKENS,
         help="Maximum number of tokens to generate",
     )
+    parser.add_argument(
+        "--system-prompt",
+        default=None,
+        help="System prompt to be used for the chat template",
+    )
     return parser
 
 
@@ -82,7 +92,9 @@ def main():
     model, tokenizer = load(
         args.model,
         adapter_path=args.adapter_path,
-        tokenizer_config={"trust_remote_code": True},
+        tokenizer_config={
+            "trust_remote_code": True if args.trust_remote_code else None
+        },
     )
 
     def print_help():
@@ -104,7 +116,10 @@ def main():
         if query == "h":
             print_help()
             continue
-        messages = [{"role": "user", "content": query}]
+        messages = []
+        if args.system_prompt is not None:
+            messages.append({"role": "system", "content": args.system_prompt})
+        messages.append({"role": "user", "content": query})
         prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
         for response in stream_generate(
             model,

mlx_lm/convert.py

@@ -49,7 +49,6 @@ def mixed_quant_predicate_builder(
     def mixed_quant_predicate(
         path: str,
         module: nn.Module,
-        config: dict,
     ) -> Union[bool, dict]:
         """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
         Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
@@ -88,6 +87,7 @@ def convert(
     quantize: bool = False,
     q_group_size: int = 64,
     q_bits: int = 4,
+    q_mode: str = "affine",
     dtype: Optional[str] = None,
     upload_repo: str = None,
     revision: Optional[str] = None,
@@ -137,7 +137,12 @@ def convert(
     if quantize:
         print("[INFO] Quantizing")
         model, config = quantize_model(
-            model, config, q_group_size, q_bits, quant_predicate=quant_predicate
+            model,
+            config,
+            q_group_size,
+            q_bits,
+            mode=q_mode,
+            quant_predicate=quant_predicate,
         )
 
     if dequantize:
@@ -183,6 +188,13 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
+    parser.add_argument(
+        "--q-mode",
+        help="The quantization mode.",
+        type=str,
+        default="affine",
+        choices=["affine", "mxfp4"],
+    )
     parser.add_argument(
         "--quant-predicate",
         help=f"Mixed-bit quantization recipe.",
@@ -210,6 +222,12 @@ def configure_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        help="Trust remote code when loading tokenizer.",
+        action="store_true",
+        default=False,
+    )
     return parser
 
 

mlx_lm/evaluate.py

@@ -20,6 +20,7 @@ import mlx.nn as nn
 import numpy as np
 from lm_eval.api.model import LM
 from lm_eval.api.registry import register_model
+from lm_eval.models import huggingface
 from tqdm import tqdm
 
 from .generate import stream_generate
@@ -62,7 +63,7 @@ def chat_template_fn(**extra_kwargs):
 @register_model("mlxlm")
 class MLXLM(LM):
 
-    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
+    tokenizer_name = huggingface.HFLM.tokenizer_name
     apply_chat_template = chat_template_fn()
 
     def __init__(
@@ -70,9 +71,13 @@ class MLXLM(LM):
         path_or_hf_repo: str,
         max_tokens: Optional[int] = None,
         use_chat_template: Optional[bool] = None,
+        trust_remote_code: bool = False,
     ) -> None:
         super().__init__()
-        self._model, self.tokenizer = load(path_or_hf_repo)
+        tokenizer_config = {"trust_remote_code": True if trust_remote_code else None}
+        self._model, self.tokenizer = load(
+            path_or_hf_repo, tokenizer_config=tokenizer_config
+        )
         self._max_tokens = max_tokens or self.tokenizer.model_max_length
         self._batch_size = 8
         self.use_chat_template = use_chat_template
@@ -165,7 +170,7 @@ class MLXLM(LM):
         indices = []
         for v in group_reqs.values():
             idx, resp = zip(*v)
-            indices.extend(idx)
+            indices.append(idx)
             responses.append(resp)
 
         # split data accross ranks
@@ -211,31 +216,36 @@ class MLXLM(LM):
                 scores[-1] += mx.sum(score).item()
                 is_greedy[-1] &= mx.all(ig).item()
 
-        scores = mx.array(scores)
-        is_greedy = mx.array(is_greedy)
-
         if long_completions > 0:
             logging.info(
                 f"Prefix eliminated for {long_completions} requests with "
                 + "completion longer than context."
             )
 
+        # All gather the results across nodes
         num_results = len(requests)
+        per_group = mx.distributed.all_max(len(scores), stream=mx.cpu).item()
+        scores = scores + [0] * (per_group - len(scores))
+        is_greedy = is_greedy + [False] * (per_group - len(is_greedy))
+        scores = mx.array(scores)
+        is_greedy = mx.array(is_greedy)
+        scores = mx.distributed.all_gather(scores, stream=mx.cpu)
+        is_greedy = mx.distributed.all_gather(is_greedy, stream=mx.cpu)
+        mx.eval(scores, is_greedy)
 
-        # all gather the results across groups
-        if group.size() > 1:
-            per_group = int(np.ceil(num_results / group.size()))
-            scores = mx.pad(scores, ((0, per_group - len(scores)),))
-            is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
-            scores = mx.distributed.all_gather(scores[mx.newaxis], stream=mx.cpu)
-            is_greedy = mx.distributed.all_gather(is_greedy[mx.newaxis], stream=mx.cpu)
-            mx.eval(scores, is_greedy)
-            scores = scores.T.reshape(-1)
-            is_greedy = is_greedy.T.reshape(-1)
-
-        inv_sort = mx.argsort(mx.array(indices))
+        # Arrange the indices to match the scores from each node and then
+        # inverse sort the scores
+        all_indices = []
+        for rank in range(group.size()):
+            rank_indices = [
+                idx for question in indices[rank :: group.size()] for idx in question
+            ]
+            rank_indices += [num_results] * (per_group - len(rank_indices))
+            all_indices.extend(rank_indices)
+        inv_sort = mx.argsort(mx.array(all_indices))
         scores = scores[:num_results][inv_sort]
         is_greedy = is_greedy[:num_results][inv_sort]
+
         return list(zip(scores.tolist(), is_greedy.tolist()))
 
     def loglikelihood_rolling(self, requests) -> list[float]:
@@ -275,8 +285,8 @@ class MLXLM(LM):
         )
         inputs = self._tokenize([req.args[0] for req in requests])
         all_scores = []
-        for i in tqdm(range(0, len(texts), self._batch_size)):
-            batch = texts[i : i + self._batch_size]
+        for i in tqdm(range(0, len(inputs), self._batch_size)):
+            batch = inputs[i : i + self._batch_size]
             scores, lengths, _ = self._score_fn(batch)
             mask = mx.arange(scores.shape[-1]) < lengths[:, None]
             all_scores.extend((mask * scores).sum(axis=-1).tolist())
@@ -371,6 +381,17 @@ def main():
         apply_chat_template, e.g. '{"enable_thinking":false}'""",
         default="{}",
     )
+    parser.add_argument(
+        "--confirm-run-unsafe-code",
+        action="store_true",
+        help="Confirm that you want to run tasks that execute untrusted code.",
+        default=False,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
 
     args = parser.parse_args()
 
@@ -382,10 +403,17 @@ def main():
 
     mx.random.seed(args.seed)
 
+    # Initialize the communication if in distributed mode
+    world = mx.distributed.init()
+    mx.eval(mx.distributed.all_sum(1, stream=mx.cpu))
+    if world.size() > 1 and world.rank() == 0:
+        print(f"Evaluating with {world.size()} nodes")
+
     lm = MLXLM(
         args.model,
         max_tokens=args.max_tokens,
         use_chat_template=args.apply_chat_template,
+        trust_remote_code=args.trust_remote_code,
     )
     MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
 
@@ -400,6 +428,7 @@ def main():
         numpy_random_seed=args.seed,
         torch_random_seed=args.seed,
         fewshot_random_seed=args.seed,
+        confirm_run_unsafe_code=args.confirm_run_unsafe_code,
     )
 
     file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
@@ -407,7 +436,7 @@ def main():
         file_keys += [f"{args.num_shots:02d}"]
     file_keys += args.tasks
     filename = "_".join(file_keys)
-    if mx.distributed.init().rank() == 0:
+    if world.rank() == 0:
         output_path = output_dir / filename
         output_path.write_text(json.dumps(results["results"], indent=4))
         print("Results:")

mlx_lm/examples/lora_config.yaml

@@ -37,8 +37,9 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5
 
-# Whether to report the logs to WandB
-# wand: "wandb-project"
+# Services to report logs to (comma-separated): wandb, swanlab, or both ('wandb,swanlab').
+# report_to: wandb,swanlab
+# project_name: "Your-awesome-mlx-project-name"
 
 # Number of training steps between loss reporting.
 steps_per_report: 10

mlx_lm/examples/pipeline_generate.py

@@ -50,7 +50,7 @@ def shard_and_load(repo):
 
     # Lazy load and shard model to figure out
     # which weights we need
-    model, _ = load_model(model_path, lazy=True, strict=False)
+    model, config = load_model(model_path, lazy=True, strict=False)
 
     group = mx.distributed.init()
     rank = group.rank()
@@ -68,7 +68,11 @@ def shard_and_load(repo):
     download(args.model, allow_patterns=local_files)
 
     # Load and shard the model, and load the weights
-    tokenizer = load_tokenizer(model_path)
+    tokenizer = load_tokenizer(
+        model_path,
+        {"trust_remote_code": True},
+        eos_token_ids=config.get("eos_token_id", None),
+    )
     model, _ = load_model(model_path, lazy=True, strict=False)
     model.model.pipeline(group)
     mx.eval(model.parameters())

mlx_lm/generate.py

@@ -61,6 +61,11 @@ def setup_arg_parser():
         ),
         default=None,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -218,31 +223,31 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
         try:
             yield
         finally:
-            return
-
-    model_bytes = tree_reduce(
-        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
-    )
-    max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
-    if model_bytes > 0.9 * max_rec_size:
-        model_mb = model_bytes // 2**20
-        max_rec_mb = max_rec_size // 2**20
-        print(
-            f"[WARNING] Generating with a model that requires {model_mb} MB "
-            f"which is close to the maximum recommended size of {max_rec_mb} "
-            "MB. This can be slow. See the documentation for possible work-arounds: "
-            "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            pass
+    else:
+        model_bytes = tree_reduce(
+            lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
         )
-    old_limit = mx.set_wired_limit(max_rec_size)
-    try:
-        yield
-    finally:
-        if streams is not None:
-            for s in streams:
-                mx.synchronize(s)
-        else:
-            mx.synchronize()
-        mx.set_wired_limit(old_limit)
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        if model_bytes > 0.9 * max_rec_size:
+            model_mb = model_bytes // 2**20
+            max_rec_mb = max_rec_size // 2**20
+            print(
+                f"[WARNING] Generating with a model that requires {model_mb} MB "
+                f"which is close to the maximum recommended size of {max_rec_mb} "
+                "MB. This can be slow. See the documentation for possible work-arounds: "
+                "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            )
+        old_limit = mx.set_wired_limit(max_rec_size)
+        try:
+            yield
+        finally:
+            if streams is not None:
+                for s in streams:
+                    mx.synchronize(s)
+            else:
+                mx.synchronize()
+            mx.set_wired_limit(old_limit)
 
 
 @dataclass
@@ -329,21 +334,25 @@ def generate_step(
            when ``kv_bits`` is non-None. Default: ``0``.
         prompt_progress_callback (Callable[int, int]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
-        input_embeddings (mx.array, optional): Input embeddings to use in conjunction
-           with prompt tokens. Default: ``None``.
+        input_embeddings (mx.array, optional): Input embeddings to use instead of or in
+          conjunction with prompt tokens. Default: ``None``.
 
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
-    if len(prompt) == 0:
-        raise ValueError("Prompt must be non-empty.")
     if input_embeddings is not None:
         if not does_model_support_input_embeddings(model):
             raise ValueError("Model does not support input embeddings.")
-        elif prompt.shape[0] != input_embeddings.shape[0]:
+        elif len(prompt) > 0 and len(prompt) != len(input_embeddings):
             raise ValueError(
-                "If using input embeddings, the sequence length must match that of the prompt."
+                f"When providing input_embeddings, their sequence length ({len(input_embeddings)}) "
+                f"must match the sequence length of the prompt ({len(prompt)}), or the "
+                "prompt must be empty."
             )
+    elif len(prompt) == 0:
+        raise ValueError(
+            "Either input_embeddings or prompt (or both) must be provided."
+        )
 
     tokens = None
 
@@ -386,7 +395,7 @@ def generate_step(
 
             logits = logits[:, -1, :]
 
-            if logits_processors:
+            if logits_processors and len(input_tokens) > 0:
                 tokens = (
                     mx.concat([tokens, input_tokens])
                     if tokens is not None
@@ -402,24 +411,28 @@ def generate_step(
             return sampled, logprobs.squeeze(0)
 
     with mx.stream(generation_stream):
-        total_prompt_tokens = prompt.shape[0]
+        total_prompt_tokens = (
+            len(input_embeddings) if input_embeddings is not None else len(prompt)
+        )
         prompt_processed_tokens = 0
-        while total_prompt_tokens - prompt_processed_tokens > prefill_step_size:
+        prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
+        while total_prompt_tokens - prompt_processed_tokens > 1:
+            n_to_process = min(prefill_step_size, prompt.size - 1)
             _model_call(
-                input_tokens=prompt[:prefill_step_size][None],
+                input_tokens=prompt[:n_to_process][None],
                 input_embeddings=(
-                    input_embeddings[:prefill_step_size][None]
+                    input_embeddings[:n_to_process][None]
                     if input_embeddings is not None
                     else None
                 ),
             )
             quantize_cache_fn(prompt_cache)
             mx.eval([c.state for c in prompt_cache])
+            prompt_processed_tokens += n_to_process
             prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            prompt = prompt[prefill_step_size:]
+            prompt = prompt[n_to_process:]
             input_embeddings = (
-                input_embeddings[prefill_step_size:]
+                input_embeddings[n_to_process:]
                 if input_embeddings is not None
                 else input_embeddings
             )
@@ -668,6 +681,7 @@ def stream_generate(
         )
     else:
         kwargs.pop("max_kv_size", None)
+        kwargs.pop("prompt_progress_callback", None)
         token_generator = speculative_generate_step(
             prompt, model, draft_model, **kwargs
         )
@@ -792,7 +806,7 @@ def main():
     tokenizer_config = (
         {} if not using_cache else json.loads(metadata["tokenizer_config"])
     )
-    tokenizer_config["trust_remote_code"] = True
+    tokenizer_config["trust_remote_code"] = True if args.trust_remote_code else None
 
     model_path = args.model
     if using_cache:

mlx_lm/lora.py

@@ -3,6 +3,7 @@ import math
 import os
 import re
 import types
+import warnings
 from pathlib import Path
 
 import mlx.core as mx
@@ -11,7 +12,7 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml
 
-from .tuner.callbacks import WandBCallback
+from .tuner.callbacks import get_reporting_callbacks
 from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
@@ -46,6 +47,9 @@ CONFIG_DEFAULTS = {
     "optimizer_config": {
         "adam": {},
         "adamw": {},
+        "muon": {},
+        "sgd": {},
+        "adafactor": {},
     },
     "data": "data/",
     "seed": 0,
@@ -67,7 +71,9 @@ CONFIG_DEFAULTS = {
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
     "mask_prompt": False,
-    "wandb": None,
+    "wandb": None,  # will be deprecated in a future release
+    "report_to": None,
+    "project_name": None,
 }
 
 
@@ -103,9 +109,9 @@ def build_parser():
     parser.add_argument(
         "--optimizer",
         type=str,
-        choices=["adam", "adamw"],
+        choices=["adam", "adamw", "muon", "sgd", "adafactor"],
         default=None,
-        help="Optimizer to use for training: adam or adamw",
+        help="Optimizer to use for training: adam, adamw, sgd, or adafactor.",
     )
     parser.add_argument(
         "--mask-prompt",
@@ -179,11 +185,26 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
-    parser.add_argument(
+    parser.add_argument(  # will be deprecated in a future release
         "--wandb",
         type=str,
         default=None,
-        help="WandB project name to report training metrics. Disabled if None.",
+        help=(
+            "The 'wandb' argument is deprecated and will be removed in a future release. "
+            "Use 'report_to: wandb' and 'project_name' in the configuration instead."
+        ),
+    )
+    parser.add_argument(
+        "--report-to",
+        type=str,
+        default=None,
+        help="Services to report logs to ('wandb', 'swanlab', or 'wandb,swanlab').",
+    )
+    parser.add_argument(
+        "--project-name",
+        type=str,
+        default=None,
+        help="Project name for logging. Defaults to the name of the root directory.",
     )
     parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
@@ -251,11 +272,16 @@ def train_model(
 
     optimizer_name = args.optimizer.lower()
     optimizer_config = args.optimizer_config.get(optimizer_name, {})
-
     if optimizer_name == "adam":
         opt_class = optim.Adam
     elif optimizer_name == "adamw":
         opt_class = optim.AdamW
+    elif optimizer_name == "muon":
+        opt_class = optim.Muon
+    elif optimizer_name == "sgd":
+        opt_class = optim.SGD
+    elif optimizer_name == "adafactor":
+        opt_class = optim.Adafactor
     else:
         raise ValueError(f"Unsupported optimizer: {optimizer_name}")
 
@@ -288,17 +314,23 @@ def evaluate_model(args, model: nn.Module, test_set):
 
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
-
     if args.wandb is not None:
-        training_callback = WandBCallback(
-            project_name=args.wandb,
-            log_dir=args.adapter_path,
-            config=vars(args),
-            wrapped_callback=training_callback,
+        warnings.warn(
+            "The 'wandb' argument is deprecated and will be removed in a future release. "
+            "Use 'report_to: wandb' and 'project_name' in the configuration instead.",
+            DeprecationWarning,
         )
+        args.report_to = "wandb"
+        args.project_name = args.wandb
+    training_callback = get_reporting_callbacks(
+        args.report_to,
+        project_name=args.project_name,
+        log_dir=args.adapter_path,
+        config=vars(args),
+    )
 
     print("Loading pretrained model")
-    model, tokenizer = load(args.model)
+    model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})
 
     print("Loading datasets")
     train_set, valid_set, test_set = load_dataset(args, tokenizer)

mlx_lm/models/apertus.py

@@ -0,0 +1,222 @@
+# 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
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    mlp_bias: bool
+    num_attention_heads: int
+    attention_bias: bool
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    rope_theta: float
+    post_norm: bool
+    qk_norm: bool
+    tie_word_embeddings: bool
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):
+    alpha_p = nn.softplus(alpha_p)
+    alpha_n = beta + nn.softplus(alpha_n)
+    return mx.where(
+        x > 0,
+        alpha_p * mx.square(x) + beta * x,
+        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
+    )
+
+
+class XieLU(nn.Module):
+    def __init__(
+        self,
+        alpha_p_init=0.8,
+        alpha_n_init=0.8,
+        beta=0.5,
+        eps=-1e-6,
+    ):
+        super().__init__()
+        alpha_p_tensor = mx.array(alpha_p_init)
+        alpha_n_tensor = mx.array(alpha_n_init - beta)
+        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
+        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
+
+        self.beta = mx.array(beta)
+        self.eps = mx.array(eps)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)
+
+
+class ApertusMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.up_proj = nn.Linear(
+            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+        )
+        self.act_fn = XieLU()
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+
+class ApertusAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+
+        self.head_dim = args.hidden_size // args.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            args.hidden_size, args.num_attention_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size, args.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size, args.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            args.num_attention_heads * self.head_dim, args.hidden_size, bias=False
+        )
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        queries = self.q_norm(
+            queries.reshape(B, L, self.num_attention_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        keys = self.k_norm(keys.reshape(B, L, self.num_key_value_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class ApertusDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = ApertusAttention(args)
+        self.mlp = ApertusMLP(args)
+
+        self.attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.feedforward_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = x + self.self_attn(self.attention_layernorm(x), mask, cache)
+        out = h + self.mlp(self.feedforward_layernorm(h))
+        return out
+
+
+class ApertusModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            ApertusDecoderLayer(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask=mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = ApertusModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if k.endswith("alpha_p") or k.endswith("alpha_n"):
+                weights[k] = v.squeeze()
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/bailing_moe.py

@@ -0,0 +1,322 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    moe_intermediate_size: int
+    num_experts: int
+    num_shared_experts: int
+    norm_topk_prob: bool
+    num_attention_heads: int
+    num_experts_per_tok: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    vocab_size: int
+    first_k_dense_replace: int
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    use_bias: bool = False
+    use_qkv_bias: bool = False
+    norm_head: bool = False
+    norm_softmax: bool = False
+
+
+class BailingMoeMLP(nn.Module):
+    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+        self.intermediate_size = (
+            intermediate_size
+            if intermediate_size is not None
+            else args.intermediate_size
+        )
+
+        self.gate_proj = nn.Linear(
+            args.hidden_size, self.intermediate_size, bias=args.use_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, args.hidden_size, bias=args.use_bias
+        )
+        self.up_proj = nn.Linear(
+            args.hidden_size, self.intermediate_size, bias=args.use_bias
+        )
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class BailingMoeAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.head_dim = args.hidden_size // self.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.query_key_value = nn.Linear(
+            args.hidden_size,
+            (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=args.use_qkv_bias,
+        )
+        self.dense = nn.Linear(
+            self.num_attention_heads * self.head_dim,
+            args.hidden_size,
+            bias=args.use_bias,
+        )
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            traditional=False,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        qkv = self.query_key_value(x)
+
+        q_size = self.num_attention_heads * self.head_dim
+        kv_size = self.num_key_value_heads * self.head_dim
+        q, k, v = mx.split(qkv, [q_size, q_size + kv_size], axis=-1)
+
+        queries = q.reshape(B, L, self.num_attention_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = k.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        values = v.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.dense(output)
+
+
+class BailingMoeGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.norm_topk_prob = config.norm_topk_prob
+        self.gating_dim = config.hidden_size
+
+        self.gate_proj = nn.Linear(self.gating_dim, self.num_experts, bias=False)
+
+    def __call__(self, hidden_states):
+        B, L, D = hidden_states.shape
+        x = hidden_states.reshape(-1, D)
+
+        logits = self.gate_proj(x)
+        scores = mx.softmax(logits, axis=-1, precise=True)
+
+        topk_idx = mx.argpartition(scores, kth=-self.top_k, axis=-1)[..., -self.top_k :]
+        topk_scores = mx.take_along_axis(scores, topk_idx, axis=-1)
+
+        if self.top_k > 1 and self.norm_topk_prob:
+            denom = mx.sum(topk_scores, axis=-1, keepdims=True)
+            topk_scores = topk_scores / mx.maximum(denom, 1e-9)
+
+        return topk_idx, topk_scores
+
+
+class BailingMoeSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_experts_per_tok = args.num_experts_per_tok
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+            bias=args.use_bias,
+        )
+
+        self.gate = BailingMoeGate(config=args)
+
+        if args.num_shared_experts > 0:
+            self.shared_experts = BailingMoeMLP(
+                args=args,
+                intermediate_size=args.moe_intermediate_size * args.num_shared_experts,
+            )
+        else:
+            self.shared_experts = None
+
+    def __call__(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+
+        if self.shared_experts is not None:
+            identity = hidden_states
+
+        x = hidden_states.reshape(-1, hidden_dim)
+
+        expert_indices, expert_weights = self.gate(hidden_states)
+        expert_outputs = self.switch_mlp(x, expert_indices)
+
+        weighted_output = mx.sum(expert_outputs * expert_weights[..., None], axis=-2)
+        output = weighted_output.reshape(batch_size, seq_len, hidden_dim)
+
+        if self.shared_experts is not None:
+            output = output + self.shared_experts(hidden_states)
+
+        return output
+
+
+class BailingMoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.attention = BailingMoeAttention(args)
+
+        self.mlp = (
+            BailingMoeSparseMoeBlock(args)
+            if (
+                args.num_experts is not None and layer_idx >= args.first_k_dense_replace
+            )
+            else BailingMoeMLP(args)
+        )
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.attention(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class BailingMoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            BailingMoeDecoderLayer(args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        h = self.word_embeddings(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.norm_head = args.norm_head
+        self.model_type = args.model_type
+        self.model = BailingMoeModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.model(inputs, mask, cache)
+        return self.lm_head(h)
+
+    def sanitize(self, weights):
+        if self.norm_head:
+            w = weights["lm_head.weight"]
+            dtype = w.dtype
+            weight_norm = (
+                mx.linalg.norm(w.astype(mx.float32), axis=0, keepdims=True) + 1e-7
+            )
+            weights["lm_head.weight"] = (w / weight_norm).astype(dtype)
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+
+            if l >= self.args.first_k_dense_replace:
+                for m in ["gate_proj", "down_proj", "up_proj"]:
+                    for k in ["weight", "scales", "biases"]:
+                        if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                            to_join = [
+                                weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                                for e in range(self.args.num_experts)
+                            ]
+                            weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(
+                                to_join
+                            )
+
+                if f"{prefix}.mlp.gate.weight" in weights:
+                    gate_weight = weights.pop(f"{prefix}.mlp.gate.weight")
+                    weights[f"{prefix}.mlp.gate.gate_proj.weight"] = gate_weight
+
+                if f"{prefix}.mlp.gate.bias" in weights:
+                    gate_bias = weights.pop(f"{prefix}.mlp.gate.bias")
+                    weights[f"{prefix}.mlp.gate.gate_proj.bias"] = gate_bias
+
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/base.py

@@ -35,7 +35,7 @@ def create_causal_mask(
     rinds = rinds[None]
     mask = linds >= rinds
     if window_size is not None:
-        mask = mask & (linds <= rinds + window_size)
+        mask = mask & (linds < rinds + window_size)
     if lengths is not None:
         lengths = lengths[:, None, None, None]
         mask = mask & (rinds < lengths)

mlx_lm/models/cache.py

@@ -453,9 +453,9 @@ class RotatingKVCache(_BaseCache):
         raise NotImplementedError("RotatingKVCache Quantization NYI")
 
 
-class MambaCache(_BaseCache):
-    def __init__(self):
-        self.cache = [None, None]
+class ArraysCache(_BaseCache):
+    def __init__(self, size):
+        self.cache = [None] * size
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -472,6 +472,11 @@ class MambaCache(_BaseCache):
         self.cache = v
 
 
+class MambaCache(ArraysCache):
+    def __init__(self):
+        super().__init__(size=2)
+
+
 class ChunkedKVCache(KVCache):
     def __init__(self, chunk_size=None):
         super().__init__()

mlx_lm/models/deepseek_v3.py

@@ -124,20 +124,6 @@ class DeepseekV3YarnRotaryEmbedding(nn.Module):
         )
 
 
-# A clipped silu to prevent fp16 from overflowing
-@partial(mx.compile, shapeless=True)
-def clipped_silu(x):
-    return mx.clip(x * mx.sigmoid(x), a_min=-100, a_max=100)
-
-
-class ClippedSilu(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def __call__(self, x):
-        return clipped_silu(x)
-
-
 class DeepseekV3Attention(nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
@@ -303,7 +289,9 @@ def group_expert_select(
     group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
     k = n_group - topk_group
     group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-    scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
+    scores = mx.put_along_axis(
+        scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+    )
     scores = mx.flatten(scores, -2, -1)
 
     k = top_k
@@ -352,7 +340,6 @@ class DeepseekV3MoE(nn.Module):
             config.hidden_size,
             config.moe_intermediate_size,
             config.n_routed_experts,
-            activation=ClippedSilu(),
         )
 
         self.gate = MoEGate(config)

mlx_lm/models/exaone4.py

@@ -0,0 +1,208 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    rope_theta: float
+    head_dim: int
+    tie_word_embeddings: bool
+    rope_scaling: Dict[str, Union[float, str]]
+    sliding_window: Optional[int]
+    sliding_window_pattern: Optional[str]
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, is_local: Optional[bool]):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.is_local = is_local or False
+        self.use_rope = is_local is None or is_local
+        if self.use_rope:
+            self.rope = initialize_rope(
+                head_dim,
+                base=args.rope_theta,
+                traditional=False,
+                scaling_config=args.rope_scaling,
+                max_position_embeddings=args.max_position_embeddings,
+            )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            if self.use_rope:
+                queries = self.rope(queries, offset=cache.offset)
+                keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        elif self.use_rope:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs, is_local: bool):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args, is_local)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(x, mask, cache)
+        h = x + self.post_attention_layernorm(r)
+        r = self.mlp(h)
+        out = h + self.post_feedforward_layernorm(r)
+        return out
+
+
+class ExaoneModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        pattern = args.sliding_window_pattern
+        self.layers = [
+            TransformerBlock(
+                args=args,
+                is_local=pattern[i % len(pattern)] == "L" if pattern else None,
+            )
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = ExaoneModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                if l.self_attn.is_local
+                else KVCache()
+            )
+            for l in self.layers
+        ]
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gemma3_text.py

@@ -215,6 +215,7 @@ class Model(nn.Module):
         self.model_type = args.model_type
         self.model = Gemma3Model(args)
         self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.tie_word_embeddings = False
 
     def __call__(
         self,
@@ -224,13 +225,16 @@ class Model(nn.Module):
         input_embeddings: Optional[mx.array] = None,
     ):
         out = self.model(inputs, mask, cache, input_embeddings)
-        out = self.lm_head(out)
+        if self.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
         return out
 
     def sanitize(self, weights):
-        weights = dict(weights)
         if "lm_head.weight" not in weights:
-            weights["lm_head.weight"] = weights["model.embed_tokens.weight"]
+            self.tie_word_embeddings = True
+            self.pop("lm_head")
         return weights
 
     @property

mlx_lm/models/gemma3n.py

@@ -25,7 +25,6 @@ class TextConfig(BaseModelArgs):
     vocab_size: int
     num_key_value_heads: int
     num_kv_shared_layers: int
-    query_pre_attn_scalar: float
     vocab_size_per_layer_input: int
     sliding_window: int
     max_position_embeddings: int
@@ -177,7 +176,11 @@ class MLP(nn.Module):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
+        self.intermediate_size = (
+            config.intermediate_size[layer_idx]
+            if isinstance(config.intermediate_size, list)
+            else config.intermediate_size
+        )
         self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
@@ -473,7 +476,7 @@ class LanguageModel(nn.Module):
         per_layer_inputs = self.project_per_layer_inputs(h, per_layer_inputs)
 
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = self.make_cache()
 
         if mask is None:
             full_mask = create_attention_mask(

mlx_lm/models/glm4_moe.py

@@ -0,0 +1,329 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    moe_intermediate_size: int
+    norm_topk_prob: bool
+    num_attention_heads: int
+    n_group: int
+    head_dim: int
+    topk_group: int
+    n_shared_experts: int
+    n_routed_experts: int
+    routed_scaling_factor: float
+    num_experts_per_tok: int
+    first_k_dense_replace: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    rope_scaling: Optional[Dict]
+    use_qk_norm: bool
+    tie_word_embeddings: bool
+    attention_bias: bool
+    partial_rotary_factor: float
+    scoring_func: str = "sigmoid"
+    topk_method: str = "noaux_tc"
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.use_qk_norm = args.use_qk_norm
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+            self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+
+        self.rope = nn.RoPE(
+            int(head_dim * args.partial_rotary_factor),
+            traditional=False,
+            base=args.rope_theta,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1)
+
+        if self.use_qk_norm:
+            queries = self.q_norm(queries)
+            keys = self.k_norm(keys)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
+        self.intermediate_size = (
+            config.intermediate_size if intermediate_size is None else intermediate_size
+        )
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def __call__(self, x):
+        down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + e_score_correction_bias
+    if n_group > 1:
+        k = top_k
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        denominator = scores.sum(axis=-1, keepdims=True)
+        scores = scores / denominator
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+        assert config.topk_method == "noaux_tc", "Unsupported topk method."
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class MoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+            self.shared_experts = MLP(
+                config=config, intermediate_size=intermediate_size
+            )
+
+    def __call__(self, x):
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(config)
+        self.mlp = (
+            MoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+            )
+            else MLP(config)
+        )
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            DecoderLayer(config, idx) for idx in range(config.num_hidden_layers)
+        ]
+        self.start_idx = 0
+        self.end_idx = len(self.layers)
+        self.num_layers = self.end_idx
+
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * self.num_layers
+
+        for i in range(self.num_layers):
+            h = self.layers[self.start_idx + i](h, mask, cache[i])
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = LanguageModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, cache, mask)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        mpt_layer = self.args.num_hidden_layers
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        # Remove multi-token prediction layer
+        return {
+            k: v
+            for k, v in weights.items()
+            if not k.startswith(f"model.layers.{mpt_layer}")
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/gpt_oss.py

@@ -0,0 +1,379 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_causal_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "gpt_oss"
+    num_hidden_layers: int = 36
+    num_local_experts: int = 128
+    num_experts_per_tok: int = 4
+    vocab_size: int = 201088
+    rms_norm_eps: float = 1e-05
+    hidden_size: int = 2880
+    intermediate_size: int = 2880
+    head_dim: int = 64
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 8
+    sliding_window: int = 128
+    rope_theta: int = 150000
+    rope_scaling: Any = None
+    layer_types: list = None
+
+
+# These operators emulate particular methods in torch that don't exist in MLX natively
+def mlx_topk(a, k, axis=-1):
+    """MLX equivalent of torch.topk"""
+    partitioned_indices = mx.argpartition(a, kth=-k, axis=axis)
+    # Extract only the top k indices (last k elements after partition)
+    top_k_indices = partitioned_indices[..., -k:]
+    # Get the corresponding values
+    top_k_values = mx.take_along_axis(a, top_k_indices, axis=axis)
+    return top_k_values, top_k_indices
+
+
+@partial(mx.compile, shapeless=True)
+def swiglu(x_linear, x_glu, alpha: float = 1.702, limit: float = 7.0):
+    # Clamp the input values
+    x_glu = mx.clip(x_glu, a_min=None, a_max=limit)
+    x_linear = mx.clip(x_linear, a_min=-limit, a_max=limit)
+
+    glu_scaled = alpha * x_glu
+    sig = mx.sigmoid(glu_scaled)
+
+    out_glu = x_glu * sig
+    # Note we add an extra bias of 1 to the linear layer
+    return out_glu * (x_linear + 1)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x, gate):
+        return swiglu(x, gate)
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.head_dim = config.head_dim
+        self.num_attention_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+
+        self.sinks = mx.zeros((config.num_attention_heads,))
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True
+        )
+
+        self.o_proj = nn.Linear(
+            self.head_dim * config.num_attention_heads, config.hidden_size, bias=True
+        )
+
+        self.sm_scale = 1 / math.sqrt(config.head_dim)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            config.rope_theta,
+            traditional=False,
+            scaling_config=config.rope_scaling,
+        )
+
+        # Cache the mask so we don't have to create it every time
+        self._previous_mask = None
+
+    def get_causal_mask(self, x, cache):
+        _, L, _ = x.shape
+        offset = cache.offset if cache is not None else 0
+        offset = max(1, offset)
+
+        def _make_mask(L, offset):
+            zero = mx.array(0, dtype=x.dtype)
+            neginf = mx.array(-mx.inf, dtype=x.dtype)
+            mask = mx.where(create_causal_mask(L, offset - 1), zero, neginf)
+            mask = mask.reshape(1, 1, L, -1)
+            mask = mx.tile(mask, (1, self.num_attention_heads, 1, 1))
+            sinks = mx.tile(self.sinks.reshape(1, -1, 1, 1), (1, 1, L, 1))
+            mask = mx.concatenate([sinks, mask], axis=-1)
+            return mask
+
+        # When training re-create the mask so that gradients flow to the sinks.
+        # When L is large then recreate the mask because otherwise it will take
+        # a pretty significant chunk of memory.
+        if self.training or L > 8:
+            self._previous_mask = None
+            return _make_mask(L, offset)
+
+        # Create the mask once and try to reuse it. For this reason we round up
+        # to the closest multiple of 512 so we can reuse the mask several times.
+        length = ((L + offset + 511) // 512) * 512
+        if (
+            self._previous_mask is None
+            or self._previous_mask.shape[-1] < length
+            or self._previous_mask.shape[-2] != L
+        ):
+            self._previous_mask = _make_mask(L, length - L)
+
+        return self._previous_mask[..., : L + offset]
+
+    def get_sliding_window_mask(self, x, cache, window_size):
+        _, L, _ = x.shape
+        offset = cache.offset if cache is not None else 0
+        offset = max(1, offset)
+
+        def _make_mask(L, offset):
+            zero = mx.array(0, dtype=x.dtype)
+            neginf = mx.array(-mx.inf, dtype=x.dtype)
+            mask = create_causal_mask(L, offset - 1, window_size)
+            mask = mx.where(mask, zero, neginf)
+            mask = mask.reshape(1, 1, L, -1)
+            mask = mx.tile(mask, (1, self.num_attention_heads, 1, 1))
+            sinks = mx.tile(self.sinks.reshape(1, -1, 1, 1), (1, 1, L, 1))
+            mask = mx.concatenate([sinks, mask], axis=-1)
+            return mask
+
+        # If we are training then simply re-create the mask every time to make
+        # sure gradients flow to the sinks.
+        #
+        # For simplicity also re-create the mask if we have more than 1 query
+        # for now.
+        if self.training or L > 1:
+            self._previous_mask = None
+            return _make_mask(L, min(window_size + 1, offset))
+
+        # We are in inference so cache the mask and try to reuse it
+        if self._previous_mask is None:
+            self._previous_mask = _make_mask(L, window_size)
+
+        return self._previous_mask[..., : min(L + offset, window_size + 1)]
+
+    def get_mask(self, x, cache, window_size):
+        if window_size is not None:
+            return self.get_sliding_window_mask(x, cache, window_size)
+        else:
+            return self.get_causal_mask(x, cache)
+
+    def __call__(self, x: mx.array, mask: mx.array, cache=None) -> mx.array:
+        B, L, _ = x.shape
+        D = self.head_dim
+        Hk = self.num_key_value_heads
+
+        q = self.q_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+        k = self.k_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+        v = self.v_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+
+        # If cache is None or the cache offset is 0 then we need to add a 0 key
+        # and value to make some space for the sink
+        if cache is None or cache.offset == 0:
+            q = self.rope(q)
+            k = self.rope(k)
+
+            zeros = mx.zeros((B, Hk, 1, D), dtype=k.dtype)
+            k = mx.concatenate([zeros, k], axis=2)
+            v = mx.concatenate([zeros, v], axis=2)
+            if cache is not None:
+                k, v = cache.update_and_fetch(k, v)
+
+        # We have already put the 0 in the cache no need to do anything special
+        else:
+            q = self.rope(q, offset=cache.offset - 1)
+            k = self.rope(k, offset=cache.offset - 1)
+            k, v = cache.update_and_fetch(k, v)
+
+        # NOTE: mask should contain the sink weights already
+        v_hat = scaled_dot_product_attention(q, k, v, cache, self.sm_scale, mask=mask)
+
+        return self.o_proj(v_hat.swapaxes(1, 2).reshape(B, L, -1))
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_local_experts = config.num_local_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        self.experts = SwitchGLU(
+            input_dims=config.hidden_size,
+            hidden_dims=config.intermediate_size,
+            num_experts=config.num_local_experts,
+            activation=SwiGLU(),
+            bias=True,
+        )
+        self.router = nn.Linear(config.hidden_size, config.num_local_experts, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        g = self.router(x)
+        experts, indices = mlx_topk(g, k=self.num_experts_per_tok, axis=-1)
+        expert_weights = mx.softmax(experts, axis=-1, precise=True)
+
+        # Experts block
+        x = self.experts(x, indices)
+
+        x = x * mx.expand_dims(expert_weights, axis=-1)
+        return x.sum(axis=-2)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.self_attn = AttentionBlock(config)
+        self.mlp = MLPBlock(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+
+    def __call__(self, x: mx.array, mask: mx.array, cache=None) -> mx.array:
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, mask, cache)
+        x = residual + x
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        return x
+
+
+class GptOssMoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+        self.layer_types = args.layer_types or [
+            "sliding_attention",
+            "full_attention",
+        ] * (args.num_hidden_layers // 2)
+        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+        self.window_size = args.sliding_window
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            x = input_embeddings
+        else:
+            x = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        if mask is None:
+            masks = [
+                l.self_attn.get_mask(
+                    x, c, self.window_size if lt == "sliding_attention" else None
+                )
+                for (l, c, lt) in zip(self.layers, cache, self.layer_types)
+            ]
+        else:
+            masks = [mask] * len(self.layers)
+
+        for i, (layer, c, m) in enumerate(zip(self.layers, cache, masks)):
+            x = layer(x, m, c)
+        x = self.norm(x)
+        return x
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GptOssMoeModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs: mx.array, mask: mx.array = None, cache=None):
+        return self.lm_head(self.model(inputs, mask, cache))
+
+    def sanitize(self, weights):
+        if any("gate_proj.weight" in k for k in weights.keys()):
+            return weights  # already sanitized
+
+        new_weights = {}
+        for k, v in weights.items():
+            if "gate_up_proj" in k and "bias" not in k:
+                if "_blocks" in k:
+                    v = v.view(mx.uint32).flatten(-2)
+                    k = k.replace("_blocks", ".weight")
+                if "_scales" in k:
+                    k = k.replace("_scales", ".scales")
+                new_weights[k.replace("gate_up_proj", "gate_proj")] = mx.contiguous(
+                    v[..., ::2, :]
+                )
+                new_weights[k.replace("gate_up_proj", "up_proj")] = mx.contiguous(
+                    v[..., 1::2, :]
+                )
+            elif "down_proj" in k and "bias" not in k:
+                if "_blocks" in k:
+                    v = v.view(mx.uint32).flatten(-2)
+                    k = k.replace("_blocks", ".weight")
+                if "_scales" in k:
+                    k = k.replace("_scales", ".scales")
+                new_weights[k] = v
+            elif "gate_up_proj_bias" in k:
+                new_weights[k.replace("gate_up_proj_bias", "gate_proj.bias")] = (
+                    mx.contiguous(v[..., ::2])
+                )
+                new_weights[k.replace("gate_up_proj_bias", "up_proj.bias")] = (
+                    mx.contiguous(v[..., 1::2])
+                )
+            elif "down_proj_bias" in k:
+                new_weights[k.replace("down_proj_bias", "down_proj.bias")] = v
+            else:
+                new_weights[k] = v
+
+        return new_weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("router"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    def make_cache(self):
+        caches = []
+        for lt in self.model.layer_types:
+            if lt == "full_attention":
+                caches.append(KVCache())
+            else:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window + 1, keep=1)
+                )
+        return caches

mlx_lm/models/granitemoe.py

@@ -0,0 +1,237 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    logits_scaling: float
+    attention_multiplier: float
+    embedding_multiplier: float
+    residual_multiplier: float
+    max_position_embeddings: int
+    num_key_value_heads: int
+    attention_bias: bool
+    rope_theta: float
+    num_local_experts: int
+    num_experts_per_tok: int
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+
+
+class GraniteMoeAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+
+        self.scale = args.attention_multiplier
+        attention_bias = args.attention_bias
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            False,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class GraniteMoeTopKGating(nn.Module):
+    def __init__(self, input_size: int, num_experts: int, top_k: int):
+        super().__init__()
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.top_k = top_k
+        self.layer = nn.Linear(input_size, num_experts, bias=False)
+
+    def __call__(self, hidden_states: mx.array):
+        logits = self.layer(hidden_states)
+        top_k_idx = mx.argpartition(logits, kth=-self.top_k, axis=-1)[
+            ..., -self.top_k :
+        ]
+        top_k_logits = mx.take_along_axis(logits, top_k_idx, axis=-1)
+        top_k_gates = mx.softmax(top_k_logits.astype(mx.float32), axis=-1)
+        return top_k_idx, top_k_gates
+
+
+class GraniteMoeMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.input_size = args.hidden_size
+        self.hidden_size = args.intermediate_size
+        self.switch_mlp = SwitchGLU(
+            self.input_size, self.hidden_size, args.num_local_experts
+        )
+        self.router = GraniteMoeTopKGating(
+            input_size=self.input_size,
+            num_experts=args.num_local_experts,
+            top_k=args.num_experts_per_tok,
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        token_ids, gates = self.router(x)
+        y = self.switch_mlp(x, token_ids)
+        return (y * gates[..., None]).sum(axis=-2).astype(y.dtype)
+
+
+class GraniteMoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = GraniteMoeAttention(args)
+        self.block_sparse_moe = GraniteMoeMoE(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.residual_multiplier = args.residual_multiplier
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r * self.residual_multiplier
+        r = self.block_sparse_moe(self.post_attention_layernorm(h))
+        out = h + r * self.residual_multiplier
+        return out
+
+
+class GraniteMoEModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            GraniteMoeDecoderLayer(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.embedding_multiplier = args.embedding_multiplier
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs) * self.embedding_multiplier
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GraniteMoEModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.logits_scaling = args.logits_scaling
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out / self.logits_scaling
+
+    def sanitize(self, weights):
+        if "model.layers.0.block_sparse_moe.input_linear.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}.block_sparse_moe"
+            key = f"{prefix}.input_linear.weight"
+            value = weights.pop(key)
+            gate_proj, up_proj = mx.split(value, 2, axis=1)
+            weights[key.replace("input_linear", "switch_mlp.gate_proj")] = gate_proj
+            weights[key.replace("input_linear", "switch_mlp.up_proj")] = up_proj
+            key = f"{prefix}.output_linear.weight"
+            weights[key.replace("output_linear", "switch_mlp.down_proj")] = weights.pop(
+                key
+            )
+
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("block_sparse_moe.router.layer"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/hunyuan_v1_dense.py

@@ -0,0 +1,224 @@
+# Copyright © 2023-2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float = 10000
+    max_position_embeddings: int = 32768
+    attention_bias: bool = False
+    use_qk_norm: bool = True
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    head_dim: Optional[int] = None
+
+    def __post_init__(self):
+        if self.rope_scaling:
+            required_keys = {"alpha", "factor", "type"}
+            if not all(key in self.rope_scaling for key in required_keys):
+                raise ValueError(f"rope_scaling must contain keys {required_keys}")
+
+
+class DynamicNTKAlphaRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        base: float = 10000,
+        scaling_alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.dims = dims
+        base = base * scaling_alpha ** (dims / (dims - 2))
+        self._freqs = base ** (mx.arange(0, self.dims, 2) / self.dims)
+
+    def __call__(self, x, offset: int = 0):
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=False,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = (
+            args.head_dim if args.head_dim is not None else args.hidden_size // n_heads
+        )
+        self.head_dim = head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
+
+        self.use_qk_norm = args.use_qk_norm
+        if self.use_qk_norm:
+            self.query_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
+            self.key_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
+
+        scaling_alpha = 1.0
+        if args.rope_scaling and "alpha" in args.rope_scaling:
+            scaling_alpha = args.rope_scaling["alpha"]
+
+        self.rope = DynamicNTKAlphaRoPE(
+            head_dim,
+            base=args.rope_theta,
+            scaling_alpha=scaling_alpha,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        if self.use_qk_norm:
+            queries = self.query_layernorm(queries)
+            keys = self.key_layernorm(keys)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class HunyuanV1DenseModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = HunyuanV1DenseModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.model.embed_tokens.as_linear(out)
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/lfm2-vl.py

@@ -0,0 +1,57 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import lfm2
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+        self.text_config["full_attn_idxs"] = [
+            i
+            for i, layer_type in enumerate(self.text_config["layer_types"])
+            if layer_type == "full_attention"
+        ]
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = lfm2.Model(lfm2.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def make_cache(self):
+        return self.language_model.make_cache()

mlx_lm/models/lfm2.py

@@ -0,0 +1,278 @@
+# Copyright © 2025 Apple Inc.
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import ArraysCache, KVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    norm_eps: float
+    conv_bias: bool
+    conv_L_cache: int
+    block_dim: int
+    block_ff_dim: int
+    block_multiple_of: int
+    block_ffn_dim_multiplier: float
+    block_auto_adjust_ff_dim: bool
+    full_attn_idxs: List[int]
+    rope_theta: float
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+
+        self.q_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps)
+        self.k_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps)
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.rope = nn.RoPE(
+            self.head_dim,
+            base=args.rope_theta,
+            traditional=False,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = self.q_layernorm(queries.reshape(B, L, self.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_layernorm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, mask=mask, scale=self.scale
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+class ShortConv(nn.Module):
+    def __init__(
+        self,
+        args: ModelArgs,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.args = args
+        self.layer_idx = layer_idx
+        self.L_cache = args.conv_L_cache
+        self.bias = args.conv_bias
+
+        self.conv = nn.Conv1d(
+            in_channels=args.hidden_size,
+            out_channels=args.hidden_size,
+            kernel_size=self.L_cache,
+            groups=args.hidden_size,
+            bias=self.bias,
+        )
+        self.in_proj = nn.Linear(args.hidden_size, 3 * args.hidden_size, bias=self.bias)
+        self.out_proj = nn.Linear(args.hidden_size, args.hidden_size, bias=self.bias)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        seqlen = x.shape[1]
+        BCx = self.in_proj(x)
+        B, C, x = mx.split(BCx, 3, axis=-1)
+        Bx = B * x
+
+        state = None
+        if cache is not None:
+            state = cache[0]
+        if state is None:
+            state = mx.zeros(
+                (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), dtype=Bx.dtype
+            )
+
+        Bx = mx.concatenate([state, Bx], axis=-2)
+        if cache is not None:
+            cache[0] = Bx[:, -(self.L_cache - 1) :]
+        conv_out = self.conv(Bx)
+
+        y = C * conv_out
+        return self.out_proj(y)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        ff_dim: int,
+        multiple_of: int,
+        auto_adjust_ff_dim: bool,
+        ffn_dim_multiplier: Optional[float],
+    ):
+        super().__init__()
+        if auto_adjust_ff_dim:
+            ff_dim = int(2 * ff_dim / 3)
+            if ffn_dim_multiplier is not None:
+                ff_dim = int(ffn_dim_multiplier * ff_dim)
+            ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, ff_dim, bias=False)
+        self.w3 = nn.Linear(dim, ff_dim, bias=False)
+        self.w2 = nn.Linear(ff_dim, dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.w2(nn.silu(self.w1(x)) * self.w3(x))
+
+
+class Lfm2DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.is_attention_layer = layer_idx in args.full_attn_idxs
+
+        if self.is_attention_layer:
+            self.self_attn = Attention(args)
+        else:
+            self.conv = ShortConv(args, layer_idx)
+        self.feed_forward = MLP(
+            dim=args.block_dim,
+            ff_dim=args.block_ff_dim,
+            multiple_of=args.block_multiple_of,
+            auto_adjust_ff_dim=args.block_auto_adjust_ff_dim,
+            ffn_dim_multiplier=args.block_ffn_dim_multiplier,
+        )
+
+        self.operator_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+        self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+
+        if self.is_attention_layer:
+            r = self.self_attn(self.operator_norm(x), mask=mask, cache=cache)
+        else:
+            r = self.conv(
+                self.operator_norm(x),
+                cache=cache,
+            )
+        h = x + r
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out
+
+
+class Lfm2Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Lfm2DecoderLayer(args, layer_idx=i) for i in range(args.num_hidden_layers)
+        ]
+
+        self.embedding_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+
+        if mask is None:
+            first_attn_idx = self.args.full_attn_idxs[0]
+            c = [cache[first_attn_idx]] if cache is not None else None
+            mask = create_attention_mask(h, c)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.embedding_norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Lfm2Model(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, mask, cache, input_embeddings)
+        return self.model.embed_tokens.as_linear(out)
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for name, param in weights.items():
+            if "conv.weight" in name:
+                if param.shape[-1] > param.shape[1]:
+                    param = param.transpose(0, 2, 1)
+
+            sanitized_weights[name] = param
+        return sanitized_weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            KVCache() if l.is_attention_layer else ArraysCache(size=1)
+            for l in self.layers
+        ]

mlx_lm/models/longcat_flash.py

@@ -0,0 +1,383 @@
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import CacheList, KVCache
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    attention_method: str
+    zero_expert_type: str
+    hidden_size: int
+    ffn_hidden_size: int
+    moe_topk: int
+    expert_ffn_hidden_size: int
+    n_routed_experts: int
+    zero_expert_num: int
+    num_layers: int
+    vocab_size: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    qk_nope_head_dim: int
+    v_head_dim: int
+    routed_scaling_factor: float
+    rms_norm_eps: float
+    rope_theta: float
+    mla_scale_q_lora: bool
+    mla_scale_kv_lora: bool
+    attention_bias: bool
+    norm_topk_prob: bool = False
+    router_bias: bool = False
+
+
+class LongcatFlashMLA(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.qk_rope_head_dim = args.qk_rope_head_dim
+        self.qk_nope_head_dim = args.qk_nope_head_dim
+        self.kv_lora_rank = args.kv_lora_rank
+        self.q_lora_rank = args.q_lora_rank
+        self.v_head_dim = args.v_head_dim
+
+        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
+        self.scale = self.qk_head_dim**-0.5
+
+        if self.q_lora_rank is None:
+            self.q_proj = nn.Linear(
+                args.hidden_size,
+                self.num_attention_heads * self.qk_head_dim,
+                bias=False,
+            )
+        else:
+            self.q_a_proj = nn.Linear(
+                args.hidden_size, self.q_lora_rank, bias=args.attention_bias
+            )
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank,
+                self.num_attention_heads * self.qk_head_dim,
+                bias=False,
+            )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            args.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=args.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_attention_heads * (self.qk_nope_head_dim + args.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_attention_heads * args.v_head_dim,
+            args.hidden_size,
+            bias=args.attention_bias,
+        )
+
+        if args.mla_scale_q_lora:
+            self.mla_scale_q_lora = (args.hidden_size / self.q_lora_rank) ** 0.5
+        if args.mla_scale_kv_lora:
+            self.mla_scale_kv_lora = (args.hidden_size / self.kv_lora_rank) ** 0.5
+
+        self.rope = nn.RoPE(
+            dims=self.qk_rope_head_dim, base=args.rope_theta, traditional=True
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        if self.q_lora_rank is None:
+            q_states = self.q_proj(x)
+        else:
+            q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+
+        q_states = q_states.reshape(B, L, -1, self.qk_head_dim).transpose(0, 2, 1, 3)
+
+        if self.mla_scale_q_lora is not None:
+            q_states = q_states * self.mla_scale_q_lora
+
+        q_pass, q_rot = mx.split(q_states, [self.qk_nope_head_dim], axis=-1)
+
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        k_pass, k_rot = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pass = self.kv_a_layernorm(k_pass)
+
+        if self.mla_scale_kv_lora is not None:
+            k_pass = k_pass * self.mla_scale_kv_lora
+
+        key_shape = (B, L, -1, self.qk_nope_head_dim + self.v_head_dim)
+        k_pass = self.kv_b_proj(k_pass).reshape(*key_shape).transpose(0, 2, 1, 3)
+        k_pass, value_states = mx.split(k_pass, [self.qk_nope_head_dim], axis=-1)
+
+        k_rot = k_rot.reshape(B, 1, L, self.qk_rope_head_dim)
+
+        if cache is not None:
+            q_rot = self.rope(q_rot, cache.offset)
+            k_rot = self.rope(k_rot, cache.offset)
+        else:
+            q_rot = self.rope(q_rot)
+            k_rot = self.rope(k_rot)
+
+        k_rot = mx.broadcast_to(k_rot, (*k_pass.shape[:-1], k_rot.shape[-1]))
+
+        query_states = mx.concatenate([q_pass, q_rot], axis=-1)
+        key_states = mx.concatenate([k_pass, k_rot], axis=-1)
+
+        if cache is not None:
+            key_states, value_states = cache.update_and_fetch(key_states, value_states)
+
+        attn_output = scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,
+        )
+
+        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(attn_output)
+
+
+class LongcatFlashMLP(nn.Module):
+    def __init__(self, args: ModelArgs, is_expert: bool = False):
+        super().__init__()
+        hidden_size = args.expert_ffn_hidden_size if is_expert else args.ffn_hidden_size
+
+        self.gate_proj = nn.Linear(args.hidden_size, hidden_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, hidden_size, bias=False)
+        self.down_proj = nn.Linear(hidden_size, args.hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LongcatFlashTopkRouter(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.config = args
+        self.top_k = args.moe_topk
+        self.n_routed_experts = args.n_routed_experts + args.zero_expert_num
+        self.routed_scaling_factor = args.routed_scaling_factor
+        self.norm_topk_prob = args.norm_topk_prob
+        self.router_bias = args.router_bias
+
+        self.classifier = nn.Linear(
+            args.hidden_size, self.n_routed_experts, bias=self.router_bias
+        )
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+
+    def __call__(self, hidden_states: mx.array) -> Tuple[mx.array, mx.array]:
+
+        dtype = hidden_states.dtype
+        router_logits = self.classifier(hidden_states)
+        scores = mx.softmax(router_logits, axis=-1)
+
+        corrected_scores = scores + self.e_score_correction_bias
+        topk_indices = mx.argpartition(corrected_scores, kth=-self.top_k, axis=-1)[
+            ..., -self.top_k :
+        ]
+        topk_weights = mx.take_along_axis(scores, topk_indices, axis=-1)
+
+        if self.norm_topk_prob:
+            denominator = mx.sum(topk_weights, axis=-1, keepdims=True) + 1e-20
+            topk_weights = topk_weights / denominator
+
+        topk_weights = topk_weights * self.routed_scaling_factor
+
+        return topk_indices, topk_weights.astype(dtype)
+
+
+class LongcatFlashMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.config = args
+        self.num_experts_per_tok = args.moe_topk
+        self.n_routed_experts = args.n_routed_experts
+        self.zero_expert_num = args.zero_expert_num
+        self.zero_expert_type = args.zero_expert_type
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.expert_ffn_hidden_size,
+            args.n_routed_experts,
+        )
+
+        self.router = LongcatFlashTopkRouter(args)
+
+    def __call__(self, hidden_states):
+
+        topk_indices, topk_weights = self.router(hidden_states)
+
+        # Process all regular experts at once
+        mask = topk_indices >= self.n_routed_experts
+        topk_indices = mx.where(mask, 0, topk_indices)
+        regular_weights = mx.where(mask, 0.0, topk_weights)
+
+        regular_outputs = self.switch_mlp(hidden_states, topk_indices)
+
+        weighted_outputs = regular_outputs * topk_weights[..., None]
+
+        # Add identity expert contribution if needed
+        assert self.zero_expert_type == "identity"
+        identity_weights = mx.where(mask, topk_weights, 0.0)
+        identity_outputs = hidden_states[..., None, :] * identity_weights[..., None]
+        weighted_outputs = weighted_outputs + identity_outputs
+
+        final_output = mx.sum(weighted_outputs, axis=-2)
+        return final_output
+
+
+class LongcatFlashDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.mlp = LongcatFlashMoE(args)
+
+        self.self_attn = [LongcatFlashMLA(args) for _ in range(2)]
+        self.mlps = [LongcatFlashMLP(args, False) for _ in range(2)]
+        self.input_layernorm = [
+            nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) for _ in range(2)
+        ]
+        self.post_attention_layernorm = [
+            nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) for _ in range(2)
+        ]
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        hidden_states = x
+        shortcut_mlp_output = None
+
+        for i in range(2):
+            residual = hidden_states
+
+            hidden_states = self.input_layernorm[i](hidden_states)
+            hidden_states = self.self_attn[i](hidden_states, mask=mask, cache=cache[i])
+            hidden_states = residual + hidden_states
+
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm[i](hidden_states)
+
+            if i == 0:
+                shortcut_mlp_output = self.mlp(hidden_states)
+
+            hidden_states = self.mlps[i](hidden_states)
+            hidden_states = residual + hidden_states
+
+            if i == 1:
+                hidden_states = hidden_states + shortcut_mlp_output
+
+        return hidden_states
+
+
+class LongcatFlashModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_layers = args.num_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [LongcatFlashDecoderLayer(args) for idx in range(args.num_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        if mask is None:
+            mask = create_attention_mask(
+                h, [cache[0][0]] if cache is not None else None
+            )
+
+        if cache is None:
+            cache = [None] * self.num_layers
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = LongcatFlashModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("classifier"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate
+
+    def sanitize(self, weights):
+        for l in range(self.args.num_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        new_weights = {}
+        for k, v in weights.items():
+            if k.startswith("model.mtp"):
+                continue
+            new_weights[k] = v
+        return new_weights
+
+    def make_cache(self):
+        return [CacheList(KVCache(), KVCache()) for _ in self.model.layers]

mlx_lm/models/nemotron_h.py

@@ -0,0 +1,387 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, MambaCache
+
+
+@dataclass()
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attention_bias: bool
+    mamba_num_heads: int
+    mamba_head_dim: int
+    mamba_proj_bias: bool
+    ssm_state_size: int
+    conv_kernel: int
+    n_groups: int
+    time_step_limit: Tuple[float, float]
+    mlp_bias: bool
+    layer_norm_epsilon: float
+    rms_norm_eps: float
+    use_bias: bool
+    use_conv_bias: bool
+    residual_in_fp32: bool
+    head_dim: Optional[int] = None
+    hybrid_override_pattern: Optional[List[str]] = None
+
+
+class MambaRMSNormGated(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = mx.ones(hidden_size)
+
+    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+        if gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+        return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
+
+
+class NemotronHMamba2Mixer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_heads = args.mamba_num_heads
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.ssm_state_size
+        self.conv_kernel_size = args.conv_kernel
+        self.intermediate_size = args.mamba_num_heads * args.mamba_head_dim
+        self.n_groups = args.n_groups
+        self.head_dim = args.mamba_head_dim
+        self.time_step_limit = args.time_step_limit
+        self.heads_per_group = self.num_heads // self.n_groups
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            kernel_size=args.conv_kernel,
+            padding=0,
+            groups=self.conv_dim,
+            bias=args.use_conv_bias,
+        )
+
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size, projection_size, bias=args.mamba_proj_bias
+        )
+
+        self.dt_bias = mx.ones(self.num_heads)
+        self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
+        self.D = mx.ones(self.num_heads)
+
+        self.norm = MambaRMSNormGated(
+            self.intermediate_size, eps=args.layer_norm_epsilon
+        )
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias
+        )
+
+    def _apply_conv(
+        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        if cache is not None:
+            if cache[0] is None:
+                conv_state = mx.zeros(
+                    (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim),
+                    dtype=conv_input.dtype,
+                )
+            else:
+                conv_state = cache[0]
+            padded_input = mx.concatenate([conv_state, conv_input], axis=1)
+            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :]
+        else:
+            padded_input = mx.pad(
+                conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)]
+            )
+
+        conv_output = self.conv1d(padded_input)
+        return nn.silu(conv_output)
+
+    def _ssm(
+        self,
+        hidden_states: mx.array,
+        B: mx.array,
+        C: mx.array,
+        dt: mx.array,
+        cache: Optional[MambaCache] = None,
+    ) -> mx.array:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        dt = nn.softplus(dt + self.dt_bias)
+        dt = mx.clip(dt, self.time_step_limit[0], self.time_step_limit[1])
+
+        hidden_states = hidden_states.reshape(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        )
+
+        B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+        B = mx.repeat(B, self.heads_per_group, axis=2)
+        C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+        C = mx.repeat(C, self.heads_per_group, axis=2)
+
+        A = -mx.exp(self.A_log.astype(mx.float32)).astype(hidden_states.dtype)
+
+        if cache is not None and cache[1] is not None:
+            h = cache[1]
+        else:
+            h = mx.zeros(
+                (batch_size, self.num_heads, self.head_dim, self.ssm_state_size),
+                dtype=hidden_states.dtype,
+            )
+
+        outputs = []
+        for t in range(seq_len):
+            dt_t = dt[:, t, :]
+            dA = mx.exp(dt_t * A)[..., None, None]
+            dB = (dt_t[..., None] * B[:, t])[..., None, :]
+
+            h = dA * h + dB * hidden_states[:, t, :, :, None]
+            y_t = (h @ C[:, t, :, :, None]).squeeze(-1) + self.D[
+                :, None
+            ] * hidden_states[:, t]
+            outputs.append(y_t)
+
+        if cache is not None:
+            cache[1] = h
+
+        y = mx.stack(outputs, axis=1)
+        return y.reshape(batch_size, seq_len, self.intermediate_size)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        cache: Optional[MambaCache] = None,
+    ) -> mx.array:
+
+        projected = self.in_proj(hidden_states)
+
+        gate, conv_input, dt = mx.split(
+            projected,
+            [self.intermediate_size, self.intermediate_size + self.conv_dim],
+            axis=-1,
+        )
+
+        conv_output = self._apply_conv(conv_input, cache)
+
+        hidden_states_ssm, B, C = mx.split(
+            conv_output,
+            [
+                self.intermediate_size,
+                self.intermediate_size + self.n_groups * self.ssm_state_size,
+            ],
+            axis=-1,
+        )
+        y = self._ssm(hidden_states_ssm, B, C, dt, cache)
+        y = self.norm(y, gate)
+        return self.out_proj(y)
+
+
+class NemotronHAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = (
+            args.head_dim
+            if args.head_dim is not None
+            else (args.hidden_size // args.num_attention_heads)
+        )
+        self.num_key_value_heads = args.num_key_value_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=args.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=args.attention_bias
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries = self.q_proj(x).reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        keys = (
+            self.k_proj(x)
+            .reshape(B, L, self.num_key_value_heads, -1)
+            .transpose(0, 2, 1, 3)
+        )
+        values = (
+            self.v_proj(x)
+            .reshape(B, L, self.num_key_value_heads, -1)
+            .transpose(0, 2, 1, 3)
+        )
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+@partial(mx.compile, shapeless=True)
+def relu2(x):
+    return mx.square(nn.relu(x))
+
+
+class NemotronHMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.up_proj = nn.Linear(
+            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+        )
+
+    def __call__(self, x):
+        return self.down_proj(relu2(self.up_proj(x)))
+
+
+class NemotronHBlock(nn.Module):
+    def __init__(self, args: ModelArgs, block_type: str):
+        super().__init__()
+        self.residual_in_fp32 = args.residual_in_fp32
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.block_type = block_type
+
+        if self.block_type == "M":
+            self.mixer = NemotronHMamba2Mixer(args)
+        elif self.block_type == "*":
+            self.mixer = NemotronHAttention(args)
+        elif self.block_type == "-":
+            self.mixer = NemotronHMLP(args)
+
+    def __call__(
+        self,
+        x,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        hidden_states = self.norm(x)
+        if self.block_type == "M":
+            hidden_states = self.mixer(hidden_states, cache=cache)
+        elif self.block_type == "*":
+            hidden_states = self.mixer(hidden_states, mask=mask, cache=cache)
+        else:
+            hidden_states = self.mixer(hidden_states)
+
+        return x + hidden_states
+
+
+class NemotronHModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            NemotronHBlock(args, block_type)
+            for block_type in args.hybrid_override_pattern
+        ]
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.fa_idx = 0
+        for b in args.hybrid_override_pattern:
+            if b == "*":
+                break
+            elif b == "M":
+                self.fa_idx += 1
+
+    def __call__(
+        self,
+        inputs,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        hidden_states = self.embeddings(inputs)
+
+        if mask is None:
+            attn_mask = create_attention_mask(
+                hidden_states, cache[self.fa_idx : self.fa_idx + 1]
+            )
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        cache_counter = 0
+        for layer in self.layers:
+            if layer.block_type == "M" or layer.block_type == "*":
+                c = cache[cache_counter]
+                cache_counter += 1
+            else:
+                c = None
+
+            if layer.block_type == "*":
+                mask = attn_mask
+            else:
+                mask = None
+            hidden_states = layer(hidden_states, mask=mask, cache=c)
+
+        return self.norm_f(hidden_states)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.backbone = NemotronHModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        out = self.backbone(inputs, mask=mask, cache=cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.backbone.layers
+
+    def make_cache(self):
+        caches = []
+        for l in self.layers:
+            if l.block_type == "M":
+                caches.append(MambaCache())
+            elif l.block_type == "*":
+                caches.append(KVCache())
+        return caches
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+        return weights

mlx_lm/models/qwen2_vl.py

@@ -0,0 +1,59 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import qwen2
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    @classmethod
+    def from_dict(cls, params):
+        if "text_config" not in params:
+            return cls(model_type=params["model_type"], text_config=params)
+        return cls(**params)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = qwen2.Model(qwen2.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("visual", None)
+        weights.pop("vision_tower", None)
+        weights = dict(tree_flatten(weights))
+
+        sanitized = {}
+        for key, value in weights.items():
+            if not key.startswith("language_model."):
+                key = "language_model." + key
+            sanitized[key] = value
+        return sanitized
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers

mlx_lm/models/qwen3_moe.py

@@ -235,6 +235,15 @@ class Model(nn.Module):
                     weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(to_join)
         return weights
 
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
     @property
     def layers(self):
         return self.model.layers

mlx_lm/models/rope_utils.py

@@ -236,6 +236,7 @@ def initialize_rope(
             dims=dims,
             max_position_embeddings=max_position_embeddings,
             traditional=traditional,
+            scaling_factor=scaling_factor,
             base=base,
             **rope_kwargs,
         )
@@ -250,6 +251,11 @@ def initialize_rope(
             short_factor=scaling_config["short_factor"],
             long_factor=scaling_config["long_factor"],
         )
-
+    elif rope_type == "mrope":
+        mrope_section = scaling_config.get("mrope_section", [])
+        assert (
+            len(mrope_section) == 3
+        ), f"MRoPE currently only supports 3 sections, got {len(mrope_section)}."
+        return nn.RoPE(dims, traditional=traditional, base=base)
     else:
         raise ValueError(f"Unsupported RoPE type {rope_type}")

mlx_lm/models/seed_oss.py

@@ -0,0 +1,186 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    head_dim: int
+    max_position_embeddings: Optional[int] = None
+    attention_bias: bool = False
+    attention_out_bias: bool = False
+    mlp_bias: bool = False
+    rope_theta: float = 10000
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        self.head_dim = head_dim = args.head_dim
+
+        self.scale = head_dim**-0.5
+
+        input_bias = args.attention_bias
+        output_bias = args.attention_out_bias
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=input_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=input_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=input_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=output_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, bias=False):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size, args.mlp_bias)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class SeedModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = SeedModel(args)
+        self.tie_word_embeddings = args.tie_word_embeddings
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.model(inputs, mask=mask, cache=cache)
+        if self.tie_word_embeddings:
+            return h @ self.model.embed_tokens.weight.T
+        else:
+            return self.lm_head(h)
+
+    def sanitize(self, weights):
+        if self.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/switch_layers.py

@@ -1,6 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import math
+from functools import partial
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -30,11 +31,12 @@ class QuantizedSwitchLinear(nn.Module):
         bias: bool = True,
         group_size: int = 64,
         bits: int = 4,
+        mode: str = "affine",
     ):
         super().__init__()
 
         scale = math.sqrt(1 / input_dims)
-        self.weight, self.scales, self.biases = mx.quantize(
+        self.weight, self.scales, *biases = mx.quantize(
             mx.random.uniform(
                 low=-scale,
                 high=scale,
@@ -42,13 +44,16 @@ class QuantizedSwitchLinear(nn.Module):
             ),
             group_size=group_size,
             bits=bits,
+            mode=mode,
         )
+        self.biases = biases[0] if biases else None
 
         if bias:
             self.bias = mx.zeros((num_experts, output_dims))
 
         self.group_size = group_size
         self.bits = bits
+        self.mode = mode
 
         # Freeze this model's parameters
         self.freeze()
@@ -70,11 +75,12 @@ class QuantizedSwitchLinear(nn.Module):
             x,
             self["weight"],
             self["scales"],
-            self["biases"],
+            self.get("biases"),
             rhs_indices=indices,
             transpose=True,
             group_size=self.group_size,
             bits=self.bits,
+            mode=self.mode,
             sorted_indices=sorted_indices,
         )
         if "bias" in self:
@@ -120,24 +126,47 @@ class SwitchLinear(nn.Module):
             x = x + mx.expand_dims(self["bias"][indices], -2)
         return x
 
-    def to_quantized(self, group_size: int = 64, bits: int = 4):
+    def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"):
         num_experts, output_dims, input_dims = self.weight.shape
         ql = QuantizedSwitchLinear(
-            input_dims, output_dims, num_experts, False, group_size, bits
+            input_dims,
+            output_dims,
+            num_experts,
+            False,
+            group_size,
+            bits,
+            mode=mode,
         )
-        ql.weight, ql.scales, ql.biases = mx.quantize(self.weight, group_size, bits)
+        ql.weight, ql.scales, *biases = mx.quantize(
+            self.weight, group_size, bits, mode=mode
+        )
+        ql.biases = biases[0] if biases else None
+
         if "bias" in self:
             ql.bias = self.bias
         return ql
 
 
+@partial(mx.compile, shapeless=True)
+def swiglu(x, gate):
+    return nn.silu(gate) * x
+
+
+class SwiGLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x, gate):
+        return swiglu(x, gate)
+
+
 class SwitchGLU(nn.Module):
     def __init__(
         self,
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.SiLU(),
+        activation=SwiGLU(),
         bias: bool = False,
     ):
         super().__init__()
@@ -162,7 +191,7 @@ class SwitchGLU(nn.Module):
         x_up = self.up_proj(x, idx, sorted_indices=do_sort)
         x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
         x = self.down_proj(
-            self.activation(x_gate) * x_up,
+            self.activation(x_up, x_gate),
             idx,
             sorted_indices=do_sort,
         )

mlx_lm/perplexity.py

@@ -0,0 +1,187 @@
+# Copyright © 2025 Apple Inc.
+"""
+Evaluate perplexity (PPL) of MLX models.
+"""
+
+import argparse
+import math
+import time
+import types
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from mlx_lm.tuner.datasets import load_dataset
+from mlx_lm.tuner.utils import get_total_parameters
+from mlx_lm.utils import load
+
+
+def load_data(
+    tokenizer,
+    data_path: str,
+    num_samples: int,
+    sequence_length: int,
+):
+    args = types.SimpleNamespace(
+        hf_dataset={
+            "path": data_path,
+            "train_split": "train",
+            "valid_split": "train[:1]",
+        },
+        train=True,
+        test=False,
+    )
+    dataset = load_dataset(args, tokenizer)[0]
+
+    perm = np.random.permutation(len(dataset)).tolist()
+
+    num_tokens = sequence_length * num_samples if num_samples > 0 else float("inf")
+    data = []
+    i = 0
+    while len(data) < num_tokens:
+        tokens, _ = dataset.process(dataset[perm[i]])
+        i += 1
+        data.extend(tokens)
+
+    data = mx.array(data[: (len(data) // sequence_length) * sequence_length])
+    data = data.reshape(-1, sequence_length)
+    if num_samples > 0:
+        data = data[:num_samples]
+    return data
+
+
+def eval_ppl(model, data, batch_size=8):
+    """
+    Evaluate perplexity on a dataset with standard error calculation.
+
+    Args:
+        model: The model to evaluate
+        data: Tokenized data tensor
+        batch_size: Batch size for evaluation
+
+    Returns:
+        tuple: (perplexity, standard_error)
+    """
+    all_losses = []
+
+    num_batches = (len(data) + batch_size - 1) // batch_size
+    for i, s in enumerate(range(0, len(data), batch_size)):
+        batch = data[s : s + batch_size]
+        # Forward pass: get logits for all tokens except last
+        logits = model(batch[:, :-1]).astype(mx.float32)
+
+        # Calculate cross-entropy loss with next tokens
+        losses = nn.losses.cross_entropy(logits, batch[:, 1:], reduction="none")
+        mx.eval(losses)
+        # Store individual token losses
+        all_losses.append(losses.flatten())
+
+        # Progress indicator
+        if (i + 1) % 1 == 0 or (i + 1) == num_batches:
+            print(f"  Processed {i + 1}/{num_batches} batches...", end="\r")
+
+    print()  # New line after progress
+
+    # Concatenate all losses into a single array
+    all_losses = mx.concatenate(all_losses)
+
+    # Calculate mean loss and perplexity
+    mean_loss = all_losses.mean().item()
+    ppl = math.exp(mean_loss)
+    # Calculate standard error
+    std_dev = mx.sqrt(mx.var(all_losses, ddof=1)).item()
+    num_tokens = all_losses.size
+    standard_error = std_dev / math.sqrt(num_tokens)
+    # Delta approximation for standard error of perplexity
+    standard_error_ppl = ppl * standard_error
+
+    return ppl, standard_error_ppl
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate perplexity of MLX models")
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to model or Hugging Face model ID",
+    )
+    parser.add_argument(
+        "--batch-size", type=int, default=8, help="Batch size for evaluation"
+    )
+    parser.add_argument(
+        "--sequence-length",
+        type=int,
+        default=512,
+        help="Sequence length for evaluation",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=256,
+        help="Number of samples to use (-1 for all available)",
+    )
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        default="allenai/tulu-3-sft-mixture",
+        help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=123, help="Random seed for data sampling"
+    )
+
+    args = parser.parse_args()
+
+    # Set random seed
+    np.random.seed(args.seed)
+    mx.random.seed(args.seed)
+
+    # Load model
+    print(f"Loading model from {args.model}...")
+    model, tokenizer = load(args.model)
+
+    # Count parameters
+    total_params = get_total_parameters(model)
+    print(f"Model loaded: {total_params/1e6:.1f}M parameters")
+
+    # Load evaluation data
+    print(f"\nLoading dataset...")
+    print(f"  Sequence length: {args.sequence_length}")
+
+    data = load_data(
+        tokenizer,
+        args.data_path,
+        num_samples=args.num_samples,
+        sequence_length=args.sequence_length,
+    )
+
+    print(f"  Loaded {len(data)} samples")
+
+    # Evaluate perplexity
+    print(f"\nEvaluating perplexity with batch size {args.batch_size}...")
+    start_time = time.time()
+
+    ppl, se = eval_ppl(model, data, batch_size=args.batch_size)
+
+    eval_time = time.time() - start_time
+    tokens_evaluated = data.shape[0] * (data.shape[1] - 1)  # B * (L - 1)
+    # Print results
+    print("\n" + "=" * 60)
+    print("EVALUATION RESULTS")
+    print("=" * 60)
+    print(f"Model: {args.model}")
+    print(f"Perplexity: {ppl:.3f} ± {se:.3f}")
+    print(f"Evaluation time: {eval_time:.2f} seconds")
+    print(f"Peak memory: {mx.get_peak_memory() / 1e9:.2f} GB")
+    print(f"Tokens per second: {tokens_evaluated / eval_time:.0f}")
+
+    # Additional statistics
+    print(f"\nDataset statistics:")
+    print(f"  Total samples: {len(data)}")
+    print(f"  Total tokens: {data.size}")
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/quant/dwq.py

@@ -30,15 +30,25 @@ class Catcher(nn.Module):
         self.module = module
 
     def __call__(self, *args, **kwargs):
-        self.outputs = self.module(*args, **kwargs)
-        return self.outputs
+        outputs = self.module(*args, **kwargs)
+        self.outputs = outputs[0] if isinstance(outputs, tuple) else outputs
+        return outputs
+
+    def __getattr__(self, key: str):
+        if (value := self.get(key)) is not None:
+            return value
+        elif (m := self.get("module")) is not None:
+            return getattr(m, key)
+        else:
+            super(nn.Module, self).__getattribute__(key)
 
 
 def dwq_quantize(
     model,
     q_model,
     opt,
-    data,
+    train_data,
+    valid_data,
     batch_size: int = 2,
     max_seq_length: int = 2048,
     activation_layer_step: float = 0.25,
@@ -50,10 +60,15 @@ def dwq_quantize(
     world_size = group.size()
     rank = group.rank()
 
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            tqdm.write(*args, **kwargs)
+
     def unfreeze(_, m):
-        if hasattr(m, "bits") and hasattr(m, "group_size"):
+        if hasattr(m, "bits") and hasattr(m, "group_size") and m.mode == "affine":
             m.unfreeze(keys=["scales", "biases"], recurse=False)
 
+    q_model.train()
     q_model.apply_to_modules(unfreeze)
     print_trainable_parameters(q_model)
 
@@ -89,17 +104,52 @@ def dwq_quantize(
                 for qe, e in zip(q_extra_targets, extra_targets)
             ]
         )
-        loss = kl_loss + activation_loss_weight * act_loss.mean()
-        return loss, ntoks
+        act_loss = act_loss.mean()
+        loss = kl_loss + activation_loss_weight * act_loss
+        return loss, ntoks, kl_loss, act_loss
 
     def step(inputs, targets, extra_targets, lengths, params):
-        (loss, ntoks), grads = mx.value_and_grad(loss_fn)(
+        (loss, ntoks, *_), grads = mx.value_and_grad(loss_fn)(
             params, inputs, targets, extra_targets, lengths
         )
         grads = nn.average_gradients(grads)
         params = opt.apply_gradients(grads, params)
         return loss, ntoks, params
 
+    def validate(params, it):
+        v_loss = 0.0
+        v_kl_loss = 0.0
+        v_act_loss = 0.0
+        v_tokens = 0
+        for it, (batch, lengths) in tqdm(
+            enumerate(iterate_batches(valid_data, batch_size, max_seq_length)),
+            total=len(valid_data) // batch_size,
+            desc="Computing validation loss",
+            leave=False,
+        ):
+            batch = batch[:, :-1]
+            targets, extra_targets = forward(model, batch)
+            mx.eval(targets, extra_targets)
+            loss, ntoks, kl_loss, act_loss = loss_fn(
+                params, batch, targets, extra_targets, lengths
+            )
+            mx.eval(loss, ntoks)
+            loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
+            kl_loss = mx.distributed.all_sum(kl_loss, stream=mx.cpu).item() / world_size
+            act_loss = (
+                mx.distributed.all_sum(act_loss, stream=mx.cpu).item() / world_size
+            )
+            ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
+            v_tokens += ntoks
+            v_loss += loss * ntoks
+            v_kl_loss += kl_loss * ntoks
+            v_act_loss += act_loss * ntoks
+        loss = v_loss / v_tokens
+        kl_loss = v_kl_loss / v_tokens
+        act_loss = v_act_loss / v_tokens
+        rprint(f"Validation: {it=}, {loss=:.3f}, {kl_loss=:.3f}, {act_loss=:.3f}")
+        return loss
+
     # Accumulate learned weights in higher precision
     params = tree_map(
         lambda x: x.astype(mx.float32),
@@ -109,11 +159,16 @@ def dwq_quantize(
     total_loss = 0.0
     total_tokens = 0
     tokens = 0
+
     tic = time.time()
+
+    # Compute initial validation loss
+    initial_valid_loss = valid_loss = validate(params, it=0)
+
     for it, (batch, lengths) in (
         pbar := tqdm(
-            enumerate(iterate_batches(data, batch_size, max_seq_length)),
-            total=len(data) // batch_size,
+            enumerate(iterate_batches(train_data, batch_size, max_seq_length)),
+            total=len(train_data) // batch_size,
         )
     ):
         batch = batch[:, :-1]
@@ -132,42 +187,74 @@ def dwq_quantize(
                 peak_memory_gb = mx.get_peak_memory() / 1e9
                 avg_loss = total_loss / tokens
                 total_tokens += tokens
-                tqdm.write(
+                rprint(
                     f"{it=}, {avg_loss=:.4f}, {total_tokens=},"
                     f" {toks_per_sec=:.3f}, {peak_memory_gb=:.3f}",
                 )
                 tic = time.time()
                 tokens = 0
                 total_loss = 0
+        if (it + 1) % 200 == 0:
+            valid_loss = validate(params, it=it)
+
+    valid_loss = validate(params, it=it)
+    if initial_valid_loss < valid_loss:
+        rprint(
+            f"❌❌❌\n[WARNING] Final validation loss {valid_loss:.3f} is "
+            f"worse than initial validation loss {initial_valid_loss:.3f}."
+            " Model quality will likely be degraded.\n❌❌❌"
+        )
+
     q_model.update(tree_map(lambda x: x.astype(dtype), params))
     for lid in layer_ids:
         q_model.layers[lid] = q_model.layers[lid].module
 
 
-def load_data(tokenizer, data_path: str, num_samples: int, max_seq_length: int):
+def load_data(
+    tokenizer,
+    data_path: str,
+    num_samples: int,
+    max_seq_length: int,
+    num_valid_samples: int = 32,
+):
     args = types.SimpleNamespace(
         hf_dataset={
             "path": data_path,
-            "train_split": f"train",
+            "train_split": "train",
             "valid_split": "train[:1]",
         },
         train=True,
         test=False,
     )
     dataset = load_dataset(args, tokenizer)[0]
-    perm = np.random.permutation(len(dataset))[:num_samples].tolist()
+    perm = np.random.permutation(len(dataset))
+    train_perm = perm[:num_samples].tolist()
+    valid_perm = perm[num_samples : num_samples + num_valid_samples].tolist()
 
     def process(idx):
         tokens, offset = dataset.process(dataset[idx])
         return (tokens[:max_seq_length], offset)
 
-    return [process(i) for i in perm]
+    train = [process(i) for i in train_perm]
+    valid = [process(i) for i in valid_perm]
+    return train, valid
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model", "-m", default="Qwen/Qwen3-4B")
-    parser.add_argument("--quantized-model", default=None)
+    parser.add_argument(
+        "--model",
+        "-m",
+        help="A model to distill from for DWQ. If `quantized-model` is not"
+        " given the student model will be this model quantized according"
+        " to `bits` and `group-size`.",
+        required=True,
+    )
+    parser.add_argument(
+        "--quantized-model",
+        default=None,
+        help="An already quantized model (the student model) to improve with DWQ.",
+    )
     parser.add_argument(
         "--mlx-path", default="mlx_model", help="Path to save the quantized model."
     )
@@ -213,22 +300,28 @@ def main():
     mx.random.seed(args.seed)
 
     model_path, hf_repo = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=True
+    )
 
-    calibration_data = load_data(
+    train_data, valid_data = load_data(
         tokenizer, args.data_path, args.num_samples, args.max_seq_length
     )
 
     if args.quantized_model is not None:
-        q_model_path = get_model_path(args.quantized_model, revision=None)
-        q_model, config, _ = fetch_from_hub(q_model_path, lazy=True)
+        q_model_path, _ = get_model_path(args.quantized_model, revision=None)
+        q_model, config, _ = fetch_from_hub(
+            q_model_path, lazy=True, trust_remote_code=True
+        )
+        if "quantization" not in config:
+            raise ValueError("Quantized model must already be quantized.")
     else:
         q_model = copy.deepcopy(model)
         _, config = quantize_model(
             q_model,
             config,
-            q_group_size=args.group_size,
-            q_bits=args.bits,
+            group_size=args.group_size,
+            bits=args.bits,
         )
 
     opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
@@ -236,7 +329,8 @@ def main():
         model,
         q_model,
         opt,
-        calibration_data,
+        train_data,
+        valid_data,
         batch_size=args.batch_size,
         max_seq_length=args.max_seq_length,
         gradient_checkpoint=args.grad_checkpoint,

mlx_lm/quant/dynamic_quant.py

@@ -240,8 +240,8 @@ def main():
     model, config = quantize_model(
         model,
         config,
-        q_group_size=args.low_group_size,
-        q_bits=args.low_bits,
+        group_size=args.low_group_size,
+        bits=args.low_bits,
         quant_predicate=quant_predicate,
     )
 

mlx_lm/sample_utils.py

@@ -16,7 +16,7 @@ def make_sampler(
     xtc_probability: float = 0.0,
     xtc_threshold: float = 0.0,
     xtc_special_tokens: List[int] = [],
-) -> Callable[mx.array, mx.array]:
+) -> Callable[[mx.array], mx.array]:
     """
     Make a sampler function for use with ``generate_step``.
 
@@ -48,8 +48,6 @@ def make_sampler(
 
     # Create sampler chain
     sampling_methods = []
-    if top_k > 0:
-        sampling_methods.append(lambda x: apply_top_k(x, top_k))
     if top_p > 0 and top_p < 1.0:
         sampling_methods.append(lambda x: apply_top_p(x, top_p))
     if min_p != 0.0:
@@ -58,14 +56,15 @@ def make_sampler(
         sampling_methods.append(
             lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens)
         )
+    if top_k > 0:
+        sampling_methods.append(lambda x: apply_top_k(x, top_k))
 
     # Apply the sampling methods
-    def sampler(logits):
+    def sampler(logprobs):
         for method in sampling_methods:
-            logits = method(logits)
-
+            logprobs = method(logprobs)
         # Return the sampled token
-        return categorical_sampling(logits, temp)
+        return categorical_sampling(logprobs, temp)
 
     return sampler
 

mlx_lm/server.py

@@ -162,8 +162,10 @@ class ModelProvider:
         self.draft_model = None
 
         # Preload the default model if it is provided
+        self.default_model_map = {}
         if self.cli_args.model is not None:
-            self.load("default_model", draft_model_path="default_model")
+            self.default_model_map[self.cli_args.model] = "default_model"
+            self.load(self.cli_args.model, draft_model_path="default_model")
 
     def _validate_model_path(self, model_path: str):
         model_path = Path(model_path)
@@ -174,6 +176,12 @@ class ModelProvider:
 
     # Added in adapter_path to load dynamically
     def load(self, model_path, adapter_path=None, draft_model_path=None):
+        model_path, adapter_path, draft_model_path = map(
+            lambda s: s.lower() if s else None,
+            (model_path, adapter_path, draft_model_path),
+        )
+
+        model_path = self.default_model_map.get(model_path, model_path)
         if self.model_key == (model_path, adapter_path, draft_model_path):
             return self.model, self.tokenizer
 
@@ -196,11 +204,10 @@ class ModelProvider:
                     "A model path has to be given as a CLI "
                     "argument or in the HTTP request"
                 )
+            adapter_path = adapter_path or self.cli_args.adapter_path
             model, tokenizer = load(
                 self.cli_args.model,
-                adapter_path=(
-                    adapter_path if adapter_path else self.cli_args.adapter_path
-                ),  # if the user doesn't change the model but adds an adapter path
+                adapter_path=adapter_path,
                 tokenizer_config=tokenizer_config,
             )
         else:
@@ -297,7 +304,23 @@ class APIHandler(BaseHTTPRequestHandler):
         # Fetch and parse request body
         content_length = int(self.headers["Content-Length"])
         raw_body = self.rfile.read(content_length)
-        self.body = json.loads(raw_body.decode())
+        try:
+            self.body = json.loads(raw_body.decode())
+        except json.JSONDecodeError as e:
+            logging.error(f"JSONDecodeError: {e} - Raw body: {raw_body.decode()}")
+            # Set appropriate headers based on streaming requirement
+            if self.stream:
+                self._set_stream_headers(400)
+                self.wfile.write(
+                    f"data: {json.dumps({'error': f'Invalid JSON in request body: {e}'})}\n\n".encode()
+                )
+            else:
+                self._set_completion_headers(400)
+                self.wfile.write(
+                    json.dumps({"error": f"Invalid JSON in request body: {e}"}).encode()
+                )
+            return
+
         indent = "\t"  # Backslashes can't be inside of f-strings
         logging.debug(f"Incoming Request Body: {json.dumps(self.body, indent=indent)}")
         assert isinstance(
@@ -330,7 +353,10 @@ class APIHandler(BaseHTTPRequestHandler):
         self.xtc_threshold = self.body.get("xtc_threshold", 0.0)
         self.logit_bias = self.body.get("logit_bias", None)
         self.logprobs = self.body.get("logprobs", -1)
+        self.seed = self.body.get("seed", None)
         self.validate_model_parameters()
+        if self.seed is not None:
+            mx.random.seed(self.seed)
         # Load the model if needed
         try:
             self.model, self.tokenizer = self.model_provider.load(
@@ -427,6 +453,8 @@ class APIHandler(BaseHTTPRequestHandler):
             raise ValueError("model must be a string")
         if self.adapter is not None and not isinstance(self.adapter, str):
             raise ValueError("adapter must be a string")
+        if self.seed is not None and not isinstance(self.seed, int):
+            raise ValueError("seed must be an integer")
 
     def generate_response(
         self,
@@ -487,16 +515,18 @@ class APIHandler(BaseHTTPRequestHandler):
             "choices": [
                 {
                     "index": 0,
-                    "logprobs": {
-                        "token_logprobs": token_logprobs,
-                        "top_logprobs": top_logprobs,
-                        "tokens": tokens,
-                    },
                     "finish_reason": finish_reason,
                 },
             ],
         }
 
+        if token_logprobs or top_logprobs or tokens:
+            response["choices"][0]["logprobs"] = {
+                "token_logprobs": token_logprobs,
+                "top_logprobs": top_logprobs,
+                "tokens": tokens,
+            }
+
         if not self.stream:
             if not (
                 isinstance(prompt_token_count, int)
@@ -525,7 +555,7 @@ class APIHandler(BaseHTTPRequestHandler):
         elif self.object_type == "text_completion":
             choice.update(text=text)
         else:
-            ValueError(f"Unsupported response type: {self.object_type}")
+            raise ValueError(f"Unsupported response type: {self.object_type}")
 
         return response
 
@@ -662,6 +692,23 @@ class APIHandler(BaseHTTPRequestHandler):
         tool_text = ""
         in_tool_call = False
         segment = ""
+
+        # Create keepalive callback to send SSE comments during long prompt processing
+        def keepalive_callback(processed_tokens, total_tokens):
+            logging.info(
+                f"Prompt processing progress: {processed_tokens}/{total_tokens}"
+            )
+            if self.stream:
+                try:
+                    # Send SSE comment for keepalive - invisible to clients but keeps connection alive
+                    self.wfile.write(
+                        f": keepalive {processed_tokens}/{total_tokens}\n\n".encode()
+                    )
+                    self.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError, OSError):
+                    # Client disconnected, ignore
+                    pass
+
         for gen_response in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
@@ -672,6 +719,7 @@ class APIHandler(BaseHTTPRequestHandler):
             prompt_cache=self.prompt_cache.cache,
             draft_model=self.model_provider.draft_model,
             num_draft_tokens=self.num_draft_tokens,
+            prompt_progress_callback=keepalive_callback,
         ):
             logging.debug(gen_response.text)
 
@@ -693,6 +741,7 @@ class APIHandler(BaseHTTPRequestHandler):
             token = gen_response.token
             logprobs = gen_response.logprobs
             tokens.append(token)
+            self.prompt_cache.tokens.append(token)
 
             if self.logprobs > 0:
                 sorted_indices = mx.argpartition(-logprobs, kth=self.logprobs - 1)
@@ -735,7 +784,8 @@ class APIHandler(BaseHTTPRequestHandler):
                     segment = ""
                     tool_calls = []
 
-        self.prompt_cache.tokens.extend(tokens)
+        if gen_response.finish_reason is not None:
+            finish_reason = gen_response.finish_reason
 
         logging.debug(f"Prompt: {gen_response.prompt_tps:.3f} tokens-per-sec")
         logging.debug(f"Generation: {gen_response.generation_tps:.3f} tokens-per-sec")
@@ -748,7 +798,12 @@ class APIHandler(BaseHTTPRequestHandler):
             self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
             self.wfile.flush()
             if self.stream_options is not None and self.stream_options["include_usage"]:
-                response = self.completion_usage_response(len(prompt), len(tokens))
+                original_prompt_length = (
+                    len(self.prompt_cache.tokens) - len(tokens) + len(prompt)
+                )
+                response = self.completion_usage_response(
+                    original_prompt_length, len(tokens)
+                )
                 self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
                 self.wfile.flush()
             self.wfile.write("data: [DONE]\n\n".encode())
@@ -839,7 +894,7 @@ class APIHandler(BaseHTTPRequestHandler):
         """
         Respond to a GET request from a client.
         """
-        if self.path == "/v1/models":
+        if self.path.startswith("/v1/models"):
             self.handle_models_request()
         elif self.path == "/health":
             self.handle_health_check()
@@ -867,11 +922,18 @@ class APIHandler(BaseHTTPRequestHandler):
 
         files = ["config.json", "model.safetensors.index.json", "tokenizer_config.json"]
 
+        parts = self.path.split("/")
+        filter_repo_id = None
+        if len(parts) > 3:
+            filter_repo_id = "/".join(parts[3:])
+
         def probably_mlx_lm(repo):
             if repo.repo_type != "model":
                 return False
             if "main" not in repo.refs:
                 return False
+            if filter_repo_id is not None and repo.repo_id != filter_repo_id:
+                return False
             file_names = {f.file_path.name for f in repo.refs["main"].files}
             return all(f in file_names for f in files)
 

mlx_lm/tokenizer_utils.py

@@ -89,7 +89,7 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
     def text(self):
         if self._current_tokens:
             self._current_text = self._tokenizer.decode(self._current_tokens)
-            if (
+            if self._current_text.endswith("\ufffd") or (
                 self._tokenizer.clean_up_tokenization_spaces
                 and len(self._current_text) > 0
                 and self._current_text[-1] == " "
@@ -283,8 +283,6 @@ class TokenizerWrapper:
                 self._think_end = think_end
                 break
         if tokenizer.chat_template and '"tool"' in tokenizer.chat_template:
-            self._tool_call_start = ""
-            self._tool_call_end = ""
             for tool_call_start, tool_call_end in TOOL_CALL_TOKENS:
                 if tool_call_start in vocab and tool_call_end in vocab:
                     self._tool_call_start = tool_call_start

mlx_lm/tuner/callbacks.py

@@ -1,10 +1,17 @@
 # Copyright © 2024 Apple Inc.
 
+import os
+
 try:
     import wandb
 except ImportError:
     wandb = None
 
+try:
+    import swanlab
+except ImportError:
+    swanlab = None
+
 
 class TrainingCallback:
 
@@ -27,17 +34,100 @@ class WandBCallback(TrainingCallback):
     ):
         if wandb is None:
             raise ImportError(
-                "wandb is not installed. Please install it to use WandBCallback."
+                "wandb is not installed. please install wandb via: pip install wandb",
             )
         self.wrapped_callback = wrapped_callback
-        wandb.init(project=project_name, dir=log_dir, config=config)
+        wandb.init(
+            project=project_name,
+            name=os.path.basename(log_dir),
+            dir=log_dir,
+            config=config,
+        )
+
+    def _convert_to_serializable(self, data: dict) -> dict:
+        return {k: v.tolist() if hasattr(v, "tolist") else v for k, v in data.items()}
 
     def on_train_loss_report(self, train_info: dict):
-        wandb.log(train_info, step=train_info.get("iteration"))
+        wandb.log(
+            self._convert_to_serializable(train_info), step=train_info.get("iteration")
+        )
         if self.wrapped_callback:
             self.wrapped_callback.on_train_loss_report(train_info)
 
     def on_val_loss_report(self, val_info: dict):
-        wandb.log(val_info, step=val_info.get("iteration"))
+        wandb.log(
+            self._convert_to_serializable(val_info), step=val_info.get("iteration")
+        )
         if self.wrapped_callback:
             self.wrapped_callback.on_val_loss_report(val_info)
+
+
+class SwanLabCallback(TrainingCallback):
+    def __init__(
+        self,
+        project_name: str,
+        log_dir: str,
+        config: dict,
+        wrapped_callback: TrainingCallback = None,
+    ):
+        if swanlab is None:
+            raise ImportError(
+                "swanlab is not installed. please install swanlab via: pip install swanlab",
+            )
+        self.wrapped_callback = wrapped_callback
+        swanlab.init(
+            project=project_name,
+            experiment_name=os.path.basename(log_dir),
+            logdir=os.path.join(log_dir, "swanlog"),
+            config=config,
+        )
+
+    def _convert_to_serializable(self, data: dict) -> dict:
+        return {k: v.tolist() if hasattr(v, "tolist") else v for k, v in data.items()}
+
+    def on_train_loss_report(self, train_info: dict):
+        swanlab.log(
+            self._convert_to_serializable(train_info), step=train_info.get("iteration")
+        )
+        if self.wrapped_callback:
+            self.wrapped_callback.on_train_loss_report(train_info)
+
+    def on_val_loss_report(self, val_info: dict):
+        swanlab.log(
+            self._convert_to_serializable(val_info), step=val_info.get("iteration")
+        )
+        if self.wrapped_callback:
+            self.wrapped_callback.on_val_loss_report(val_info)
+
+
+SUPPORT_CALLBACK = {
+    "wandb": WandBCallback,
+    "swanlab": SwanLabCallback,
+}
+
+
+def get_reporting_callbacks(
+    report_to: str = None,
+    project_name: str = None,
+    log_dir: str = None,
+    config: str = None,
+):
+    if report_to is None or report_to == "":
+        return None
+    report_to = [item.strip().lower() for item in report_to.split(",") if item.strip()]
+    training_callback = None
+    for callback in report_to:
+        try:
+            training_callback = SUPPORT_CALLBACK[callback](
+                project_name=project_name,
+                log_dir=log_dir,
+                config=config,
+                wrapped_callback=training_callback,
+            )
+        except KeyError as e:
+            raise ValueError(
+                f"{callback} callback doesn't exist "
+                f"choose from {', '.join(SUPPORT_CALLBACK.keys())}"
+            ) from e
+
+    return training_callback

mlx_lm/tuner/losses.py

@@ -4,7 +4,13 @@ import mlx.core as mx
 import mlx.nn as nn
 
 
+def can_run_metal():
+    return mx.default_device() == mx.gpu and mx.metal.is_available()
+
+
 def _make_kl_forward_kernel():
+    if not can_run_metal():
+        return
     source = """
     constexpr int M = 4;
     constexpr int block = 1024 * M;
@@ -171,6 +177,8 @@ def _make_kl_forward_kernel():
 
 
 def _make_kl_backward_kernel():
+    if not can_run_metal():
+        return
     source = """
     constexpr int M = 4;
     constexpr int block = 1024 * M;
@@ -367,7 +375,7 @@ def _kl_div_loss(primals, cotangent, output):
 
 
 def kl_div_loss(logits_q, logits_p):
-    if mx.metal.is_available():
+    if can_run_metal():
         return _kl_div_loss(logits_q, logits_p)
     else:
         return nn.losses.kl_div_loss(
@@ -376,3 +384,415 @@ def kl_div_loss(logits_q, logits_p):
             axis=-1,
             reduction="none",
         )
+
+
+def _make_js_forward_kernel():
+    if not can_run_metal():
+        return
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx;
+    out_kl_q += out_idx;
+
+    float lse_p;
+    float lse_q;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q = max_q + metal::fast::log(sum_exp_q);
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl_p = 0;
+        float kl_q = 0;
+        const float logtwo = metal::fast::log(static_cast<float>(2));
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl_p and kl_q
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                float logp_j = vals_p[j] - lse_p;
+                float logq_j = vals_q[j] - lse_q;
+                float p_j = metal::fast::exp(logp_j);
+                float q_j = metal::fast::exp(logq_j);
+                kl_p += p_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logq_j - logp_j)));
+                kl_q += q_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logp_j - logq_j)));
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                float logp_j = vals_p[j] - lse_p;
+                float logq_j = vals_q[j] - lse_q;
+                float p_j = metal::fast::exp(logp_j);
+                float q_j = metal::fast::exp(logq_j);
+                kl_p += p_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logq_j - logp_j)));
+                kl_q += q_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logp_j - logq_j)));
+            }
+        }
+
+        // Add the kl_p and kl_q across the threadgroup
+        kl_p = simd_sum(kl_p);
+        kl_q = simd_sum(kl_q);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = kl_p;
+            shared[simd_group_id * 2 + 1] = kl_q;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        kl_p = shared[simd_lane_id * 2 + 0];
+        kl_q = shared[simd_lane_id * 2 + 1];
+        kl_p = simd_sum(kl_p);
+        kl_q = simd_sum(kl_q);
+
+        if (thread_index_in_threadgroup == 0) {
+            out[0] = static_cast<T>(0.5 * kl_p + 0.5 * kl_q);
+            out_kl_q[0] = static_cast<T>(kl_q);
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="js_forward",
+        input_names=["logits_q", "logits_p"],
+        output_names=["out", "out_kl_q"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+def _make_js_backward_kernel():
+    if not can_run_metal():
+        return
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out_q += out_idx * V;
+    cotan += out_idx;
+    output_kl_q += out_idx;
+
+    float lse_q;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q = max_q + metal::fast::log(sum_exp_q);
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float c = cotan[0];
+        const float logtwo = metal::fast::log(static_cast<float>(2));
+        float kl_q = output_kl_q[0];
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and compute vjp for logits_q
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                float logp_j = vals_p[j] - lse_p;
+                float logq_j = vals_q[j] - lse_q;
+                float q_j = metal::fast::exp(logq_j);
+                out_q[offset + j] = static_cast<T>(
+                    c * 0.5 * q_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logp_j - logq_j)) - kl_q)
+                );
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                if (offset + j < V) {
+                    float logp_j = vals_p[j] - lse_p;
+                    float logq_j = vals_q[j] - lse_q;
+                    float q_j = metal::fast::exp(logq_j);
+                    out_q[offset + j] = static_cast<T>(
+                        c * 0.5 * q_j * (logtwo - metal::fast::log(1 + metal::fast::exp(logp_j - logq_j)) - kl_q)
+                    );
+                }
+            }
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="js_backward",
+        input_names=["logits_q", "logits_p", "cotan", "output_kl_q"],
+        output_names=["out_q"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+_js_forward_kernel = _make_js_forward_kernel()
+_js_backward_kernel = _make_js_backward_kernel()
+
+
+@mx.custom_function
+def _js_div_loss(logits_q, logits_p):
+    n_outs = logits_q.size // logits_q.shape[-1]
+    dt = logits_q.dtype
+
+    outputs = _js_forward_kernel(
+        inputs=[logits_q, logits_p],
+        output_shapes=[logits_q.shape[:-1], logits_q.shape[:-1]],
+        output_dtypes=[dt, dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, n_outs, 1),
+        threadgroup=(1024, 1, 1),
+    )
+    return outputs[0], mx.stop_gradient(outputs[1])
+
+
+@_js_div_loss.vjp
+def _js_div_loss(primals, cotangents, outputs):
+    logits_q, logits_p = primals
+    cotan, _ = cotangents
+    _, kl_q = outputs
+    dt = logits_q.dtype
+
+    dp = mx.zeros_like(logits_p)
+    dq = _js_backward_kernel(
+        inputs=[logits_q, logits_p, cotan, kl_q],
+        output_shapes=[logits_q.shape],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, cotan.size, 1),
+        threadgroup=(1024, 1, 1),
+    )
+    return dq, dp
+
+
+def js_div_loss(logits_q, logits_p):
+    if can_run_metal():
+        return _js_div_loss(logits_q, logits_p)[0]
+    else:
+        logprobs_p = logits_p - mx.logsumexp(logits_p, axis=-1, keepdims=True)
+        logprobs_q = logits_q - mx.logsumexp(logits_q, axis=-1, keepdims=True)
+        logprobs_m = (
+            logprobs_p
+            + mx.log(1 + mx.exp(logprobs_q - logprobs_p))
+            - mx.log(2).astype(logits_q.dtype)
+        )
+        kl_p = nn.losses.kl_div_loss(logprobs_m, logprobs_p, axis=-1, reduction="none")
+        kl_q = nn.losses.kl_div_loss(logprobs_m, logprobs_q, axis=-1, reduction="none")
+        return 0.5 * (kl_p + kl_q)

mlx_lm/tuner/trainer.py

@@ -102,13 +102,14 @@ def iterate_batches(
 
     # If running in distributed mode (N machines) then each one should skip N-1
     # samples
+    offset = mx.distributed.init().rank()
     step = mx.distributed.init().size()
     if batch_size % step != 0:
         raise ValueError("The batch size must be divisible by the number of workers")
 
     # Make the batches:
     batch_idx = [
-        idx[i : i + batch_size : step]
+        idx[i + offset : i + offset + batch_size : step]
         for i in range(0, len(idx) - batch_size + 1, batch_size)
     ]
 
@@ -196,7 +197,8 @@ def train(
     iterate_batches: callable = iterate_batches,
     training_callback: TrainingCallback = None,
 ):
-    mx.set_wired_limit(mx.metal.device_info()["max_recommended_working_set_size"])
+    if mx.metal.is_available():
+        mx.set_wired_limit(mx.metal.device_info()["max_recommended_working_set_size"])
     print(f"Starting training..., iters: {args.iters}")
     world = mx.distributed.init()
     world_size = world.size()

mlx_lm/tuner/utils.py

@@ -88,6 +88,7 @@ def linear_to_lora_layers(
         "mistral",
         "mistral3",
         "llama",
+        "lfm2",
         "phi",
         "mixtral",
         "nemotron",
@@ -119,6 +120,14 @@ def linear_to_lora_layers(
         "ernie4_5",
         "dots1",
         "smollm3",
+        "exaone4",
+        "hunyuan_v1_dense",
+        "gpt_oss",
+        "ernie4_5_moe",
+        "granitemoe",
+        "longcat_flash",
+        "seed_oss",
+        "apertus",
     }:
         keys = {"self_attn.q_proj", "self_attn.v_proj"}
         if model.model_type in ["mixtral", "phimoe"]:
@@ -128,7 +137,8 @@ def linear_to_lora_layers(
             keys.add("mlp.shared_expert_gate")
         if model.model_type in ["olmoe", "qwen3_moe", "dots1"]:
             keys.add("mlp.gate")
-
+        if model.model_type in ["longcat_flash"]:
+            keys.add("mlp.router.classifier")
     elif model.model_type == "gpt_bigcode":
         keys = {"attn.c_attn"}
     elif model.model_type == "gpt2":
@@ -147,7 +157,12 @@ def linear_to_lora_layers(
         keys = {"norm_attn_norm.attn.Wqkv", "ffn.router.layer"}
     elif model.model_type == "internlm2":
         keys = {"attention.wqkv", "attention.wo"}
-    elif model.model_type in {"deepseek_v2", "deepseek_v3", "minicpm3"}:
+    elif model.model_type in {
+        "deepseek_v2",
+        "deepseek_v3",
+        "longcat_flash",
+        "minicpm3",
+    }:
         keys = {
             "self_attn.q_proj",
             "self_attn.q_a_proj",
@@ -164,6 +179,15 @@ def linear_to_lora_layers(
         }
     elif model.model_type == "exaone":
         keys = {"attn.attention.q_proj", "attn.attention.v_proj"}
+    elif model.model_type == "bailing_moe":
+        keys = {"attention.query_key_value", "attention.dense"}
+    elif model.model_type == "nemotron_h":
+        keys.add("mixer.in_proj")
+        keys.add("mixer.out_proj")
+        keys.add("mixer.q_proj")
+        keys.add("mixer.k_proj")
+        keys.add("mixer.v_proj")
+        keys.add("mixer.o_proj")
     else:
         raise ValueError(f"Lora does not support {model.model_type}")
 

mlx_lm/utils.py

@@ -41,9 +41,11 @@ from .tuner.utils import get_total_parameters, load_adapters
 
 # Constants
 MODEL_REMAPPING = {
-    "mistral": "llama",  # mistral is compatible with llama
+    "mistral": "llama",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
+    "kimi_k2": "deepseek_v3",
+    "qwen2_5_vl": "qwen2_vl",
 }
 
 MAX_FILE_SIZE_GB = 5
@@ -79,7 +81,9 @@ def compute_bits_per_weight(model):
     return model_bytes * 8 / model_params
 
 
-def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path:
+def get_model_path(
+    path_or_hf_repo: str, revision: Optional[str] = None
+) -> Tuple[Path, Optional[str]]:
     """
     Ensures the model is available locally. If the path does not exist locally,
     it is downloaded from the Hugging Face Hub.
@@ -101,7 +105,7 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
                 revision=revision,
                 allow_patterns=[
                     "*.json",
-                    "*.safetensors",
+                    "model*.safetensors",
                     "*.py",
                     "tokenizer.model",
                     "*.tiktoken",
@@ -140,7 +144,7 @@ def load_model(
     strict: bool = True,
     model_config: dict = {},
     get_model_classes: Callable[[dict], Tuple[Type[nn.Module], Type]] = _get_classes,
-) -> nn.Module:
+) -> Tuple[nn.Module, dict]:
     """
     Load and initialize the model from a given path.
 
@@ -158,7 +162,7 @@ def load_model(
             Defaults to the ``_get_classes`` function.
 
     Returns:
-        nn.Module: The loaded and initialized model.
+        Tuple[nn.Module, dict[str, Any]]: The loaded and initialized model and config.
 
     Raises:
         FileNotFoundError: If the weight files (.safetensors) are not found.
@@ -169,10 +173,6 @@ def load_model(
 
     weight_files = glob.glob(str(model_path / "model*.safetensors"))
 
-    if not weight_files:
-        # Try weight for back-compat
-        weight_files = glob.glob(str(model_path / "weight*.safetensors"))
-
     if not weight_files and strict:
         logging.error(f"No safetensors found in {model_path}")
         raise FileNotFoundError(f"No safetensors found in {model_path}")
@@ -189,8 +189,7 @@ def load_model(
     if hasattr(model, "sanitize"):
         weights = model.sanitize(weights)
 
-    if (quantization := config.get("quantization", None)) is not None:
-
+    def _quantize(quantization):
         def class_predicate(p, m):
             # Handle custom per layer quantizations
             if p in config["quantization"]:
@@ -203,8 +202,13 @@ def load_model(
             model,
             group_size=quantization["group_size"],
             bits=quantization["bits"],
+            mode=quantization.get("mode", "affine"),
             class_predicate=class_predicate,
         )
+
+    if (quantization := config.get("quantization", None)) is not None:
+        _quantize(quantization)
+
     elif quantization_config := config.get("quantization_config", False):
         # Handle legacy quantization config
         quant_method = quantization_config["quant_method"]
@@ -212,8 +216,13 @@ def load_model(
             from .models.bitlinear_layers import bitnet_quantize
 
             model = bitnet_quantize(model, quantization_config)
+        elif quant_method == "mxfp4":
+            quantization = {"group_size": 32, "bits": 4, "mode": "mxfp4"}
+            config["quantization"] = quantization
+            config["quantization_config"] = quantization
+            _quantize(quantization)
         else:
-            raise ValueError(f"Unsupported quantization method {quant_method}")
+            raise ValueError(f"Unknown quantization method {quant_method}.")
 
     model.load_weights(list(weights.items()), strict=strict)
 
@@ -445,20 +454,22 @@ def save_model(
 def quantize_model(
     model: nn.Module,
     config: dict,
-    q_group_size: int,
-    q_bits: int,
+    group_size: int,
+    bits: int,
+    mode: str = "affine",
     quant_predicate: Optional[
         Callable[[str, nn.Module, dict], Union[bool, dict]]
     ] = None,
-) -> Tuple:
+) -> Tuple[nn.Module, dict]:
     """
     Applies quantization to the model weights.
 
     Args:
         model (nn.Module): The model to be quantized.
         config (dict): Model configuration.
-        q_group_size (int): Group size for quantization.
-        q_bits (int): Bits per weight for quantization.
+        group_size (int): Group size for quantization.
+        bits (int): Bits per weight for quantization.
+        mode (str): The quantization mode.
         quant_predicate (Callable): A callable that decides how
             to quantize each layer based on the path.
             Accepts the layer `path`, the `module` and the model `config`.
@@ -468,31 +479,38 @@ def quantize_model(
     Returns:
         Tuple: Tuple containing quantized model and config.
     """
-    if "quantization" in config:
-        raise ValueError("Cannot quantize already quantized model")
     quantized_config = copy.deepcopy(config)
-    quantized_config["quantization"] = {"group_size": q_group_size, "bits": q_bits}
 
-    def base_predicate(path, module):
+    quant_predicate = quant_predicate or getattr(model, "quant_predicate", None)
+    quant_params = {"group_size": group_size, "bits": bits, "mode": mode}
+    if "quantization" in quantized_config:
+        # If the model is already partially quantized, return params so that
+        # the config is set on a per-layer basis
+        fine_grained_config = True
+    else:
+        fine_grained_config = False
+        quantized_config["quantization"] = quant_params
+
+    def wrapped_predicate(path, module):
         if not hasattr(module, "to_quantized"):
             return False
-        if module.weight.shape[-1] % q_group_size != 0:
+        if module.weight.shape[-1] % group_size != 0:
             return False
-        return True
-
-    # Add any custom quantization parameters to the config as we go
-    def wrapped_predicate(p, m):
-        bool_or_params = base_predicate(p, m)
-        if bool_or_params:
-            bool_or_params = quant_predicate(p, m, config)
-        quantized_config["quantization"][p] = bool_or_params
+        bool_or_params = True
+        if quant_predicate is not None:
+            bool_or_params = quant_predicate(path, module)
+        if isinstance(bool_or_params, dict):
+            quantized_config["quantization"][path] = bool_or_params
+        elif fine_grained_config and bool_or_params:
+            quantized_config["quantization"][path] = quant_params
         return bool_or_params
 
     nn.quantize(
         model,
-        q_group_size,
-        q_bits,
-        class_predicate=wrapped_predicate if quant_predicate else base_predicate,
+        group_size,
+        bits,
+        mode=mode,
+        class_predicate=wrapped_predicate,
     )
     # support hf model tree #957
     quantized_config["quantization_config"] = quantized_config["quantization"]

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.25.0
+mlx>=0.29.0
 numpy
 transformers>=4.39.3
 protobuf

setup.py

@@ -15,7 +15,7 @@ from _version import __version__
 setup(
     name="mlx-lm",
     version=__version__,
-    description="LLMs on Apple silicon with MLX and the Hugging Face Hub",
+    description="LLMs with MLX and the Hugging Face Hub",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
     readme="README.md",
@@ -27,9 +27,9 @@ setup(
     packages=["mlx_lm", "mlx_lm.models", "mlx_lm.quant", "mlx_lm.tuner"],
     python_requires=">=3.8",
     extras_require={
-        "test": ["datasets"],
+        "test": ["datasets", "lm-eval"],
+        "train": ["datasets", "tqdm"],
         "evaluate": ["lm-eval", "tqdm"],
-        "quant": ["datasets", "tqdm"],
     },
     entry_points={
         "console_scripts": [
@@ -37,6 +37,7 @@ setup(
             "mlx_lm.dwq = mlx_lm.quant.dwq:main",
             "mlx_lm.dynamic_quant = mlx_lm.quant.dynamic_quant:main",
             "mlx_lm.gptq = mlx_lm.quant.gptq:main",
+            "mlx_lm.benchmark = mlx_lm.benchmark:main",
             "mlx_lm.cache_prompt = mlx_lm.cache_prompt:main",
             "mlx_lm.chat = mlx_lm.chat:main",
             "mlx_lm.convert = mlx_lm.convert:main",
@@ -44,6 +45,7 @@ setup(
             "mlx_lm.fuse = mlx_lm.fuse:main",
             "mlx_lm.generate = mlx_lm.generate:main",
             "mlx_lm.lora = mlx_lm.lora:main",
+            "mlx_lm.perplexity = mlx_lm.perplexity:main",
             "mlx_lm.server = mlx_lm.server:main",
             "mlx_lm.manage = mlx_lm.manage:main",
             "mlx_lm.upload = mlx_lm.upload:main",

tests/test_chat.py

@@ -0,0 +1,167 @@
+import argparse
+import unittest
+from unittest.mock import MagicMock, patch
+
+from mlx_lm.chat import setup_arg_parser
+
+
+class TestChat(unittest.TestCase):
+
+    def test_setup_arg_parser_system_prompt(self):
+        parser = setup_arg_parser()
+
+        # Test default (no system prompt)
+        args = parser.parse_args([])
+        self.assertIsNone(args.system_prompt)
+
+        # Test with system prompt
+        args = parser.parse_args(["--system-prompt", "You are a helpful assistant."])
+        self.assertEqual(args.system_prompt, "You are a helpful assistant.")
+
+    def test_setup_arg_parser_all_args(self):
+        parser = setup_arg_parser()
+        args = parser.parse_args(
+            [
+                "--model",
+                "test-model",
+                "--adapter-path",
+                "/path/to/adapter",
+                "--temp",
+                "0.7",
+                "--top-p",
+                "0.9",
+                "--xtc-probability",
+                "0.1",
+                "--xtc-threshold",
+                "0.1",
+                "--seed",
+                "42",
+                "--max-kv-size",
+                "1024",
+                "--max-tokens",
+                "512",
+                "--system-prompt",
+                "You are a helpful assistant.",
+            ]
+        )
+
+        self.assertEqual(args.model, "test-model")
+        self.assertEqual(args.adapter_path, "/path/to/adapter")
+        self.assertEqual(args.temp, 0.7)
+        self.assertEqual(args.top_p, 0.9)
+        self.assertEqual(args.xtc_probability, 0.1)
+        self.assertEqual(args.xtc_threshold, 0.1)
+        self.assertEqual(args.seed, 42)
+        self.assertEqual(args.max_kv_size, 1024)
+        self.assertEqual(args.max_tokens, 512)
+        self.assertEqual(args.system_prompt, "You are a helpful assistant.")
+
+    @patch("mlx_lm.chat.load")
+    @patch("mlx_lm.chat.make_prompt_cache")
+    @patch("mlx_lm.chat.stream_generate")
+    @patch("builtins.input")
+    @patch("builtins.print")
+    def test_system_prompt_in_messages(
+        self,
+        mock_print,
+        mock_input,
+        mock_stream_generate,
+        mock_make_prompt_cache,
+        mock_load,
+    ):
+        from mlx_lm.chat import main
+
+        # Mock the model and tokenizer
+        mock_model = MagicMock()
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.return_value = "processed_prompt"
+        mock_load.return_value = (mock_model, mock_tokenizer)
+
+        # Mock prompt cache
+        mock_prompt_cache = MagicMock()
+        mock_make_prompt_cache.return_value = mock_prompt_cache
+
+        # Mock stream_generate to return some responses
+        mock_response = MagicMock()
+        mock_response.text = "Hello there!"
+        mock_stream_generate.return_value = [mock_response]
+
+        # Mock user input: first a question, then 'q' to quit
+        mock_input.side_effect = ["What is the weather?", "q"]
+
+        # Test with system prompt
+        with patch(
+            "sys.argv", ["chat.py", "--system-prompt", "You are a weather assistant."]
+        ):
+            try:
+                main()
+            except SystemExit:
+                pass
+
+        # Verify that apply_chat_template was called with system prompt
+        mock_tokenizer.apply_chat_template.assert_called()
+        call_args = mock_tokenizer.apply_chat_template.call_args[0][
+            0
+        ]  # First positional arg (messages)
+
+        # Check that the messages contain both system and user messages
+        self.assertEqual(len(call_args), 2)
+        self.assertEqual(call_args[0]["role"], "system")
+        self.assertEqual(call_args[0]["content"], "You are a weather assistant.")
+        self.assertEqual(call_args[1]["role"], "user")
+        self.assertEqual(call_args[1]["content"], "What is the weather?")
+
+    @patch("mlx_lm.chat.load")
+    @patch("mlx_lm.chat.make_prompt_cache")
+    @patch("mlx_lm.chat.stream_generate")
+    @patch("builtins.input")
+    @patch("builtins.print")
+    def test_no_system_prompt_in_messages(
+        self,
+        mock_print,
+        mock_input,
+        mock_stream_generate,
+        mock_make_prompt_cache,
+        mock_load,
+    ):
+        from mlx_lm.chat import main
+
+        # Mock the model and tokenizer
+        mock_model = MagicMock()
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.return_value = "processed_prompt"
+        mock_load.return_value = (mock_model, mock_tokenizer)
+
+        # Mock prompt cache
+        mock_prompt_cache = MagicMock()
+        mock_make_prompt_cache.return_value = mock_prompt_cache
+
+        # Mock stream_generate to return some responses
+        mock_response = MagicMock()
+        mock_response.text = "Hello there!"
+        mock_stream_generate.return_value = [mock_response]
+
+        # Mock user input: first a question, then 'q' to quit
+        mock_input.side_effect = ["What is the weather?", "q"]
+
+        # Test without system prompt
+        with patch("sys.argv", ["chat.py"]):
+            try:
+                main()
+            except SystemExit:
+                pass
+
+        # Verify that apply_chat_template was called without system prompt
+        mock_tokenizer.apply_chat_template.assert_called()
+        call_args = mock_tokenizer.apply_chat_template.call_args[0][
+            0
+        ]  # First positional arg (messages)
+
+        # Check that the messages contain only user message
+        self.assertEqual(len(call_args), 1)
+        self.assertEqual(call_args[0]["role"], "user")
+        self.assertEqual(call_args[0]["content"], "What is the weather?")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_evaluate.py

@@ -0,0 +1,59 @@
+# Copyright © 2024 Apple Inc.
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import mlx.core as mx
+
+from mlx_lm.evaluate import MLXLM
+
+
+class TestMLXLM(unittest.TestCase):
+    def setUp(self):
+        # Mock the load function to avoid loading actual models
+        self.mock_model = MagicMock()
+        self.mock_tokenizer = MagicMock()
+        self.mock_tokenizer.model_max_length = 2048
+        self.mock_tokenizer.chat_template = None
+        self.mock_tokenizer.encode = MagicMock(return_value=[1, 2, 3, 4])
+
+        with patch("mlx_lm.evaluate.load") as mock_load:
+            mock_load.return_value = (self.mock_model, self.mock_tokenizer)
+            self.mlx_lm = MLXLM("test_model")
+
+    def test_loglikelihood_rolling_processes_all_inputs(self):
+        """Test that loglikelihood_rolling processes all inputs correctly when batching."""
+        # Create 5 mock requests to test batching with batch_size=2
+        mock_requests = [MagicMock(args=(f"text {i}",)) for i in range(5)]
+
+        # Mock inputs
+        test_inputs = [(i, i + 1, i + 2) for i in range(5)]
+        self.mlx_lm._tokenize = MagicMock(return_value=test_inputs)
+
+        # Mock _score_fn to return different scores for each batch
+        def mock_score_fn(batch):
+            batch_size = len(batch)
+            scores = mx.array([[0.1] * 3 for _ in range(batch_size)])
+            lengths = mx.array([3] * batch_size)
+            return scores, lengths, None
+
+        self.mlx_lm._score_fn = MagicMock(side_effect=mock_score_fn)
+        self.mlx_lm._batch_size = 2
+
+        result = self.mlx_lm.loglikelihood_rolling(mock_requests)
+
+        # Should return 5 results (one per request)
+        self.assertEqual(len(result), 5)
+
+        # Should have called _score_fn 3 times (batches of 2, 2, 1)
+        self.assertEqual(self.mlx_lm._score_fn.call_count, 3)
+
+        # Verify the batches were correct sizes
+        call_args_list = self.mlx_lm._score_fn.call_args_list
+        self.assertEqual(len(call_args_list[0][0][0]), 2)  # First batch: 2 items
+        self.assertEqual(len(call_args_list[1][0][0]), 2)  # Second batch: 2 items
+        self.assertEqual(len(call_args_list[2][0][0]), 1)  # Third batch: 1 item
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_generate.py

@@ -147,8 +147,8 @@ class TestGenerate(unittest.TestCase):
 
         self.assertEqual("TEST", response)
         num_embeddings = prompt_embeddings.shape[0]
-        self.assertEqual(
-            num_embeddings / prefill_step_size, num_prompt_processing_callbacks
+        self.assertTrue(
+            num_embeddings / prefill_step_size < num_prompt_processing_callbacks
         )
 
 

tests/test_losses.py

@@ -0,0 +1,64 @@
+# Copyright © 2025 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+
+from mlx_lm.tuner.losses import can_run_metal, js_div_loss, kl_div_loss
+
+
+class TestLosses(unittest.TestCase):
+
+    def test_kl_div_loss(self):
+        self.assertTrue(can_run_metal())
+
+        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+
+        with mx.stream(mx.cpu):
+            expected = kl_div_loss(logits_q, logits_p)
+        kl = kl_div_loss(logits_q, logits_p)
+
+        self.assertTrue(mx.allclose(kl, expected, rtol=1e-4))
+
+    def test_js_div_loss(self):
+        self.assertTrue(can_run_metal())
+
+        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+
+        with mx.stream(mx.cpu):
+            expected = js_div_loss(logits_q, logits_p)
+        js = js_div_loss(logits_q, logits_p)
+
+        self.assertTrue(mx.allclose(js, expected))
+
+    def test_kl_div_loss_vjp(self):
+        self.assertTrue(can_run_metal())
+
+        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
+
+        with mx.stream(mx.cpu):
+            expected = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]
+        vjp_q = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]
+
+        self.assertTrue(mx.allclose(vjp_q, expected))
+
+    def test_js_div_loss_vjp(self):
+        self.assertTrue(can_run_metal())
+
+        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
+        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)
+
+        with mx.stream(mx.cpu):
+            expected = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]
+        vjp_q = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]
+
+        self.assertTrue(mx.allclose(vjp_q, expected))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_models.py

@@ -147,6 +147,24 @@ class TestModels(unittest.TestCase):
         out2 = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask=mask)
         self.assertTrue(mx.allclose(out1[1, :, :2], out2[1, :, :2]))
 
+    def test_mask_with_window(self):
+        mask = create_causal_mask(5, 0, window_size=3)
+        expected_sums = mx.array([1, 2, 3, 3, 3])
+        sums = mask.sum(axis=1)
+        self.assertTrue(mx.array_equal(sums, expected_sums))
+
+        mask = create_causal_mask(5, 1, window_size=3)
+        self.assertEqual(mask.shape, (5, 6))
+        expected_sums = mx.array([2, 3, 3, 3, 3])
+        sums = mask.sum(axis=1)
+        self.assertTrue(mx.array_equal(sums, expected_sums))
+
+        mask = create_causal_mask(5, 2, window_size=3)
+        self.assertEqual(mask.shape, (5, 7))
+        expected_sums = mx.array([3, 3, 3, 3, 3])
+        sums = mask.sum(axis=1)
+        self.assertTrue(mx.array_equal(sums, expected_sums))
+
     def test_rope(self):
         rope = rope_utils.initialize_rope(32, base=100, traditional=False)
         self.assertTrue(isinstance(rope, nn.RoPE))
@@ -221,7 +239,7 @@ class TestModels(unittest.TestCase):
             self.assertEqual(outputs.shape, (1, 2, vocab_size))
             self.assertEqual(outputs.dtype, t)
 
-            if model_type not in ("mamba", "plamo2"):
+            if model_type not in ("mamba", "plamo2", "gpt_oss"):
                 mask = create_causal_mask(inputs.shape[1], 0).astype(t)
                 outputs = model(inputs, mask=mask)
                 self.assertEqual(outputs.shape, (1, 2, vocab_size))
@@ -251,6 +269,33 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_lfm2(self):
+        from mlx_lm.models import lfm2
+
+        args = lfm2.ModelArgs(
+            model_type="lfm2",
+            hidden_size=1024,
+            num_hidden_layers=4,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            norm_eps=1e-5,
+            vocab_size=10_000,
+            full_attn_idxs=[0, 1, 2],
+            rope_theta=10000,
+            block_dim=1024,
+            block_ffn_dim_multiplier=1.5,
+            block_auto_adjust_ff_dim=True,
+            block_ff_dim=2048,
+            block_multiple_of=256,
+            max_position_embeddings=1000,
+            conv_bias=True,
+            conv_L_cache=3,
+        )
+        model = lfm2.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_bitnet(self):
         from mlx_lm.models import bitnet
 
@@ -1012,6 +1057,38 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_hunyuan_v1_dense(self):
+        from mlx_lm.models import hunyuan_v1_dense
+
+        args = hunyuan_v1_dense.ModelArgs(
+            model_type="hunyuan_v1_dense",
+            hidden_size=128,
+            attention_bias=False,
+            intermediate_size=256,
+            num_attention_heads=4,
+            num_hidden_layers=4,
+            num_key_value_heads=2,
+            rms_norm_eps=1e-4,
+            rope_theta=1000,
+            vocab_size=1000,
+            use_qk_norm=True,
+            rope_scaling={
+                "alpha": 1000.0,
+                "factor": 1.0,
+                "type": "dynamic",
+                "beta_fast": 32,
+                "beta_slow": 1,
+                "mscale": 1.0,
+                "mscale_all_dim": 0.0,
+                "original_max_position_embeddings": 8192,
+            },
+            max_position_embeddings=32768,
+        )
+        model = hunyuan_v1_dense.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_olmo2(self):
         from mlx_lm.models import olmo2
 
@@ -1099,6 +1176,33 @@ class TestModels(unittest.TestCase):
             model, "smollm3", args.vocab_size, args.num_hidden_layers
         )
 
+    def test_gpt_oss(self):
+        from mlx_lm.models import gpt_oss
+
+        args = gpt_oss.ModelArgs(
+            model_type="gpt_oss",
+            hidden_size=1024,
+            num_hidden_layers=4,
+            intermediate_size=2048,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            num_local_experts=16,
+            num_experts_per_tok=2,
+            sliding_window=128,
+            rope_theta=10000,
+            vocab_size=10_000,
+            layer_types=[
+                "sliding_attention",
+                "full_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+        )
+        model = gpt_oss.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_server.py

@@ -1,6 +1,7 @@
 # Copyright © 2024 Apple Inc.
 
 import http
+import io
 import json
 import threading
 import unittest
@@ -79,15 +80,21 @@ class TestServer(unittest.TestCase):
             "top_p": 0.9,
             "repetition_penalty": 1.1,
             "repetition_context_size": 20,
+            "seed": 999,
             "stop": "stop sequence",
         }
 
         response = requests.post(url, json=post_data)
 
-        response_body = response.text
+        response_body = json.loads(response.text)
 
         self.assertIn("id", response_body)
         self.assertIn("choices", response_body)
+        first_text = response_body["choices"][0]["text"]
+        self.assertEqual(
+            first_text,
+            json.loads(requests.post(url, json=post_data).text)["choices"][0]["text"],
+        )
 
     def test_handle_chat_completions(self):
         url = f"http://localhost:{self.port}/v1/chat/completions"
@@ -508,5 +515,55 @@ class TestGetPromptCache(unittest.TestCase):
         self.assertEqual(self.handler.prompt_cache.model_key, ("model_v2", None, None))
 
 
+class TestKeepalive(unittest.TestCase):
+
+    def test_keepalive_callback(self):
+        """Test keepalive callback sends SSE comments and handles errors"""
+        from unittest.mock import Mock
+
+        # Mock handler
+        mock_wfile = io.BytesIO()
+        handler = Mock()
+        handler.wfile = mock_wfile
+
+        # Test callback logic (same as in server.py)
+        def keepalive_callback(processed_tokens, total_tokens):
+            if handler.stream:
+                try:
+                    handler.wfile.write(
+                        f": keepalive {processed_tokens}/{total_tokens}\n\n".encode()
+                    )
+                    handler.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError, OSError):
+                    pass
+
+        # Test streaming enabled
+        handler.stream = True
+        keepalive_callback(1024, 4096)
+
+        output = mock_wfile.getvalue().decode("utf-8")
+        self.assertEqual(output, ": keepalive 1024/4096\n\n")
+
+        # Test streaming disabled
+        handler.stream = False
+        mock_wfile.seek(0)
+        mock_wfile.truncate(0)
+        keepalive_callback(2048, 4096)
+
+        output = mock_wfile.getvalue().decode("utf-8")
+        self.assertEqual(output, "")
+
+        # Test error handling
+        handler.stream = True
+        handler.wfile = Mock()
+        handler.wfile.write.side_effect = BrokenPipeError("Connection broken")
+
+        # Should not raise exception
+        try:
+            keepalive_callback(3072, 4096)
+        except Exception as e:
+            self.fail(f"Callback should handle BrokenPipeError: {e}")
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_tokenizers.py

@@ -48,6 +48,9 @@ class TestTokenizers(unittest.TestCase):
         tokens = tokenizer.encode("こんにちは！私の名前はAI")
         check(tokens)
 
+        tokens = tokenizer.encode("⊕ ⊻ ∧ ¬")
+        check(tokens)
+
         tokens = tokenizer.encode("a ,b")
         check(tokens)
 
@@ -104,7 +107,7 @@ class TestTokenizers(unittest.TestCase):
 
         tokenizer_repo = "mlx-community/Llama-3.2-1B-Instruct-4bit"
         tokenizer = self.download_tokenizer(tokenizer_repo)
-        self.assertTrue(tokenizer.has_tool_calling, False)
+        self.assertFalse(tokenizer.has_tool_calling)
 
     def test_thinking(self):
         tokenizer_repo = "mlx-community/Qwen3-4B-4bit"

tests/test_tuner_trainer.py

@@ -0,0 +1,66 @@
+# Copyright © 2025 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+
+from mlx_lm.tuner.trainer import iterate_batches
+
+
+class MockDistributedGroup:
+    def __init__(self, rank, size):
+        self._rank = rank
+        self._size = size
+
+    def rank(self):
+        return self._rank
+
+    def size(self):
+        return self._size
+
+
+class MockDistributed:
+    def __init__(self):
+        self.rank = 0
+        self.size = 1
+
+    def init(self):
+        return MockDistributedGroup(self.rank, self.size)
+
+
+class TestTunerTrainer(unittest.TestCase):
+    def test_iterate_batches_ddp(self):
+        olddist = mx.distributed
+        try:
+            mx.distributed = MockDistributed()
+
+            def run(rank, size, batch):
+                mx.distributed.rank = rank
+                mx.distributed.size = size
+
+                data = mx.arange(128).reshape(-1, 1).tolist()
+                data = [(d, 0) for d in data]
+
+                samples = set()
+                for i, (b, l) in enumerate(iterate_batches(data, batch, 1)):
+                    samples.add(tuple(mx.flatten(b).tolist()))
+
+                ref_batches = mx.arange(128).reshape(-1, batch).tolist()
+                for b in ref_batches:
+                    self.assertTrue(tuple(b[rank::size]) in samples)
+
+            run(0, 1, 4)
+            run(0, 1, 8)
+            run(0, 2, 8)
+            run(1, 2, 8)
+            run(0, 4, 8)
+            run(1, 4, 8)
+            run(2, 4, 8)
+            run(3, 4, 8)
+
+        finally:
+            mx.distributed = olddist
+
+
+if __name__ == "__main__":
+    unittest.main()