some cleanup

.circleci/config.yml

@@ -1,66 +0,0 @@
-version: 2.1
-
-orbs:
-  apple: ml-explore/pr-approval@0.1.0
- 
-jobs:
-  linux_build_and_test:
-    docker:
-      - image: cimg/python:3.9
-
-    steps:
-      - checkout
-      - run:
-          name: Run style checks
-          command: |
-            pip install pre-commit
-            pre-commit run --all
-            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
-
-  mlx_lm_build_and_test:
-    macos:
-      xcode: "15.2.0"
-    resource_class: macos.m1.large.gen1
-    steps:
-      - checkout
-      - run:
-          name: Install dependencies
-          command: |
-            brew install python@3.9
-            python3.9 -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install unittest-xml-reporting
-            pip install -e ".[test]"
-      - run:
-          name: Run Python tests
-          command: |
-            source env/bin/activate
-            python -m xmlrunner discover -v tests -o test-results/
-      - store_test_results:
-          path: test-results
-
-workflows:
-  build_and_test:
-    when:
-      matches:
-        pattern: "^(?!pull/)[-\\w]+$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - mlx_lm_build_and_test
-      - linux_build_and_test
-
-  prb:
-    when:
-      matches:
-        pattern: "^pull/\\d+(/head)?$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - hold:
-          type: approval
-      - apple/authenticate:
-          context: pr-approval
-      - mlx_lm_build_and_test:
-          requires: [ hold ]
-      - linux_build_and_test:
-          requires: [ hold ]

.gitignore

@@ -1,139 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Vim
-*.swp
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-pip-wheel-metadata/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-.python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# IDE files
-.idea/
-.vscode/
-
-# .DS_Store files
-.DS_Store

.pre-commit-config.yaml

@@ -1,11 +0,0 @@
-repos:
--   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 25.1.0
-    hooks:
-    -   id: black
--   repo: https://github.com/pycqa/isort
-    rev: 6.0.0
-    hooks:
-    -   id: isort
-        args:
-            - --profile=black

ACKNOWLEDGMENTS.md

@@ -5,8 +5,13 @@ with a short description of your contribution(s) below. For example:
 
 - Jane Smith: Added the `foo` example.
 
-MLX LM was developed with contributions from the following individuals:
+MLX Examples was developed with contributions from the following individuals:
 
+- Juarez Bochi: Added support for T5 models.
+- Sarthak Yadav: Added the `cifar` and `speechcommands` examples.
 - Shunta Saito: Added support for PLaMo models.
+- Gabrijel Boduljak: Implemented `CLIP`.
+- Markus Enzweiler: Added the `cvae` examples.
 - Prince Canuma: Helped add support for `Starcoder2` 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`, 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`.
+- Shiyu Li: Added the `Segment Anything Model`.
+- Gökdeniz Gülmez: Added support for `MiniCPM`, `Mamba` and support for `full-fine-tuning`.

CONTRIBUTING.md

@@ -1,54 +1,11 @@
 # Contributing to MLX LM 
 
-We want to make contributing to this project as easy and transparent as
-possible.
-
-## Pull Requests
-
-1. Fork and submit pull requests to the repo. 
-2. If you've added code that should be tested, add tests.
-3. Every PR should have passing tests and at least one review. 
-4. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`.
-   This should install hooks for running `black` and `clang-format` to ensure
-   consistent style for C++ and python code.
- 
-   You can also run the formatters manually as follows on individual files:
- 
-     ```bash
-     clang-format -i file.cpp
-     ```
- 
-     ```bash
-     black file.py
-     ```
-
-     or,
-
-     ```bash
-     # single file
-     pre-commit run --files file1.py 
-
-     # specific files
-     pre-commit run --files file1.py file2.py
-     ```
- 
-   or run `pre-commit run --all-files` to check all files in the repo.
-
-## Issues
-
-We use GitHub issues to track public bugs. Please ensure your description is
-clear and has sufficient instructions to be able to reproduce the issue.
-
-## License
-
-By contributing to mlx-lm, you agree that your contributions will be licensed
-under the LICENSE file in the root directory of this source tree.
-
-## Adding New Models
-
 Below are some tips to port LLMs available on Hugging Face to MLX.
 
-From this directory, do an editable install:
+Before starting checkout the [general contribution
+guidelines](https://github.com/ml-explore/mlx-examples/blob/main/CONTRIBUTING.md).
+
+Next, from this directory, do an editable install:
 
 ```shell
 pip install -e .
@@ -60,7 +17,7 @@ Then check if the model has weights in the
 convert it.
 
 After that, add the model file to the
-[`mlx_lm/models`](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/models)
+[`mlx_lm/models`](https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/models)
 directory. You can see other examples there. We recommend starting from a model
 that is similar to the model you are porting.
 
@@ -78,12 +35,12 @@ To determine the model layer names, we suggest either:
   in the Hugging Face repo.
 
 To add LoRA support edit
-[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/utils.py#L27-L60)
+[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/tuner/utils.py#L27-L60)
 
 Finally, add a test for the new modle type to the [model
-tests](https://github.com/ml-explore/mlx-lm/blob/main/tests/test_models.py).
+tests](https://github.com/ml-explore/mlx-examples/blob/main/llms/tests/test_models.py).
 
-You can run the tests with:
+From the `llms/` directory, you can run the tests with:
 
 ```shell
 python -m unittest discover tests/

MANIFEST.in

@@ -1,2 +1,2 @@
-include requirements.txt
+include mlx_lm/requirements.txt
 recursive-include mlx_lm/ *.py

README.md

@@ -1,17 +1,4 @@
-## MLX LM 
-
-MLX LM is a Python package for generating text and fine-tuning large language
-models on Apple silicon with MLX.
-
-Some key features include:
-
-* Integration with the Hugging Face Hub to easily use thousands of LLMs with a
-  single command. 
-* Support for quantizing and uploading models to the Hugging Face Hub.
-* [Low-rank and full model
-  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`
+## Generate Text with LLMs and MLX
 
 The easiest way to get started is to install the `mlx-lm` package:
 
@@ -27,12 +14,18 @@ pip install mlx-lm
 conda install -c conda-forge mlx-lm
 ```
 
+The `mlx-lm` package also has:
+
+- [LoRA, QLoRA, and full fine-tuning](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md)
+- [Merging models](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/MERGE.md)
+- [HTTP model serving](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/SERVER.md)
+
 ### Quick Start
 
 To generate text with an LLM use:
 
 ```bash
-mlx_lm.generate --prompt "How tall is Mt Everest?"
+mlx_lm.generate --prompt "Hi!"
 ```
 
 To chat with an LLM use:
@@ -78,7 +71,7 @@ To see a description of all the arguments you can do:
 ```
 
 Check out the [generation
-example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/generate_response.py)
+example](https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/examples/generate_response.py)
 to see how to use the API in more detail.
 
 The `mlx-lm` package also comes with functionality to quantize and optionally
@@ -130,18 +123,6 @@ for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
 print()
 ```
 
-#### Sampling
-
-The `generate` and `stream_generate` functions accept `sampler` and
-`logits_processors` keyword arguments. A sampler is any callable which accepts
-a possibly batched logits array and returns an array of sampled tokens.  The
-`logits_processors` must be a list of callables which take the token history
-and current logits as input and return the processed logits. The logits
-processors are applied in order.
-
-Some standard sampling functions and logits processors are provided in
-`mlx_lm.sample_utils`.
-
 ### Command Line
 
 You can also use `mlx-lm` from the command line with:
@@ -183,7 +164,7 @@ mlx_lm.convert \
 ```
 
 Models can also be converted and quantized directly in the
-[mlx-my-repo](https://huggingface.co/spaces/mlx-community/mlx-my-repo) Hugging
+[mlx-my-repo]https://huggingface.co/spaces/mlx-community/mlx-my-repo) Hugging
 Face Space.
 
 ### Long Prompts and Generations 
@@ -220,17 +201,17 @@ The cached prompt is treated as a prefix to the supplied prompt. Also notice
 when using a cached prompt, the model to use is read from the cache and need
 not be supplied explicitly.
 
-Prompt caching can also be used in the Python API in order to avoid
+Prompt caching can also be used in the Python API in order to to avoid
 recomputing the prompt. This is useful in multi-turn dialogues or across
 requests that use the same context. See the
-[example](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/examples/chat.py)
+[example](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/examples/chat.py)
 for more usage details.
 
 ### Supported Models
 
 `mlx-lm` supports thousands of Hugging Face format LLMs. If the model you want to
 run is not supported, file an
-[issue](https://github.com/ml-explore/mlx-lm/issues/new) or better yet,
+[issue](https://github.com/ml-explore/mlx-examples/issues/new) or better yet,
 submit a pull request.
 
 Here are a few examples of Hugging Face models that work with this example:

mlx_lm/LEARNED_QUANTS.md

@@ -1,93 +0,0 @@
-# Learned Quantization 
-
-To reduce the quality loss from quantization MLX LM has two options:
-
-- Distilled Weight Quantization (DWQ)
-- Activation-aware Weight Quantization (AWQ)[^1].
-
-Both DWQ and AWQ use an example dataset to tune parameters of the model. DWQ
-fine-tunes non-quantized parameters (including quantization scales and biases)
-using the non-quantized model as a teacher. AWQ scales and clips the weights
-prior to quantization. The scaling and clipping values are found with a grid
-search minimizing the distance from the quantized hidden activations to the
-non-quantized hidden activations
-
-To get started, first install the requirements:
-
-```
-pip install mlx-lm[lwq]
-```
-
-### DWQ
-
-Use `mlx_lm.dwq` to run DWQ on a given model. For example:
-
-```bash
-mlx_lm.dwq --model mistralai/Mistral-7B-Instruct-v0.3
-```
-
-Some important options, along with their default values are:
-
-- `--mlx-path mlx_model`: The location to save the DWQ model.
-- `--bits 4`: Precision of the quantization.
-- `--num-samples 1024`: Number of samples to use. Using more samples can lead to
-  better results but takes longer.
-- `--batch-size 8`: Use a smaller batch size to reduce the memory footprint.
-
-For a full list of options run:
-
-```bash
-mlx_lm.dwq --help
-```
-
-### AWQ 
-
-Use `mlx_lm.awq` to run AWQ on a given model. For example:
-
-```bash
-mlx_lm.awq --model mistralai/Mistral-7B-Instruct-v0.3
-```
-
-The script can take anywhere form a few minutes to several hours to run
-depending on the model size and the number of samples.
-
-Some important options, along with their default values, are:
-
-- `--mlx-path mlx_model`: The location to save the AWQ model.
-- `--bits 4`: Precision of the quantization.
-- `--num-samples 32`: Number of samples to use. Using more samples can lead to
-  better results but takes longer.
-- `--n-grid 10`: The granularity of the AWQ search. A larger grid can lead to
-  better results but takes longer.
-
-For a full list of options run:
-
-```bash
-mlx_lm.awq --help
-```
-
-### Evaluate
-
-Once the training script finishes, you can evaluate the quality of the model
-on downstream tasks using `mlx_lm.evaluate`. For example:
-
-```bash
-mlx_lm.evaluate \
-    --model mlx_model \
-    --tasks winogrande boolq arc_challenge arc_easy hellaswag openbookqa piqa social_iqa                     
-```
-
-### Upload to Hugging Face
-
-Use `mlx_lm.upload` to upload the quantized model to the Hugging Face Hub. For
-example:
-
-```bash
-mlx_lm.upload \
-    --path mlx_model \
-    --upload-repo mlx-community/Mistral-7B-Instruct-v0.3-3bit-DWQ
-```
-
-[^1]: Refer to the [paper](https://arxiv.org/abs/2306.00978)
-and [github repository](https://github.com/mit-han-lab/llm-awq) for more
-details.

mlx_lm/LORA.md

@@ -76,14 +76,6 @@ You can specify the output location with `--adapter-path`.
 You can resume fine-tuning with an existing adapter with
 `--resume-adapter-file <path_to_adapters.safetensors>`.
 
-#### Prompt Masking
-
-The default training computes a loss for every token in the sample. You can
-ignore the prompt and compute loss for just the completion by passing
-`--mask-prompt`. Note this is only supported for `chat` and `completion`
-datasets. For `chat` datasets the final message in the message list is
-considered the completion. See the [dataset section](#Data) for more details. 
-
 ### Evaluate
 
 To compute test set perplexity use:
@@ -249,25 +241,14 @@ Refer to the documentation for the model you are fine-tuning for more details.
 {"prompt": "What is the capital of France?", "completion": "Paris."}
 ```
 
-For the `completions` data format, a different key can be used for the prompt
-and completion by specifying the following in the YAML config:
-
-```yaml
-prompt_feature: "input"
-completion_feature: "output"
-```
-
-Here, `"input"` is the expected key instead of the default `"prompt"`, and
-`"output"` is the expected key instead of `"completion"`. 
-
 `text`:
 
 ```jsonl
 {"text": "This is an example for the model."}
 ```
 
-Note, the format is automatically determined by the dataset. Note also, keys
-in each line not expected by the loader will be ignored.
+Note, the format is automatically determined by the dataset. Note also, keys in
+each line not expected by the loader will be ignored.
 
 > [!NOTE]
 > Each example in the datasets must be on a single line. Do not put more than
@@ -289,36 +270,20 @@ Otherwise, provide a mapping of keys in the dataset to the features MLX LM
 expects. Use a YAML config to specify the Hugging Face dataset arguments. For
 example:
 
-```yaml
+```
 hf_dataset:
-  path: "billsum"
+  name: "billsum"
   prompt_feature: "text"
   completion_feature: "summary"
 ```
 
 - Use `prompt_feature` and `completion_feature` to specify keys for a
   `completions` dataset. Use `text_feature` to specify the key for a `text`
-  dataset. Use `chat_feature` to specify the key for a chat dataset.
+  dataset. 
 
 - To specify the train, valid, or test splits, set the corresponding
   `{train,valid,test}_split` argument. 
 
-You can specify a list of Hugging Face datasets with a list of records each
-with the same structure as above. For example:
-
-```yaml
-hf_dataset:
-  - path: "Open-Orca/OpenOrca"
-    train_split: "train[:90%]"
-    valid_split: "train[-10%:]"
-    prompt_feature: "question"
-    completion_feature: "response"
-  - path: "trl-lib/ultrafeedback_binarized"
-    train_split: "train[:90%]"
-    valid_split: "train[-10%:]"
-    chat_feature: "chosen"
-```
-
 - Arguments specified in `config` will be passed as keyword arguments to
   [`datasets.load_dataset`](https://huggingface.co/docs/datasets/v2.20.0/en/package_reference/loading_methods#datasets.load_dataset).
 

mlx_lm/SERVER.md

@@ -86,13 +86,6 @@ curl localhost:8080/v1/chat/completions \
 - `adapters`: (Optional) A string path to low-rank adapters. The path must be
   relative to the directory the server was started in.
 
-- `draft_model`: (Optional) Specifies a smaller model to use for speculative
-  decoding. Set to `null` to unload.
-
-- `num_draft_tokens`: (Optional) The number of draft tokens the draft model
-  should predict at once. Defaults to `3`.
-
-
 ### Response Fields
 
 - `id`: A unique identifier for the chat.

mlx_lm/__init__.py

@@ -6,6 +6,4 @@ from ._version import __version__
 
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 
-from .convert import convert
-from .generate import generate, stream_generate
-from .utils import load
+from .utils import convert, generate, load, stream_generate

mlx_lm/__main__.py

@@ -1,28 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import importlib
-import sys
-
-if __name__ == "__main__":
-    subcommands = {
-        "awq",
-        "dwq",
-        "cache_prompt",
-        "chat",
-        "convert",
-        "evaluate",
-        "fuse",
-        "generate",
-        "lora",
-        "merge",
-        "server",
-        "manage",
-        "upload",
-    }
-    if len(sys.argv) < 2:
-        raise ValueError(f"CLI requires a subcommand in {subcommands}")
-    subcommand = sys.argv.pop(1)
-    if subcommand not in subcommands:
-        raise ValueError(f"CLI requires a subcommand in {subcommands}")
-    submodule = importlib.import_module(f"mlx_lm.{subcommand}")
-    submodule.main()

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.
 
-__version__ = "0.24.0"
+__version__ = "0.20.4"

mlx_lm/awq.py

@@ -1,605 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-import copy
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Callable, Dict
-from urllib import request
-
-import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map, tree_map_with_path
-from tqdm import tqdm
-
-from mlx_lm.models.base import create_attention_mask
-from mlx_lm.models.switch_layers import SwitchLinear
-from mlx_lm.utils import (
-    fetch_from_hub,
-    get_model_path,
-    save,
-)
-
-
-@dataclass
-class ScaleConfig:
-    prev: nn.Module
-    layers: list[nn.Module]
-    block: nn.Module | None = None
-    kwargs: list = field(default_factory=list)
-    use_config: Callable[[nn.Module], bool] | None = None
-
-
-@dataclass
-class AWQConfig:
-    embed: str
-    lm_head: str
-    no_clip: list[str]
-    scale_configs: list[ScaleConfig]
-    lm_key: str | None = None
-
-
-def update(cfg, **kwargs):
-    cfg = copy.deepcopy(cfg)
-    for k, v in kwargs.items():
-        setattr(cfg, k, v)
-    return cfg
-
-
-llama_awq = AWQConfig(
-    embed="embed_tokens",
-    lm_head="lm_head",
-    no_clip=["q_proj", "k_proj"],
-    scale_configs=[
-        ScaleConfig(
-            block="self_attn",
-            prev="input_layernorm",
-            layers=["q_proj", "k_proj", "v_proj"],
-            kwargs=["mask"],
-        ),
-        ScaleConfig(prev="mlp.up_proj", layers=["mlp.down_proj"]),
-        ScaleConfig(
-            block="mlp",
-            prev="post_attention_layernorm",
-            layers=["gate_proj", "up_proj"],
-        ),
-    ],
-)
-
-gemma3_text_awq = AWQConfig(
-    embed="embed_tokens",
-    lm_head="lm_head",
-    no_clip=["q_proj", "k_proj"],
-    scale_configs=[
-        ScaleConfig(
-            block="self_attn",
-            prev="input_layernorm",
-            layers=["q_proj", "k_proj", "v_proj"],
-            kwargs=["mask"],
-        ),
-        ScaleConfig(prev="mlp.up_proj", layers=["mlp.down_proj"]),
-        ScaleConfig(
-            block="mlp",
-            prev="pre_feedforward_layernorm",
-            layers=["gate_proj", "up_proj"],
-        ),
-    ],
-)
-
-gemma3_awq = update(gemma3_text_awq, lm_key="language_model")
-
-deepseek_v2_awq = AWQConfig(
-    embed="embed_tokens",
-    lm_head="lm_head",
-    no_clip=["q_proj", "q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"],
-    scale_configs=[
-        ScaleConfig(
-            block="self_attn",
-            prev="input_layernorm",
-            layers=["q_proj", "kv_a_proj_with_mqa"],
-            kwargs=["mask"],
-        ),
-        ScaleConfig(
-            prev="self_attn.kv_a_layernorm",
-            layers=["self_attn.kv_b_proj"],
-        ),
-        ScaleConfig(
-            prev="mlp.up_proj",
-            layers=["mlp.down_proj"],
-            use_config=lambda block: not "switch_mlp" in block.mlp,
-        ),
-        ScaleConfig(
-            prev="mlp.shared_experts.up_proj",
-            layers=["mlp.shared_experts.down_proj"],
-            use_config=lambda block: "switch_mlp" in block.mlp,
-        ),
-        ScaleConfig(
-            prev="mlp.switch_mlp.up_proj",
-            layers=["mlp.switch_mlp.down_proj"],
-            use_config=lambda block: "switch_mlp" in block.mlp,
-            kwargs=["indices"],
-        ),
-        ScaleConfig(
-            block="mlp",
-            prev="post_attention_layernorm",
-            layers=["gate_proj", "up_proj"],
-            use_config=lambda block: not "switch_mlp" in block.mlp,
-        ),
-        ScaleConfig(
-            block="mlp",
-            prev="post_attention_layernorm",
-            layers=[
-                "switch_mlp.gate_proj",
-                "switch_mlp.up_proj",
-                "shared_experts.gate_proj",
-                "shared_experts.up_proj",
-                "gate",  # not quantized, just scaled
-            ],
-            use_config=lambda block: "switch_mlp" in block.mlp,
-        ),
-    ],
-)
-
-AWQ_MODEL_CONFIGS = {
-    "llama": llama_awq,
-    "mistral": llama_awq,
-    "qwen2": llama_awq,
-    "qwen3": llama_awq,
-    "gemma3_text": gemma3_text_awq,
-    "gemma3": update(gemma3_text_awq, lm_key="language_model"),
-    "deepseek_v2": deepseek_v2_awq,
-}
-
-
-def mse(x, y):
-    return ((x - y).astype(mx.float32)) ** 2
-
-
-def submodule_from_key(module, key):
-    keys = key.split(".")
-    for k in keys:
-        module = module[k]
-    return module
-
-
-def run_layer(
-    layer: nn.Module,
-    x: mx.array,
-    indices: mx.array | None = None,
-    batch_size: int = 32,
-    **kwargs,
-):
-    y = []
-    for i in range(0, x.shape[0], batch_size):
-        if indices is not None:
-            y.append(
-                layer(x[i : i + batch_size], indices[i : i + batch_size], **kwargs)
-            )
-        else:
-            y.append(layer(x[i : i + batch_size], **kwargs))
-        mx.eval(y)
-    y = mx.concatenate(y, axis=0)
-    return y
-
-
-def dist_split(x: mx.array, group: mx.distributed.Group):
-    N = group.size()
-    if N == 1:
-        return x
-    B = x.shape[0]
-    assert B % N == 0
-    r = group.rank()
-    local_B = (B + N - 1) // N
-    return x[r * local_B : (r + 1) * local_B]
-
-
-def search_best_scale(
-    layers: list[nn.Module],
-    quantize_func: Callable,
-    block: nn.Module | None,
-    layer_kwargs: dict,
-    n_grid: int,
-):
-    group = mx.distributed.init()
-
-    layer_kwargs = layer_kwargs or {}
-
-    x = layers[0].input_feat
-
-    block = block or layers[0]
-    out = block(x, **layer_kwargs)
-
-    x_max = x.abs().mean(axis=(0, 1))
-
-    best_error = float("inf")
-    best_scales = None
-
-    weights = tree_flatten(block.parameters())
-
-    # Search across different scaling ratios
-    # and take the best loss.
-    for ratio in range(n_grid):
-        ratio = ratio / n_grid
-        scales = mx.maximum(x_max**ratio, 1e-4).reshape(-1)
-        scales = scales / (scales.max() * scales.min()).sqrt()
-        for layer in layers:
-            if isinstance(layer, (nn.Linear, SwitchLinear)):
-                layer.weight = quantize_func(layer.weight * scales) / scales
-
-        out_q = run_layer(block, x, **layer_kwargs)
-        loss = mse(out, out_q).sum()
-        if group is not None:
-            loss = mx.distributed.all_sum(loss) / group.size()
-        loss /= out.size
-        mx.eval(loss)
-        if loss.item() < best_error:
-            best_error = loss.item()
-            best_scales = scales
-
-        # reload the original weights
-        block.load_weights(weights)
-
-    best_scales = best_scales.reshape(-1)
-    mx.eval(best_scales)
-    return best_scales
-
-
-def apply_scale(prev_op, layers, scales):
-    # Fuse the scales into the previous op
-    if isinstance(prev_op, (nn.Linear, SwitchLinear)):
-        assert len(layers) == 1
-        prev_op.weight = prev_op.weight / scales[:, mx.newaxis]
-        if hasattr(prev_op, "bias"):
-            prev_op.bias = prev_op.bias / scales
-        layers[0].weight = layers[0].weight * scales
-    elif isinstance(prev_op, (nn.LayerNorm, nn.RMSNorm)):
-        prev_op.weight = prev_op.weight / scales
-        if hasattr(prev_op, "bias"):
-            prev_op.bias = prev_op.bias / scales
-        for layer in layers:
-            layer.weight = layer.weight * scales
-    elif prev_op.__class__.__name__ == "RMSNorm":  # For gemma models
-        dt = prev_op.weight.dtype
-        prev_op.weight = (
-            (1.0 + prev_op.weight.astype(mx.float32)) / scales - 1.0
-        ).astype(dt)
-        for layer in layers:
-            layer.weight = layer.weight * scales
-    else:
-        raise NotImplementedError(f"Could not apply scale to prev_op: {prev_op}")
-
-    for layer in layers:
-        if hasattr(layer, "input_feat"):
-            layer.input_feat = layer.input_feat / scales
-
-
-def scale_block(
-    block: nn.Module,
-    configs: list[ScaleConfig],
-    quantize_func: Callable,
-    layer_kwargs: dict,
-    n_grid: int,
-):
-    for conf in configs:
-        if conf.use_config is not None and not conf.use_config(block):
-            continue
-        if conf.block is not None:
-            local_block = block[conf.block]
-            layers = [submodule_from_key(local_block, l) for l in conf.layers]
-        else:
-            local_block = None
-            layers = [submodule_from_key(block, l) for l in conf.layers]
-        local_kwargs = {k: layer_kwargs[k] for k in conf.kwargs if k in layer_kwargs}
-        for k in conf.kwargs:
-            if hasattr(layers[0], k):
-                local_kwargs[k] = getattr(layers[0], k)
-
-        scales = search_best_scale(
-            layers=layers,
-            block=local_block,
-            layer_kwargs=local_kwargs,
-            quantize_func=quantize_func,
-            n_grid=n_grid,
-        )
-        apply_scale(submodule_from_key(block, conf.prev), layers, scales)
-
-
-def search_best_clip(
-    module: nn.Module,
-    quantize_func: Callable,
-    group_size: int,
-    n_grid: int,
-    max_shrink: float = 0.5,
-    batch_size: int = 64,
-    n_frames: int = 512,
-):
-    group = mx.distributed.init()
-
-    # subsample the input features
-    x = module.input_feat.flatten(0, 1)
-    stride = (x.shape[0] + n_frames - 1) // n_frames
-    x = x[::stride]
-
-    w = module.weight
-    x = x.reshape(x.shape[0], -1, group_size)
-
-    w_init_shape = w.shape
-    w_all = mx.flatten(w, 0, w.ndim - 2)
-    w_max_all = []
-
-    # batch across W to save memory
-    for b in range(0, w_all.shape[0], batch_size):
-        w = w_all[b : b + batch_size]
-
-        group_shape = (w.shape[0], w.shape[-1] // group_size)
-        best_error = mx.full(group_shape, float("inf"))
-        best_w_max = mx.zeros((*group_shape, 1), dtype=x.dtype)
-
-        w_shape = w.shape
-
-        w = w.reshape(*w.shape[:-1], -1, group_size)
-        out = mx.einsum("bdg,odg->bod", x, w)
-        init_max = w.abs().max(axis=-1, keepdims=True)
-
-        # try a range of clips and pick the one with the smallest loss
-        for i in range(int(max_shrink * n_grid)):
-            p = 1 - i / n_grid
-            w_max = p * init_max
-            w_m = mx.clip(w, -w_max, w_max).reshape(w_shape)
-
-            w_q = quantize_func(w_m)
-
-            w_q = w_q.reshape(*w_q.shape[:-1], -1, group_size)
-            out_q = mx.einsum("bdg,odg->bod", x, w_q)
-
-            # Take the mean across the input batch
-            loss = mse(out, out_q).sum(axis=0)
-            if group is not None:
-                loss = mx.distributed.all_sum(loss) / group.size()
-            loss /= out.shape[0]
-            best_indices = loss < best_error
-            best_error = mx.where(best_indices, loss, best_error)
-            best_w_max = mx.where(best_indices[..., mx.newaxis], w_max, best_w_max)
-            mx.eval(best_w_max, best_error)
-
-        w_max_all.append(best_w_max)
-
-    best_w_max = mx.concatenate(w_max_all, axis=0)
-
-    w_r = w_all.reshape(*w_all.shape[:-1], -1, group_size)
-    best_w = mx.clip(w_r, -best_w_max, best_w_max)
-    best_w = best_w.reshape(w_init_shape)
-
-    mx.eval(best_w)
-    return best_w
-
-
-def clip_block(
-    block: nn.Module,
-    no_clip_keys: list[str],
-    quantize_func: Callable,
-    group_size: int,
-    n_grid: int = 20,
-):
-    def apply_clip(path, module):
-        if isinstance(module, (nn.Linear, SwitchLinear)) and all(
-            k not in path for k in no_clip_keys
-        ):
-            best_weight = search_best_clip(
-                module,
-                quantize_func=quantize_func,
-                group_size=group_size,
-                n_grid=n_grid,
-            )
-            module.weight = best_weight
-
-    tree_map_with_path(apply_clip, block.leaf_modules(), is_leaf=nn.Module.is_module)
-
-
-def awq_quantize(
-    model,
-    inputs: mx.array,
-    awq_config: AWQConfig,
-    group_size: int = 64,
-    bits: int = 3,
-    embed_group_size: int = 32,
-    embed_bits: int = 4,
-    n_grid: int = 20,
-):
-    if awq_config.lm_key is not None:
-        model = model[awq_config.lm_key]
-
-    group = mx.distributed.init()
-
-    def quantize_func(w):
-        wq = mx.quantize(w, bits=bits, group_size=group_size)
-        return mx.dequantize(*wq, bits=bits, group_size=group_size)
-
-    mask = create_attention_mask(inputs)
-
-    embed_key = awq_config.embed
-    model.model[embed_key] = model.model[embed_key].to_quantized(
-        group_size=embed_group_size, bits=embed_bits
-    )
-    inputs = model.model[embed_key](inputs)
-
-    def capture(module):
-        if not isinstance(module, (nn.Linear, SwitchLinear)):
-            return module
-
-        class Catcher(nn.Module):
-            def __call__(self, x: mx.array, *args, **kwargs):
-                # Store the input features on the original modules.
-                if hasattr(module, "input_feat"):
-                    module.input_feat = mx.concatenate([module.input_feat, x], axis=0)
-                else:
-                    module.input_feat = x
-
-                # Also store the MOE indices if applicabale
-                if isinstance(module, SwitchLinear):
-                    indices = args[0]
-                    if hasattr(module, "indices"):
-                        module.indices = mx.concatenate(
-                            [module.indices, indices], axis=0
-                        )
-                    else:
-                        module.indices = indices
-
-                return module(x, *args, **kwargs)
-
-        return Catcher()
-
-    for e, block in enumerate(tqdm(model.layers)):
-        # Capture the input features for each of the layers in the transformer block
-        orig_leaves = block.leaf_modules()
-        capture_leaves = tree_map(capture, orig_leaves, is_leaf=nn.Module.is_module)
-        block.update_modules(capture_leaves)
-        outputs = run_layer(block, inputs, mask=mask)
-        block.update_modules(orig_leaves)
-        del capture_leaves
-
-        # Quantize the block without AWQ to obtain a reference loss
-        nn.quantize(block, group_size=group_size, bits=bits)
-        outputs_q = run_layer(block, inputs, mask=mask)
-        before_loss = mse(outputs, outputs_q).sum()
-        if group is not None:
-            before_loss = mx.distributed.all_sum(before_loss) / group.size()
-        before_loss /= outputs.size
-        block.update_modules(orig_leaves)
-        orig_params = block.parameters()
-
-        scale_block(
-            block=block,
-            configs=awq_config.scale_configs,
-            quantize_func=quantize_func,
-            n_grid=n_grid,
-            layer_kwargs={"mask": mask},
-        )
-
-        clip_block(
-            block=block,
-            no_clip_keys=awq_config.no_clip,
-            quantize_func=quantize_func,
-            group_size=group_size,
-            n_grid=n_grid,
-        )
-
-        # Quantize the scaled and clipped block
-        nn.quantize(block, group_size=group_size, bits=bits)
-        outputs_q = run_layer(block, inputs, mask=mask)
-        after_loss = mse(outputs, outputs_q).sum()
-        if group is not None:
-            after_loss = mx.distributed.all_sum(after_loss) / group.size()
-        after_loss /= outputs.size
-        tqdm.write(f"Loss reduction: {after_loss / before_loss}")
-        if after_loss > before_loss:
-            # Reload original weights and quantize
-            block.update_modules(orig_leaves)
-            block.update(orig_params)
-            nn.quantize(block, group_size=group_size, bits=bits)
-            tqdm.write("Loss is not reduced, falling back to original weights.")
-
-        inputs = outputs
-
-        mx.eval(block)
-        mx.clear_cache()
-
-    if (lm_head := awq_config.lm_head) in model:
-        model[lm_head] = model[lm_head].to_quantized(
-            group_size=embed_group_size, bits=embed_bits
-        )
-
-
-def load_dataset(tokenizer, num_samples: int, sequence_length: int) -> mx.array:
-    save_dir = Path.home() / ".cache/mlx-lm/calibration_v5.txt"
-    if not save_dir.exists():
-        save_dir.parent.mkdir(parents=True, exist_ok=True)
-        url = "https://gist.githubusercontent.com/tristandruyen/9e207a95c7d75ddf37525d353e00659c/raw/571fda718462de863e5a0171078c175420c7649a/calibration_data_v5_rc.txt"
-        request.urlretrieve(url, save_dir)
-    with open(save_dir) as fid:
-        texts = fid.read()
-    tokens = tokenizer.encode(texts, return_tensors="mlx")[0]
-
-    # select random non-overlapping chunks
-    tokens = tokens[: (tokens.size // sequence_length) * sequence_length]
-    tokens = tokens.reshape(-1, sequence_length)
-    segments = mx.random.permutation(tokens.shape[0])[:num_samples]
-    return tokens[segments]
-
-
-def update_config(
-    model: nn.Module,
-    config: Dict[str, Any],
-):
-    # dummy
-    config["quantization"] = {"group_size": 64, "bits": 4}
-
-    def update_config(path, module):
-        if hasattr(module, "bits"):
-            config["quantization"][path] = {
-                "group_size": module.group_size,
-                "bits": module.bits,
-            }
-        else:
-            config["quantization"][path] = False
-
-    tree_map_with_path(update_config, model.leaf_modules(), is_leaf=nn.Module.is_module)
-    return config
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--model", "-m", default="mlx-community/Qwen2.5-7B-Instruct-bf16"
-    )
-    parser.add_argument("--mlx-path", default="mlx_model")
-    parser.add_argument("--bits", type=int, default=4)
-    parser.add_argument("--group-size", type=int, default=64)
-    parser.add_argument("--embed-bits", type=int, default=4)
-    parser.add_argument("--embed-group-size", type=int, default=32)
-    parser.add_argument("--num-samples", type=int, default=128)
-    parser.add_argument("--sequence-length", type=int, default=512)
-    parser.add_argument("--n-grid", type=int, default=20)
-    parser.add_argument("--seed", type=int, default=123)
-    args = parser.parse_args()
-
-    group = mx.distributed.init()
-
-    num_samples = args.num_samples
-    if group is not None and num_samples % group.size() > 0:
-        num_samples += group.size() - num_samples % group.size()
-
-    mx.random.seed(args.seed)
-
-    model_path = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
-
-    model_type = config["model_type"]
-    if (awq_config := AWQ_MODEL_CONFIGS.get(model_type, None)) is None:
-        raise NotImplementedError(f"AWQ support for {model_type} models NYI.")
-
-    calibration_data = load_dataset(tokenizer, args.num_samples, args.sequence_length)
-
-    calibration_data = dist_split(calibration_data, group)
-
-    awq_quantize(
-        model,
-        calibration_data,
-        awq_config,
-        bits=args.bits,
-        group_size=args.group_size,
-        embed_bits=args.embed_bits,
-        embed_group_size=args.embed_group_size,
-        n_grid=args.n_grid,
-    )
-
-    config = update_config(model, config)
-    weights = dict(tree_flatten(model.parameters()))
-    save(
-        args.mlx_path,
-        model_path,
-        weights,
-        tokenizer,
-        config,
-        hf_repo=args.model,
-    )

mlx_lm/cache_prompt.py

@@ -7,9 +7,8 @@ import time
 
 import mlx.core as mx
 
-from .generate import generate_step
 from .models.cache import make_prompt_cache, save_prompt_cache
-from .utils import load
+from .utils import generate_step, load
 
 DEFAULT_QUANTIZED_KV_START = 5000
 
@@ -148,19 +147,15 @@ def main():
         pass
 
     print()
-    print(f"Peak memory: {mx.get_peak_memory() / 1e9:.3f} GB")
+    print(f"Peak memory: {mx.metal.get_peak_memory() / 1e9:.3f} GB")
 
     print("Saving...")
     metadata = {}
     metadata["model"] = args.model
-    metadata["chat_template"] = json.dumps(tokenizer.chat_template)
+    metadata["chat_template"] = tokenizer.chat_template
     metadata["tokenizer_config"] = json.dumps(tokenizer_config)
     save_prompt_cache(args.prompt_cache_file, cache, metadata)
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.cache_prompt...` directly is deprecated."
-        " Use `mlx_lm.cache_prompt...` or `python -m mlx_lm cache_prompt ...` instead."
-    )
     main()

mlx_lm/chat.py

@@ -5,16 +5,13 @@ import json
 
 import mlx.core as mx
 
-from .generate import stream_generate
 from .models.cache import make_prompt_cache
 from .sample_utils import make_sampler
-from .utils import load
+from .utils import load, stream_generate
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
-DEFAULT_XTC_PROBABILITY = 0.0
-DEFAULT_XTC_THRESHOLD = 0.0
-DEFAULT_SEED = None
+DEFAULT_SEED = 0
 DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 
@@ -39,24 +36,7 @@ def setup_arg_parser():
     parser.add_argument(
         "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
     )
-    parser.add_argument(
-        "--xtc-probability",
-        type=float,
-        default=DEFAULT_XTC_PROBABILITY,
-        help="Probability of XTC sampling to happen each next token",
-    )
-    parser.add_argument(
-        "--xtc-threshold",
-        type=float,
-        default=0.0,
-        help="Thresold the probs of each next token candidate to be sampled by XTC",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=DEFAULT_SEED,
-        help="PRNG seed",
-    )
+    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
     parser.add_argument(
         "--max-kv-size",
         type=int,
@@ -77,8 +57,7 @@ def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
-    if args.seed is not None:
-        mx.random.seed(args.seed)
+    mx.random.seed(args.seed)
 
     model, tokenizer = load(
         args.model,
@@ -86,25 +65,12 @@ def main():
         tokenizer_config={"trust_remote_code": True},
     )
 
-    def print_help():
-        print("The command list:")
-        print("- 'q' to exit")
-        print("- 'r' to reset the chat")
-        print("- 'h' to display these commands")
-
-    print(f"[INFO] Starting chat session with {args.model}.")
-    print_help()
+    print(f"[INFO] Starting chat session with {args.model}. To exit, enter 'q'.")
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
         query = input(">> ")
         if query == "q":
             break
-        if query == "r":
-            prompt_cache = make_prompt_cache(model, args.max_kv_size)
-            continue
-        if query == "h":
-            print_help()
-            continue
         messages = [{"role": "user", "content": query}]
         prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
         for response in stream_generate(
@@ -112,15 +78,7 @@ def main():
             tokenizer,
             prompt,
             max_tokens=args.max_tokens,
-            sampler=make_sampler(
-                args.temp,
-                args.top_p,
-                xtc_threshold=args.xtc_threshold,
-                xtc_probability=args.xtc_probability,
-                xtc_special_tokens=(
-                    tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
-                ),
-            ),
+            sampler=make_sampler(args.temp, args.top_p),
             prompt_cache=prompt_cache,
         ):
             print(response.text, flush=True, end="")
@@ -128,8 +86,4 @@ def main():
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.chat...` directly is deprecated."
-        " Use `mlx_lm.chat...` or `python -m mlx_lm chat ...` instead."
-    )
     main()

mlx_lm/convert.py

@@ -1,165 +1,8 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-from pathlib import Path
-from typing import Callable, Optional, Union
 
-import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_flatten
-
-from .utils import (
-    dequantize_model,
-    fetch_from_hub,
-    get_model_path,
-    quantize_model,
-    save,
-    upload_to_hub,
-)
-
-
-def mixed_quant_predicate_builder(
-    recipe: str, model: nn.Module
-) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
-
-    high_bits = 6
-    group_size = 64
-
-    if recipe == "mixed_2_6":
-        low_bits = 2
-    elif recipe == "mixed_3_4":
-        low_bits = 3
-        high_bits = 4
-    elif recipe == "mixed_3_6":
-        low_bits = 3
-    elif recipe == "mixed_4_6":
-        low_bits = 4
-    else:
-        raise ValueError("Invalid quant recipe {recipe}")
-
-    down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
-    if len(down_keys) == 0:
-        raise ValueError("Model does not have expected keys for mixed quant.")
-
-    # Look for the layer index location in the path:
-    for layer_location, k in enumerate(down_keys[0].split(".")):
-        if k.isdigit():
-            break
-    num_layers = len(model.layers)
-
-    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
-        By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
-        """
-
-        if not hasattr(module, "to_quantized"):
-            return False
-
-        index = (
-            int(path.split(".")[layer_location])
-            if len(path.split(".")) > layer_location
-            else 0
-        )
-        use_more_bits = (
-            index < num_layers // 8
-            or index >= 7 * num_layers // 8
-            or (index - num_layers // 8) % 3 == 2
-        )
-        if "v_proj" in path and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
-        if "down_proj" in path and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
-        if "lm_head" in path:
-            return {"group_size": group_size, "bits": high_bits}
-
-        return {"group_size": group_size, "bits": low_bits}
-
-    return mixed_quant_predicate
-
-
-QUANT_RECIPES = ["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"]
-
-MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]
-
-
-def convert(
-    hf_path: str,
-    mlx_path: str = "mlx_model",
-    quantize: bool = False,
-    q_group_size: int = 64,
-    q_bits: int = 4,
-    dtype: Optional[str] = None,
-    upload_repo: str = None,
-    revision: Optional[str] = None,
-    dequantize: bool = False,
-    quant_predicate: Optional[
-        Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
-    ] = None,
-):
-    # Check the save path is empty
-    if isinstance(mlx_path, str):
-        mlx_path = Path(mlx_path)
-
-    if mlx_path.exists():
-        raise ValueError(
-            f"Cannot save to the path {mlx_path} as it already exists."
-            " Please delete the file/directory or specify a new path to save to."
-        )
-
-    print("[INFO] Loading")
-    model_path = get_model_path(hf_path, revision=revision)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
-
-    if isinstance(quant_predicate, str):
-        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
-
-    if dtype is None:
-        dtype = config.get("torch_dtype", None)
-    weights = dict(tree_flatten(model.parameters()))
-    if dtype in MODEL_CONVERSION_DTYPES:
-        print("[INFO] Using dtype:", dtype)
-        dtype = getattr(mx, dtype)
-
-        if hasattr(model, "cast_predicate"):
-            cast_predicate = model.cast_predicate()
-        else:
-            cast_predicate = lambda _: True
-        weights = {
-            k: v.astype(dtype) if cast_predicate(k) else v for k, v in weights.items()
-        }
-
-    if quantize and dequantize:
-        raise ValueError("Choose either quantize or dequantize, not both.")
-
-    if quantize:
-        print("[INFO] Quantizing")
-        model.load_weights(list(weights.items()))
-        weights, config = quantize_model(
-            model, config, q_group_size, q_bits, quant_predicate=quant_predicate
-        )
-
-    if dequantize:
-        print("[INFO] Dequantizing")
-        model = dequantize_model(model)
-        weights = dict(tree_flatten(model.parameters()))
-
-    del model
-    save(
-        mlx_path,
-        model_path,
-        weights,
-        tokenizer,
-        config,
-        hf_repo=hf_path,
-    )
-
-    if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo)
+from .utils import convert
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -186,19 +29,12 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
-    parser.add_argument(
-        "--quant-predicate",
-        help=f"Mixed-bit quantization recipe.",
-        choices=QUANT_RECIPES,
-        type=str,
-        required=False,
-    )
     parser.add_argument(
         "--dtype",
-        help="Type to save the non-quantized parameters. Defaults to config.json's `torch_dtype` or the current model weights dtype.",
+        help="Type to save the non-quantized parameters.",
         type=str,
-        choices=MODEL_CONVERSION_DTYPES,
-        default=None,
+        choices=["float16", "bfloat16", "float32"],
+        default="float16",
     )
     parser.add_argument(
         "--upload-repo",
@@ -223,8 +59,4 @@ def main():
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.convert ...` directly is deprecated."
-        " Use `mlx_lm.convert ...` or `python -m mlx_lm convert ...` instead."
-    )
     main()

mlx_lm/dwq.py

@@ -1,254 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-import copy
-import glob
-import shutil
-import time
-import types
-from pathlib import Path
-
-import mlx.core as mx
-import mlx.nn as nn
-import mlx.optimizers as optimizers
-import numpy as np
-from mlx.utils import tree_flatten, tree_map
-
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-from mlx_lm.tuner.datasets import load_dataset
-from mlx_lm.tuner.trainer import iterate_batches
-from mlx_lm.tuner.utils import print_trainable_parameters
-from mlx_lm.utils import (
-    create_model_card,
-    fetch_from_hub,
-    get_model_path,
-    quantize_model,
-    save_config,
-    save_weights,
-)
-
-
-class Catcher(nn.Module):
-    def __init__(self, module):
-        super().__init__()
-        self.module = module
-
-    def __call__(self, *args, **kwargs):
-        self.outputs = self.module(*args, **kwargs)
-        return self.outputs
-
-
-def dwq_quantize(
-    model,
-    q_model,
-    opt,
-    data,
-    batch_size: int = 2,
-    max_seq_length: int = 2048,
-    temperature: float = 1.0,
-    activation_layer_step: float = 0.25,
-    activation_loss_weight: float = 1e-1,
-    dtype: mx.Dtype = mx.bfloat16,
-):
-    group = mx.distributed.init()
-    world_size = group.size()
-    rank = group.rank()
-
-    def unfreeze(_, m):
-        if hasattr(m, "bits") and hasattr(m, "group_size"):
-            m.unfreeze(keys=["scales", "biases"], recurse=False)
-
-    q_model.apply_to_modules(unfreeze)
-    print_trainable_parameters(q_model)
-
-    layer_id_step = int(activation_layer_step * len(model.layers))
-    layer_ids = list(range(len(model.layers)))[layer_id_step::layer_id_step]
-
-    for lid in layer_ids:
-        model.layers[lid] = Catcher(model.layers[lid])
-        q_model.layers[lid] = Catcher(q_model.layers[lid])
-
-    def log_norm(x):
-        if temperature != 1.0:
-            x = x * (1 / temperature)
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
-
-    def forward(model, inputs):
-        logprobs = log_norm(model(inputs).astype(mx.float32))
-        extra_targets = [
-            model.layers[lid].outputs.astype(mx.float32) for lid in layer_ids
-        ]
-        for lid in layer_ids:
-            model.layers[lid].outputs = None
-        return logprobs, extra_targets
-
-    def loss_fn(params, x, targets, extra_targets, lengths):
-        q_model.update(tree_map(lambda x: x.astype(dtype), params))
-        logprobs, q_extra_targets = forward(q_model, x)
-        losses = nn.losses.kl_div_loss(logprobs, targets, reduction="none")
-        mask = mx.arange(targets.shape[1]) < lengths[:, 1:]
-        ntoks = mask.sum()
-        kl_loss = (mask * losses).sum() / ntoks
-        act_loss = mx.stack(
-            [
-                (mask * (qe - e).abs().mean(axis=-1)).sum() / ntoks
-                for qe, e in zip(q_extra_targets, extra_targets)
-            ]
-        )
-        loss = kl_loss + activation_loss_weight * act_loss.mean()
-        return loss, ntoks
-
-    def step(inputs, targets, extra_targets, lengths, params):
-        (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
-
-    # Accumulate learned weights in higher precision
-    params = tree_map(
-        lambda x: x.astype(mx.float32),
-        q_model.trainable_parameters(),
-    )
-
-    avg_loss = None
-    tokens = 0
-    tic = time.time()
-    for it, (batch, lengths) in enumerate(
-        iterate_batches(data, batch_size, max_seq_length)
-    ):
-        targets, extra_targets = forward(model, batch)
-        mx.eval(targets, extra_targets)
-        loss, ntoks, params = step(batch, targets, extra_targets, lengths, params)
-        mx.eval(loss, params)
-        loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
-        ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
-        tokens += ntoks
-        toks_per_sec = tokens / (time.time() - tic)
-        avg_loss = 0.95 * (avg_loss or loss) + 0.05 * loss
-        if rank == 0:
-            print(
-                f"{it=}, {loss=:.3f}, {avg_loss=:.4f}, {tokens=}, {toks_per_sec=:.3f}",
-                flush=True,
-            )
-    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 save_model(
-    model: nn.Module,
-    tokenizer: TokenizerWrapper,
-    config,
-    model_path: Path,
-    mlx_path: str,
-    hf_path: str,
-):
-    weights = dict(tree_flatten(model.parameters()))
-
-    mlx_path = Path(mlx_path)
-    save_weights(mlx_path, weights, donate_weights=True)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, mlx_path)
-
-    tokenizer.save_pretrained(mlx_path)
-
-    save_config(config, config_path=mlx_path / "config.json")
-    create_model_card(mlx_path, hf_path)
-
-
-def load_data(tokenizer, data_path: str, num_samples: int):
-    args = types.SimpleNamespace(
-        hf_dataset={
-            "path": data_path,
-            "train_split": f"train",
-            "valid_split": "train[:1]",
-        },
-        train=True,
-        test=False,
-    )
-    dataset = load_dataset(args, tokenizer)[0]
-    perm = np.random.permutation(len(dataset))[:num_samples].tolist()
-    return [dataset.process(dataset[i]) for i in perm]
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", "-m", default="Qwen/Qwen3-4B")
-    parser.add_argument("--quantized-model", default=None)
-    parser.add_argument(
-        "--mlx-path", default="mlx_model", help="Path to save the quantized model."
-    )
-    parser.add_argument(
-        "--bits",
-        type=int,
-        default=4,
-        help="Bits per weight for quantization.",
-    )
-    parser.add_argument(
-        "--group-size", type=int, default=64, help="Group size for quantization."
-    )
-    parser.add_argument(
-        "--num-samples",
-        type=int,
-        default=1024,
-        help="Number of samples to use for training.",
-    )
-    parser.add_argument("--max-seq-length", type=int, default=2048)
-    parser.add_argument("--seed", type=int, default=123)
-    parser.add_argument("--learning-rate", type=float, default=1e-6)
-    parser.add_argument("--batch-size", type=int, default=4)
-    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(
-        "--temperature",
-        type=float,
-        default=1.0,
-        help="Temperature scaling for the loss.",
-    )
-    args = parser.parse_args()
-
-    group = mx.distributed.init()
-
-    num_samples = args.num_samples
-    if num_samples % group.size() > 0:
-        num_samples += group.size() - num_samples % group.size()
-
-    np.random.seed(args.seed)
-    mx.random.seed(args.seed)
-
-    model_path = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
-
-    calibration_data = load_data(tokenizer, args.data_path, args.num_samples)
-
-    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)
-    else:
-        q_model = copy.deepcopy(model)
-        _, config = quantize_model(
-            q_model,
-            config,
-            q_group_size=args.group_size,
-            q_bits=args.bits,
-        )
-
-    opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
-    dwq_quantize(
-        model,
-        q_model,
-        opt,
-        calibration_data,
-        batch_size=args.batch_size,
-        max_seq_length=args.max_seq_length,
-        temperature=args.temperature,
-    )
-    save_model(q_model, tokenizer, config, model_path, args.mlx_path, args.model)

mlx_lm/evaluate.py

@@ -5,14 +5,12 @@ Adapted from a PyTorch implementation by David Grangier
 """
 
 import argparse
-import collections
-import copy
 import json
 import logging
 import os
 from importlib.metadata import version
 from pathlib import Path
-from typing import Any, Optional
+from typing import Optional, Union
 
 import lm_eval
 import mlx.core as mx
@@ -22,10 +20,19 @@ from lm_eval.api.model import LM
 from lm_eval.api.registry import register_model
 from tqdm import tqdm
 
-from .generate import stream_generate
-from .models.base import create_causal_mask
 from .models.cache import make_prompt_cache
-from .utils import common_prefix_len, load
+from .utils import load, stream_generate
+
+PAD = 0
+
+
+def _len_longest_common_prefix(a, b):
+    l = 0
+    for item_a, item_b in zip(a, b):
+        if item_a != item_b:
+            break
+        l += 1
+    return l
 
 
 def _rstrip_until(s, untils):
@@ -36,85 +43,75 @@ def _rstrip_until(s, untils):
     return s[: min(f)]
 
 
-def _pad_inputs(inputs):
-    lengths = np.array([len(x) for x in inputs])
-    maxlen = lengths.max()
-    padded = np.stack(
-        [np.pad(x, (0, maxlen - len(x))) for x in inputs],
+def _pad_inputs(
+    inputs,
+    maxlen,
+    genlen=0,
+    pad_left=False,
+    pad_multiple=32,
+    truncate=False,
+):
+    # pad the prompts to the left with at least genlen tokens.
+    actual_maxlen = max(len(p) for p in inputs) + genlen
+    if actual_maxlen > maxlen:
+        if not truncate:
+            raise ValueError("Inputs are too long.")
+        else:  # drop begining
+            actual_maxlen = maxlen
+            inputs = [p[max(0, len(p) - maxlen) :] for p in inputs]
+    if pad_multiple > 0:
+        maxlen = (actual_maxlen + pad_multiple - 1) // pad_multiple
+        maxlen *= pad_multiple
+    assert PAD == 0
+    lr = np.array((1, 0) if pad_left else (0, 1))
+    return np.stack(
+        [np.pad(np.array(x, np.int32), lr * (maxlen - len(x))) for x in inputs],
         axis=0,
     )
-    return mx.array(padded), mx.array(lengths)
-
-
-def chat_template_fn(**extra_kwargs):
-    def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
-        return self.tokenizer.apply_chat_template(
-            chat_history,
-            tokenize=False,
-            add_generation_prompt=add_generation_prompt,
-            continue_final_message=not add_generation_prompt,
-            **extra_kwargs,
-        )
-
-    return apply_chat_template
 
 
 @register_model("mlxlm")
 class MLXLM(LM):
-
-    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
-    apply_chat_template = chat_template_fn()
-
     def __init__(
         self,
         path_or_hf_repo: str,
+        batch_size: int = 16,
         max_tokens: Optional[int] = None,
         use_chat_template: Optional[bool] = None,
     ) -> None:
         super().__init__()
+        self._batch_size = batch_size
         self._model, self.tokenizer = load(path_or_hf_repo)
         self._max_tokens = max_tokens or self.tokenizer.model_max_length
-        self._batch_size = 8
-        self.use_chat_template = use_chat_template
-        if use_chat_template is None:
-            self.use_chat_template = self.tokenizer.chat_template is not None
+        self.use_chat_template = use_chat_template or (
+            self.tokenizer.chat_template is not None
+        )
 
-    def _process_prompt(self, prompt, step_size: int = 2048):
-        prompt = mx.array(prompt)[None]
-        cache = make_prompt_cache(self._model)
-        for i in range(0, prompt.shape[1], step_size):
-            logits = self._model(prompt[:, i : i + step_size], cache=cache)
-            mx.eval([c.state for c in cache])
-            mx.clear_cache()
-        logprobs = nn.log_softmax(logits[:, -1, :].astype(mx.float32))
-        return logprobs, cache
-
-    def _score_fn(self, inputs, cache: Optional[Any] = None, step_size: int = 2048):
-        inputs, lengths = _pad_inputs(inputs)
+    def _score_fn(self, inputs, tokenize=True, step_size=32):
+        if tokenize:
+            inputs = self._tokenize(inputs)
+        inputs = _pad_inputs(inputs, self._max_tokens, truncate=False)
+        inputs = mx.array(inputs)
         inputs, targets = inputs[..., :-1], inputs[..., 1:]
 
-        cache = cache or make_prompt_cache(self._model)
-        lengths += cache[0].offset
+        cache = make_prompt_cache(self._model)
+
+        mask = targets != PAD
 
         scores, is_greedy = [], []
         for i in range(0, inputs.shape[1], step_size):
-            inp = inputs[:, i : i + step_size]
-            T = inp.shape[1]
+            logits = self._model(inputs[:, i : i + step_size], cache=cache)
 
-            offset = cache[0].offset
-            mask = create_causal_mask(T, offset, lengths=lengths)
-
-            logits = self._model(inp, cache=cache, mask=mask)
             log_probs = nn.log_softmax(logits.astype(mx.float32))
-
             score = mx.take_along_axis(
                 log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
             )[..., 0]
-            ig = targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
-            ig = mx.where(mx.arange(T) + offset < lengths[:, None], ig, False)
+            ig = mask[:, i : i + step_size] * (
+                targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
+            )
 
             mx.eval(score, ig)
-            mx.clear_cache()
+            mx.metal.clear_cache()
 
             is_greedy.append(ig)
             scores.append(score)
@@ -122,7 +119,38 @@ class MLXLM(LM):
         scores = mx.concatenate(scores, axis=1)
         is_greedy = mx.concatenate(is_greedy, axis=1)
 
-        return scores, lengths, is_greedy
+        return scores, mask.sum(axis=-1), is_greedy
+
+    def _loglikelihood(self, texts, score_spans=None, tokenize=True):
+        # sort by length to get batches with little padding.
+        sorted_indices = sorted(range(len(texts)), key=lambda i: -len(texts[i]))
+        sorted_inputs = [texts[sorted_indices[i]] for i in range(len(texts))]
+        sorted_spans = None
+        if score_spans is not None:
+            sorted_spans = [score_spans[sorted_indices[i]] for i in range(len(texts))]
+
+        results = []
+        for i in tqdm(range(0, len(sorted_inputs), self._batch_size)):
+            batch = sorted_inputs[i : i + self._batch_size]
+            scores, length, is_greedy = self._score_fn(batch, tokenize=tokenize)
+            for j in range(len(batch)):
+                if sorted_spans is None:  # full sequence score
+                    mask = mx.arange(scores[j].shape[-1]) < length
+                    score = (scores[j].astype(mx.float32) * mask).sum(axis=-1)
+                    ig = (is_greedy[j].astype(mx.int32) * mask).sum(axis=-1)
+                else:  # subsequence score
+                    start, end = sorted_spans[i + j]
+                    score = scores[j][start:end].astype(mx.float32).sum()
+                    ig = is_greedy[j][start:end].astype(mx.int32).sum()
+                    length = end - start
+
+                results.append((score.item(), ig.item(), length))
+
+        # reorder the outputs
+        inv_sort = np.argsort(sorted_indices)
+        results = [results[inv_sort[i]] for i in range(len(results))]
+
+        return results
 
     def _tokenize(self, texts):
         return [
@@ -154,65 +182,39 @@ class MLXLM(LM):
         """
         logging.info("Estimating loglikelihood for %d pairs." % len(requests))
 
-        group = mx.distributed.init()
+        # tokenize prefix and prefix + completion for all requests.
+        tokenized = self._tokenize(
+            [t for r in requests for t in [r.args[0], r.args[0] + r.args[1]]]
+        )
 
-        # Group by common prefix
-        group_reqs = collections.defaultdict(list)
-        for idx, req in enumerate(requests):
-            group_reqs[req.args[0]].append((idx, req.args[1]))
-        questions = list(group_reqs.keys())
-        responses = []
-        indices = []
-        for v in group_reqs.values():
-            idx, resp = zip(*v)
-            indices.extend(idx)
-            responses.append(resp)
-
-        # split data accross ranks
-        questions = questions[group.rank() :: group.size()]
-        responses = responses[group.rank() :: group.size()]
+        # max length (prefix + completion) and longest common prefix per question.
+        length_stats = {}
+        for prefix, completed in zip(tokenized[0::2], tokenized[1::2]):
+            max_completed_l, min_prefix_l = length_stats.get(prefix, (0, 1e8))
+            length_stats[prefix] = (
+                max(max_completed_l, len(completed)),
+                min(min_prefix_l, _len_longest_common_prefix(prefix, completed)),
+            )
 
+        # truncate requests for completed sequences longer than model context.
+        shortened = []
+        completion_spans = []
         long_completions = 0
-        scores, is_greedy = [], []
-        for q, rs in tqdm(zip(questions, responses), total=len(questions)):
-            prefix = self._tokenize([q])[0]
-            full_sequences = self._tokenize([q + r for r in rs])
-            max_completed_l = max(len(s) for s in full_sequences)
-
+        for prefix, completed in zip(tokenized[0::2], tokenized[1::2]):
+            max_completed_l, prefix_l = length_stats[prefix]
             # compute truncation length
             truncation = max(0, max_completed_l - self._max_tokens - 1)
-            orig_prefix_l = len(prefix)
-            prefix_l = max(len(prefix) - truncation, 0)
-            prefix = prefix[len(prefix) - prefix_l :]
-
-            # If the entire prompt got truncated ignore the question
-            if prefix_l == 0:
+            prefix_l = prefix_l - truncation
+            if prefix_l <= 0:
+                # completion too long, prefix is eliminated for some requests.
                 long_completions += 1
-                all_scores.extend([-float("inf")] * len(rs))
-                all_is_greedy.extend([False] * len(rs))
-                continue
-
-            # model scoring, returns num_requests x (logp, is_greedy, length).
-            logprobs, cache = self._process_prompt(prefix)
-            max_idx = mx.argmax(logprobs).item()
-
-            for s in full_sequences:
-                inputs = s[len(prefix) :]
-                # The logprobs from the last token of the prompt are
-                # for the first input token
-                scores.append(logprobs[0, inputs[0]].item())
-                is_greedy.append((inputs[0] == max_idx))
-
-                if len(inputs) == 1:
-                    continue
-                score, _, ig = self._score_fn(
-                    mx.array(inputs)[None, :], cache=copy.deepcopy(cache)
-                )
-                scores[-1] += mx.sum(score).item()
-                is_greedy[-1] &= mx.all(ig).item()
-
-        scores = mx.array(scores)
-        is_greedy = mx.array(is_greedy)
+                truncation = max(0, len(completed) - self._max_tokens - 1)
+                prefix_l = 1
+            # truncate the completed sequence
+            completed = completed[truncation:]
+            shortened.append(completed)
+            # scores do not include initial bos, substract 1 to span bounds
+            completion_spans.append((prefix_l - 1, len(completed) - 1))
 
         if long_completions > 0:
             logging.info(
@@ -220,23 +222,16 @@ class MLXLM(LM):
                 + "completion longer than context."
             )
 
-        num_results = len(requests)
+        # model scoring, returns num_requests x (logp, is_greedy, length).
+        results = self._loglikelihood(
+            shortened,
+            score_spans=completion_spans,
+            tokenize=False,
+        )
+        return [(r[0], r[1] == r[2]) for r in results]
 
-        # 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))
-        scores = scores[:num_results][inv_sort]
-        is_greedy = is_greedy[:num_results][inv_sort]
-        return list(zip(scores.tolist(), is_greedy.tolist()))
+    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
+    apply_chat_template = lm_eval.models.huggingface.HFLM.apply_chat_template
 
     def loglikelihood_rolling(self, requests) -> list[float]:
         """Compute full log-likelihood of a string, with no truncation, for perplexity computation
@@ -273,15 +268,8 @@ class MLXLM(LM):
         logging.info(
             "Estimating loglikelihood rolling for %d sequences." % len(requests)
         )
-        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]
-            scores, lengths, _ = self._score_fn(batch)
-            mask = mx.arange(scores.shape[-1]) < lengths[:, None]
-            all_scores.extend((mask * scores).sum(axis=-1).tolist())
-
-        return all_scores
+        inputs = [req.args[0] for req in requests]
+        return [t[0] for t in self._loglikelihood(inputs)]
 
     def generate_until(self, requests) -> list[str]:
         """Generate greedily until a stopping sequence
@@ -301,15 +289,15 @@ class MLXLM(LM):
         contexts, options = zip(*[req.args for req in requests])
         # contrary to the doc the second element of the tuple contains
         # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
+        keys = list(options[0].keys())
+        assert "until" in keys
+        untils = [x["until"] for x in options]
         completions = []
 
-        for context, opt in tqdm(zip(contexts, options), total=len(contexts)):
-            until = opt["until"]
-            context = self.tokenizer.encode(
-                context, add_special_tokens=not self.use_chat_template
-            )
+        for context, until in tqdm(zip(contexts, untils), total=len(contexts)):
+            context = self._tokenize(context)
             max_tokens = min(
-                opt.get("max_gen_tokens", self._max_tokens),
+                self._max_tokens,
                 self.tokenizer.model_max_length - len(context),
             )
             text = ""
@@ -336,7 +324,7 @@ def main():
         "--output-dir", default=".", help="Output directory for result files."
     )
     parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
-    parser.add_argument("--num-shots", type=int, default=None, help="Number of shots")
+    parser.add_argument("--num-shots", type=int, default=0, help="Number of shots")
     parser.add_argument(
         "--max-tokens",
         type=int,
@@ -344,9 +332,9 @@ def main():
     )
     parser.add_argument(
         "--limit",
-        default=None,
+        default=1.0,
         help="Limit the number of examples per task.",
-        type=int,
+        type=float,
     )
     parser.add_argument("--seed", type=int, default=123, help="Random seed.")
     parser.add_argument(
@@ -364,14 +352,6 @@ def main():
         "otherwise `False`.",
         default=None,
     )
-    parser.add_argument(
-        "--chat-template-args",
-        type=json.loads,
-        help="""A JSON formatted string of arguments for the tokenizer's "
-        "apply_chat_template, e.g. '{"enable_thinking":false}'""",
-        default="{}",
-    )
-
     args = parser.parse_args()
 
     output_dir = Path(args.output_dir)
@@ -384,11 +364,10 @@ def main():
 
     lm = MLXLM(
         args.model,
+        batch_size=args.batch_size,
         max_tokens=args.max_tokens,
         use_chat_template=args.apply_chat_template,
     )
-    MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
-
     results = lm_eval.simple_evaluate(
         model=lm,
         tasks=args.tasks,
@@ -402,14 +381,12 @@ def main():
         fewshot_random_seed=args.seed,
     )
 
-    file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
-    if args.num_shots is not None:
-        file_keys += [f"{args.num_shots:02d}"]
-    file_keys += args.tasks
-    filename = "_".join(file_keys)
-    if mx.distributed.init().rank() == 0:
-        output_path = output_dir / filename
-        output_path.write_text(json.dumps(results["results"], indent=4))
-        print("Results:")
-        for result in results["results"].values():
-            print(json.dumps(result, indent=4))
+    model_name = args.model.replace("/", "_")
+    task_names = "_".join(args.tasks)
+    ver = version("lm_eval")
+    filename = f"eval_{model_name}_{task_names}_{args.num_shots:02d}_v_{ver}.json"
+    output_path = output_dir / filename
+    output_path.write_text(json.dumps(results["results"], indent=4))
+    print("Results:")
+    for result in results["results"].values():
+        print(json.dumps(result, indent=4))

mlx_lm/examples/chat.py

@@ -23,6 +23,7 @@ response = generate(
     tokenizer,
     prompt=prompt,
     verbose=True,
+    temp=0.0,
     prompt_cache=prompt_cache,
 )
 

mlx_lm/examples/lora_config.yaml

@@ -7,15 +7,6 @@ train: true
 # The fine-tuning method: "lora", "dora", or "full".
 fine_tune_type: lora
 
-# The Optimizer with its possible inputs
-optimizer: adamw
-# optimizer_config:
-#   adamw:
-#     betas: [0.9, 0.98]
-#     eps: 1e-6
-#     weight_decay: 0.05
-#     bias_correction: true
-
 # Directory with {train, valid, test}.jsonl files
 data: "/path/to/training/data"
 
@@ -81,7 +72,7 @@ lora_parameters:
 #  arguments: [1e-5, 1000, 1e-7] # passed to scheduler
 
 #hf_dataset:
-#  path: "billsum"
+#  name: "billsum"
 #  train_split: "train[:1000]"
 #  valid_split: "train[-100:]"
 #  prompt_feature: "text"

mlx_lm/examples/pipeline_generate.py

@@ -1,131 +0,0 @@
-# Copyright © 2024 Apple Inc.
-
-"""
-Run with:
-
-```
-mlx.launch \
- --hostfile /path/to/hosts.json \
- /path/to/pipeline_generate.py \
- --prompt "hello world"
-```
-
-Make sure you can run MLX over MPI on two hosts. For more information see the
-documentation:
-
-https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
-"""
-
-import argparse
-import json
-import resource
-from pathlib import Path
-
-import mlx.core as mx
-from huggingface_hub import snapshot_download
-from mlx.utils import tree_flatten
-
-from mlx_lm import load, stream_generate
-from mlx_lm.utils import load_model, load_tokenizer
-
-# Needed for 8 bit model
-resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
-
-
-def download(repo: str, allow_patterns: list[str]) -> Path:
-    return Path(
-        snapshot_download(
-            repo,
-            allow_patterns=allow_patterns,
-        )
-    )
-
-
-def shard_and_load(repo):
-    # Get model path with everything but weight safetensors
-    model_path = download(
-        args.model,
-        allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
-    )
-
-    # Lazy load and shard model to figure out
-    # which weights we need
-    model, _ = load_model(model_path, lazy=True, strict=False)
-
-    group = mx.distributed.init()
-    rank = group.rank()
-    model.model.pipeline(group)
-
-    # Figure out which files we need for the local shard
-    with open(model_path / "model.safetensors.index.json", "r") as fid:
-        weight_index = json.load(fid)["weight_map"]
-
-    local_files = set()
-    for k, _ in tree_flatten(model.parameters()):
-        local_files.add(weight_index[k])
-
-    # Download weights for local shard
-    download(args.model, allow_patterns=local_files)
-
-    # Load and shard the model, and load the weights
-    tokenizer = load_tokenizer(model_path)
-    model, _ = load_model(model_path, lazy=True, strict=False)
-    model.model.pipeline(group)
-    mx.eval(model.parameters())
-
-    # Synchronize processes before generation to avoid timeout if downloading
-    # model for the first time.
-    mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
-    return model, tokenizer
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
-    parser.add_argument(
-        "--model",
-        default="mlx-community/DeepSeek-R1-3bit",
-        help="HF repo or path to local model.",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        default="Write a quicksort in C++.",
-        help="Message to be processed by the model ('-' reads from stdin)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        "-m",
-        type=int,
-        default=256,
-        help="Maximum number of tokens to generate",
-    )
-    args = parser.parse_args()
-
-    group = mx.distributed.init()
-    rank = group.rank()
-
-    def rprint(*args, **kwargs):
-        if rank == 0:
-            print(*args, **kwargs)
-
-    model, tokenizer = shard_and_load(args.model)
-
-    messages = [{"role": "user", "content": args.prompt}]
-    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-
-    for response in stream_generate(
-        model, tokenizer, prompt, max_tokens=args.max_tokens
-    ):
-        rprint(response.text, end="", flush=True)
-
-    rprint()
-    rprint("=" * 10)
-    rprint(
-        f"Prompt: {response.prompt_tokens} tokens, "
-        f"{response.prompt_tps:.3f} tokens-per-sec"
-    )
-    rprint(
-        f"Generation: {response.generation_tokens} tokens, "
-        f"{response.generation_tps:.3f} tokens-per-sec"
-    )
-    rprint(f"Peak memory: {response.peak_memory:.3f} GB")

mlx_lm/examples/tool_use.py

@@ -1,73 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import json
-
-from mlx_lm import generate, load
-from mlx_lm.models.cache import make_prompt_cache
-
-# Specify the checkpoint
-checkpoint = "mlx-community/Qwen2.5-32B-Instruct-4bit"
-
-# Load the corresponding model and tokenizer
-model, tokenizer = load(path_or_hf_repo=checkpoint)
-
-
-# An example tool, make sure to include a docstring and type hints
-def multiply(a: float, b: float):
-    """
-    A function that multiplies two numbers
-
-    Args:
-        a: The first number to multiply
-        b: The second number to multiply
-    """
-    return a * b
-
-
-tools = {"multiply": multiply}
-
-# Specify the prompt and conversation history
-prompt = "Multiply 12234585 and 48838483920."
-messages = [{"role": "user", "content": prompt}]
-
-prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, tools=list(tools.values())
-)
-
-prompt_cache = make_prompt_cache(model)
-
-# Generate the initial tool call:
-response = generate(
-    model=model,
-    tokenizer=tokenizer,
-    prompt=prompt,
-    max_tokens=2048,
-    verbose=True,
-    prompt_cache=prompt_cache,
-)
-
-# Parse the tool call:
-# (Note, the tool call format is model specific)
-tool_open = "<tool_call>"
-tool_close = "</tool_call>"
-start_tool = response.find(tool_open) + len(tool_open)
-end_tool = response.find(tool_close)
-tool_call = json.loads(response[start_tool:end_tool].strip())
-tool_result = tools[tool_call["name"]](**tool_call["arguments"])
-
-# Put the tool result in the prompt
-messages = [{"role": "tool", "name": tool_call["name"], "content": tool_result}]
-prompt = tokenizer.apply_chat_template(
-    messages,
-    add_generation_prompt=True,
-)
-
-# Generate the final response:
-response = generate(
-    model=model,
-    tokenizer=tokenizer,
-    prompt=prompt,
-    max_tokens=2048,
-    verbose=True,
-    prompt_cache=prompt_cache,
-)

mlx_lm/fuse.py

@@ -1,4 +1,6 @@
 import argparse
+import glob
+import shutil
 from pathlib import Path
 
 from mlx.utils import tree_flatten, tree_unflatten
@@ -10,7 +12,8 @@ from .tuner.utils import dequantize, load_adapters
 from .utils import (
     fetch_from_hub,
     get_model_path,
-    save,
+    save_config,
+    save_weights,
     upload_to_hub,
 )
 
@@ -86,21 +89,23 @@ def main() -> None:
     if args.de_quantize:
         print("De-quantizing model")
         model = dequantize(model)
-        config.pop("quantization", None)
 
     weights = dict(tree_flatten(model.parameters()))
 
     save_path = Path(args.save_path)
-    hf_path = args.hf_path or (args.model if not Path(args.model).exists() else None)
-    save(
-        save_path,
-        model_path,
-        weights,
-        tokenizer,
-        config,
-        hf_repo=hf_path,
-        donate_weights=False,
-    )
+
+    save_weights(save_path, weights)
+
+    py_files = glob.glob(str(model_path / "*.py"))
+    for file in py_files:
+        shutil.copy(file, save_path)
+
+    tokenizer.save_pretrained(save_path)
+
+    if args.de_quantize:
+        config.pop("quantization", None)
+
+    save_config(config, config_path=save_path / "config.json")
 
     if args.export_gguf:
         model_type = config["model_type"]
@@ -111,6 +116,9 @@ def main() -> None:
         convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not None:
+        hf_path = args.hf_path or (
+            args.model if not Path(args.model).exists() else None
+        )
         if hf_path is None:
             raise ValueError(
                 "Must provide original Hugging Face repo to upload local model."
@@ -119,8 +127,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.fuse...` directly is deprecated."
-        " Use `mlx_lm.fuse...` or `python -m mlx_lm fuse ...` instead."
-    )
     main()

mlx_lm/generate.py

@@ -1,48 +1,22 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-import contextlib
-import functools
 import json
 import sys
-import time
-from dataclasses import dataclass
-from typing import (
-    Any,
-    Callable,
-    Generator,
-    List,
-    Optional,
-    Tuple,
-    Union,
-)
 
 import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_reduce
-from transformers import PreTrainedTokenizer
 
-from .models import cache
-from .models.cache import (
-    QuantizedKVCache,
-    load_prompt_cache,
-    make_prompt_cache,
-    trim_prompt_cache,
-)
+from .models.cache import QuantizedKVCache, load_prompt_cache
 from .sample_utils import make_sampler
-from .tokenizer_utils import TokenizerWrapper
-from .utils import load
+from .utils import generate, load
 
 DEFAULT_PROMPT = "hello"
 DEFAULT_MAX_TOKENS = 100
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
 DEFAULT_MIN_P = 0.0
-DEFAULT_TOP_K = 0
-DEFAULT_XTC_PROBABILITY = 0.0
-DEFAULT_XTC_THRESHOLD = 0.0
 DEFAULT_MIN_TOKENS_TO_KEEP = 1
-DEFAULT_SEED = None
+DEFAULT_SEED = 0
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 DEFAULT_QUANTIZED_KV_START = 5000
 
@@ -86,11 +60,6 @@ def setup_arg_parser():
         default=DEFAULT_PROMPT,
         help="Message to be processed by the model ('-' reads from stdin)",
     )
-    parser.add_argument(
-        "--prefill-response",
-        default=None,
-        help="Prefill response to be used for the chat template",
-    )
     parser.add_argument(
         "--max-tokens",
         "-m",
@@ -107,33 +76,13 @@ def setup_arg_parser():
     parser.add_argument(
         "--min-p", type=float, default=DEFAULT_MIN_P, help="Sampling min-p"
     )
-    parser.add_argument(
-        "--top-k", type=int, default=DEFAULT_TOP_K, help="Sampling top-k"
-    )
-    parser.add_argument(
-        "--xtc-probability",
-        type=float,
-        default=DEFAULT_XTC_PROBABILITY,
-        help="Probability of XTC sampling to happen each next token",
-    )
-    parser.add_argument(
-        "--xtc-threshold",
-        type=float,
-        default=0.0,
-        help="Thresold the probs of each next token candidate to be sampled by XTC",
-    )
     parser.add_argument(
         "--min-tokens-to-keep",
         type=int,
         default=DEFAULT_MIN_TOKENS_TO_KEEP,
         help="Minimum tokens to keep for min-p sampling.",
     )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=DEFAULT_SEED,
-        help="PRNG seed",
-    )
+    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
     parser.add_argument(
         "--ignore-chat-template",
         action="store_true",
@@ -144,12 +93,6 @@ def setup_arg_parser():
         action="store_true",
         help="Use the default chat template",
     )
-    parser.add_argument(
-        "--chat-template-config",
-        help="Additional config for `apply_chat_template`. Should be a dictionary of"
-        " string keys to values represented as a JSON decodable string.",
-        default=None,
-    )
     parser.add_argument(
         "--verbose",
         type=str2bool,
@@ -188,547 +131,14 @@ def setup_arg_parser():
         type=int,
         default=DEFAULT_QUANTIZED_KV_START,
     )
-    parser.add_argument(
-        "--draft-model",
-        type=str,
-        help="A model to be used for speculative decoding.",
-        default=None,
-    )
-    parser.add_argument(
-        "--num-draft-tokens",
-        type=int,
-        help="Number of tokens to draft when using speculative decoding.",
-        default=3,
-    )
     return parser
 
 
-# A stream on the default device just for generation
-generation_stream = mx.new_stream(mx.default_device())
-
-
-@contextlib.contextmanager
-def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
-    """
-    A context manager to temporarily change the wired limit.
-
-    Note, the wired limit should not be changed during an async eval.  If an
-    async eval could be running pass in the streams to synchronize with prior
-    to exiting the context manager.
-    """
-    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"
-        )
-    old_limit = mx.set_wired_limit(max_rec_size)
-    try:
-        yield None
-    finally:
-        if streams is not None:
-            for s in streams:
-                mx.synchronize(s)
-        else:
-            mx.synchronize()
-        mx.set_wired_limit(old_limit)
-
-
-@dataclass
-class GenerationResponse:
-    """
-    The output of :func:`stream_generate`.
-
-    Args:
-        text (str): The next segment of decoded text. This can be an empty string.
-        token (int): The next token.
-        from_draft (bool): Whether the token was generated by the draft model.
-        logprobs (mx.array): A vector of log probabilities.
-        prompt_tokens (int): The number of tokens in the prompt.
-        prompt_tps (float): The prompt processing tokens-per-second.
-        generation_tokens (int): The number of generated tokens.
-        generation_tps (float): The tokens-per-second for generation.
-        peak_memory (float): The peak memory used so far in GB.
-        finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
-    """
-
-    text: str
-    token: int
-    logprobs: mx.array
-    from_draft: bool
-    prompt_tokens: int
-    prompt_tps: float
-    generation_tokens: int
-    generation_tps: float
-    peak_memory: float
-    finish_reason: Optional[str] = None
-
-
-def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
-    if (
-        kv_bits is not None
-        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
-        and prompt_cache[0].offset > quantized_kv_start
-    ):
-        for i in range(len(prompt_cache)):
-            if isinstance(prompt_cache[i], cache.KVCache):
-                prompt_cache[i] = prompt_cache[i].to_quantized(
-                    group_size=kv_group_size, bits=kv_bits
-                )
-
-
-def generate_step(
-    prompt: mx.array,
-    model: nn.Module,
-    *,
-    max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
-    max_kv_size: Optional[int] = None,
-    prompt_cache: Optional[Any] = None,
-    prefill_step_size: int = 2048,
-    kv_bits: Optional[int] = None,
-    kv_group_size: int = 64,
-    quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[int, int]] = None,
-) -> Generator[Tuple[mx.array, mx.array], None, None]:
-    """
-    A generator producing token ids based on the given prompt from the model.
-
-    Args:
-        prompt (mx.array): The input prompt.
-        model (nn.Module): The model to use for generation.
-        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
-          generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
-          token from a vector of log probabilities. Default: ``None``.
-        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed
-          logits. Default: ``None``.
-        max_kv_size (int, optional): Maximum size of the key-value cache. Old
-          entries (except the first 4 tokens) will be overwritten.
-        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
-          provided, the cache will be updated in place.
-        prefill_step_size (int): Step size for processing the prompt.
-        kv_bits (int, optional): Number of bits to use for KV cache quantization.
-          None implies no cache quantization. Default: ``None``.
-        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
-        quantized_kv_start (int): Step to begin using a quantized KV cache.
-           when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_prorgress_callback (Callable[int, int]): A call-back which takes the
-           prompt tokens processed so far and the total number of prompt tokens.
-
-    Yields:
-        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
-    """
-
-    y = prompt
-    tokens = None
-
-    # Create the KV cache for generation
-    if prompt_cache is None:
-        prompt_cache = cache.make_prompt_cache(
-            model,
-            max_kv_size=max_kv_size,
-        )
-    elif len(prompt_cache) != len(model.layers):
-        raise ValueError("Wrong number of layers in the prompt cache.")
-
-    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
-
-    quantize_cache_fn = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=quantized_kv_start,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
-
-    def _step(y):
-        with mx.stream(generation_stream):
-            logits = model(y[None], cache=prompt_cache)
-            logits = logits[:, -1, :]
-
-            if logits_processors:
-                nonlocal tokens
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
-
-                for processor in logits_processors:
-                    logits = processor(tokens, logits)
-
-            quantize_cache_fn(prompt_cache)
-
-            logprobs = logits - mx.logsumexp(logits, keepdims=True)
-            y = sampler(logprobs)
-            return y, logprobs.squeeze(0)
-
-    with mx.stream(generation_stream):
-        total_prompt_tokens = y.size
-        prompt_processed_tokens = 0
-        while y.size > prefill_step_size:
-            model(y[:prefill_step_size][None], cache=prompt_cache)
-            quantize_cache_fn(prompt_cache)
-            mx.eval([c.state for c in prompt_cache])
-            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            y = y[prefill_step_size:]
-            mx.clear_cache()
-
-        y, logprobs = _step(y)
-
-    mx.async_eval(y, logprobs)
-    n = 0
-    while True:
-        if n != max_tokens:
-            next_y, next_logprobs = _step(y)
-            mx.async_eval(next_y, next_logprobs)
-        if n == 0:
-            mx.eval(y)
-            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
-        if n == max_tokens:
-            break
-        yield y.item(), logprobs
-        if n % 256 == 0:
-            mx.clear_cache()
-        y, logprobs = next_y, next_logprobs
-        n += 1
-
-
-def speculative_generate_step(
-    prompt: mx.array,
-    model: nn.Module,
-    draft_model: nn.Module,
-    *,
-    num_draft_tokens=2,
-    max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
-    prompt_cache: Optional[Any] = None,
-    prefill_step_size: int = 512,
-    kv_bits: Optional[int] = None,
-    kv_group_size: int = 64,
-    quantized_kv_start: int = 0,
-) -> Generator[Tuple[mx.array, mx.array, bool], None, None]:
-    """
-    A generator producing token ids based on the given prompt from the model.
-
-    Args:
-        prompt (mx.array): The input prompt.
-        model (nn.Module): The model to use for generation.
-        draft_model (nn.Module): The draft model for speculative decoding.
-        num_draft_tokens (int, optional): The number of draft tokens for
-          speculative decoding. Default: ``2``.
-        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
-          generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
-          token from a vector of log probabilities. Default: ``None``.
-        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed
-          logits. Default: ``None``.
-        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
-          provided, the cache will be updated in place. The cache must be trimmable.
-        prefill_step_size (int): Step size for processing the prompt.
-        kv_bits (int, optional): Number of bits to use for KV cache quantization.
-          None implies no cache quantization. Default: ``None``.
-        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
-        quantized_kv_start (int): Step to begin using a quantized KV cache.
-           when ``kv_bits`` is non-None. Default: ``0``.
-
-    Yields:
-        Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
-          and a bool indicating if the token was generated by the draft model
-    """
-
-    y = prompt.astype(mx.uint32)
-    prev_tokens = None
-
-    # Create the KV cache for generation
-    if prompt_cache is None:
-        model_cache = cache.make_prompt_cache(model)
-        draft_cache = cache.make_prompt_cache(draft_model)
-    elif len(prompt_cache) != (len(model.layers) + len(draft_model.layers)):
-        raise ValueError("Wrong number of layers in the prompt cache.")
-    else:
-        model_cache = prompt_cache[: len(model.layers)]
-        draft_cache = prompt_cache[len(model.layers) :]
-
-    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
-
-    quantize_cache_fn = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=quantized_kv_start,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    def _process_and_sample(tokens, logits):
-        if logits_processors:
-            for processor in logits_processors:
-                logits = processor(tokens, logits)
-
-        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-        y = sampler(logprobs)
-        return y, logprobs
-
-    def _step(model, cache, y, n_predict=1):
-        with mx.stream(generation_stream):
-            logits = model(y[None], cache=cache)
-            logits = logits[:, -n_predict:, :]
-
-            quantize_cache_fn(cache)
-            if logits_processors:
-                nonlocal prev_tokens
-                out_y, out_logprobs = [], []
-                if n_predict > 1:
-                    y = y[: -(n_predict - 1)]
-                for i in range(n_predict):
-                    prev_tokens = (
-                        mx.concat([prev_tokens, y]) if prev_tokens is not None else y
-                    )
-                    y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
-                    out_y.append(y)
-                    out_logprobs.append(logprobs)
-                return mx.concatenate(out_y, axis=0), mx.concatenate(
-                    out_logprobs, axis=0
-                )
-            else:
-                return _process_and_sample(None, logits.squeeze(0))
-
-    def _prefill(model, cache, y):
-        while y.size > prefill_step_size:
-            model(y[:prefill_step_size][None], cache=cache)
-            quantize_cache_fn(cache)
-            mx.eval([c.state for c in cache])
-            y = y[prefill_step_size:]
-            mx.clear_cache()
-        return y
-
-    def _rewind_cache(num_draft, num_accept):
-        cache.trim_prompt_cache(model_cache, num_draft - num_accept)
-        cache.trim_prompt_cache(draft_cache, max(num_draft - num_accept - 1, 0))
-
-    def _draft_generate(y, num_draft):
-        if num_draft == 0:
-            return mx.array([], mx.uint32)
-        ys = []
-        for _ in range(num_draft):
-            y, _ = _step(draft_model, draft_cache, y)
-            mx.async_eval(y)
-            ys.append(y)
-        return mx.concatenate(ys)
-
-    with mx.stream(generation_stream):
-        draft_y = _prefill(draft_model, draft_cache, y)
-        y = _prefill(model, model_cache, y)
-
-    ntoks = 0
-    # Set these so the finally block doesn't raise
-    num_draft = 0
-    n = 0
-    try:
-        while True:
-            num_draft = min(max_tokens - ntoks, num_draft_tokens)
-            draft_tokens = _draft_generate(draft_y, num_draft)
-            if prev_tokens is not None:
-                prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1]
-            y = mx.concatenate([y, draft_tokens])
-            tokens, logprobs = _step(model, model_cache, y, num_draft + 1)
-            mx.eval(tokens, draft_tokens)
-            draft_tokens = draft_tokens.tolist()
-            tokens = tokens.tolist()
-            n = 0
-            while n < num_draft:
-                tn, dtn, lpn = tokens[n], draft_tokens[n], logprobs[n]
-                if tn != dtn:
-                    break
-                n += 1
-                ntoks += 1
-                yield tn, lpn, True
-                if ntoks == max_tokens:
-                    break
-            if ntoks < max_tokens:
-                ntoks += 1
-                yield tokens[n], logprobs[n], False
-
-            if ntoks == max_tokens:
-                break
-
-            y = mx.array([tokens[n]], mx.uint32)
-            draft_y = y
-
-            # If we accepted all the draft tokens, include the last
-            # draft token in the next draft step since it hasn't been
-            # processed yet by the draft model
-            if n == num_draft:
-                draft_y = mx.concatenate(
-                    [mx.array(draft_tokens[-1:], mx.uint32), draft_y]
-                )
-
-            if prev_tokens is not None:
-                prev_tokens = prev_tokens[: -max(num_draft - n, 1)]
-            _rewind_cache(num_draft, n)
-    finally:
-        _rewind_cache(num_draft, n)
-
-
-def stream_generate(
-    model: nn.Module,
-    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
-    prompt: Union[str, mx.array, List[int]],
-    draft_model: Optional[nn.Module] = None,
-    **kwargs,
-) -> Generator[GenerationResponse, None, None]:
-    """
-    A generator producing text based on the given prompt from the model.
-
-    Args:
-        model (nn.Module): The model to use for generation.
-        tokenizer (PreTrainedTokenizer): The tokenizer.
-        prompt (Union[str, mx.array, List[int]]): The input prompt string or
-          integer tokens.
-        draft_model (Optional[nn.Module]): An optional draft model. If provided
-          then speculative decoding is used. The draft model must use the same
-          tokenizer as the main model. Default: ``None``.
-        kwargs: The remaining options get passed to :func:`generate_step`.
-          See :func:`generate_step` for more details.
-
-    Yields:
-        GenerationResponse: An instance containing the generated text segment and
-            associated metadata. See :class:`GenerationResponse` for details.
-    """
-    if not isinstance(tokenizer, TokenizerWrapper):
-        tokenizer = TokenizerWrapper(tokenizer)
-
-    if not isinstance(prompt, mx.array):
-        if isinstance(prompt, str):
-            # Try to infer if special tokens are needed
-            add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
-                tokenizer.bos_token
-            )
-            prompt = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
-        prompt = mx.array(prompt)
-
-    detokenizer = tokenizer.detokenizer
-
-    if draft_model is None:
-        kwargs.pop("num_draft_tokens", None)
-        token_generator = generate_step(prompt, model, **kwargs)
-        # from_draft always false for non-speculative generation
-        token_generator = (
-            (token, logprobs, False) for token, logprobs in token_generator
-        )
-    else:
-        kwargs.pop("max_kv_size", None)
-        token_generator = speculative_generate_step(
-            prompt, model, draft_model, **kwargs
-        )
-    with wired_limit(model, [generation_stream]):
-        detokenizer.reset()
-        tic = time.perf_counter()
-        for n, (token, logprobs, from_draft) in enumerate(token_generator):
-            if n == 0:
-                prompt_time = time.perf_counter() - tic
-                prompt_tps = prompt.size / prompt_time
-                tic = time.perf_counter()
-            if token in tokenizer.eos_token_ids:
-                break
-
-            detokenizer.add_token(token)
-
-            yield GenerationResponse(
-                text=detokenizer.last_segment,
-                token=token,
-                logprobs=logprobs,
-                from_draft=from_draft,
-                prompt_tokens=prompt.size,
-                prompt_tps=prompt_tps,
-                generation_tokens=n + 1,
-                generation_tps=(n + 1) / (time.perf_counter() - tic),
-                peak_memory=mx.get_peak_memory() / 1e9,
-                finish_reason=None,
-            )
-
-        detokenizer.finalize()
-        yield GenerationResponse(
-            text=detokenizer.last_segment,
-            token=token,
-            logprobs=logprobs,
-            from_draft=from_draft,
-            prompt_tokens=prompt.size,
-            prompt_tps=prompt_tps,
-            generation_tokens=n + 1,
-            generation_tps=(n + 1) / (time.perf_counter() - tic),
-            peak_memory=mx.get_peak_memory() / 1e9,
-            finish_reason="stop" if token in tokenizer.eos_token_ids else "length",
-        )
-
-
-def generate(
-    model: nn.Module,
-    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
-    prompt: Union[str, List[int]],
-    verbose: bool = False,
-    formatter: Optional[Callable] = None,
-    **kwargs,
-) -> str:
-    """
-    Generate a complete response from the model.
-
-    Args:
-       model (nn.Module): The language model.
-       tokenizer (PreTrainedTokenizer): The tokenizer.
-       prompt (Union[str, List[int]]): The input prompt string or integer tokens.
-       verbose (bool): If ``True``, print tokens and timing information.
-           Default: ``False``.
-       kwargs: The remaining options get passed to :func:`stream_generate`.
-          See :func:`stream_generate` for more details.
-    """
-    if formatter is not None:
-        print(
-            "[Warning] Text formatting is deprecated and no longer used. "
-            "The argument will be removed in a future version."
-        )
-    if verbose:
-        print("=" * 10)
-
-    text = ""
-    for response in stream_generate(model, tokenizer, prompt, **kwargs):
-        if verbose:
-            print(response.text, end="", flush=True)
-        text += response.text
-
-    if verbose:
-        print()
-        print("=" * 10)
-        if len(text) == 0:
-            print("No text generated for this prompt")
-            return
-        print(
-            f"Prompt: {response.prompt_tokens} tokens, "
-            f"{response.prompt_tps:.3f} tokens-per-sec"
-        )
-        print(
-            f"Generation: {response.generation_tokens} tokens, "
-            f"{response.generation_tps:.3f} tokens-per-sec"
-        )
-        print(f"Peak memory: {response.peak_memory:.3f} GB")
-    return text
-
-
 def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
-    if args.seed is not None:
-        mx.random.seed(args.seed)
+    mx.random.seed(args.seed)
 
     # Load the prompt cache and metadata if a cache file is provided
     using_cache = args.prompt_cache_file is not None
@@ -773,15 +183,11 @@ def main():
     for eos_token in args.extra_eos_token:
         tokenizer.add_eos_token(eos_token)
 
-    template_kwargs = {}
-    if args.chat_template_config is not None:
-        template_kwargs = json.loads(args.chat_template_config)
-
     if args.use_default_chat_template:
         if tokenizer.chat_template is None:
             tokenizer.chat_template = tokenizer.default_chat_template
     elif using_cache:
-        tokenizer.chat_template = json.loads(metadata["chat_template"])
+        tokenizer.chat_template = metadata["chat_template"]
 
     prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
     prompt = sys.stdin.read() if prompt == "-" else prompt
@@ -791,16 +197,8 @@ def main():
         else:
             messages = []
         messages.append({"role": "user", "content": prompt})
-
-        has_prefill = args.prefill_response is not None
-        if has_prefill:
-            messages.append({"role": "assistant", "content": args.prefill_response})
         prompt = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            continue_final_message=has_prefill,
-            add_generation_prompt=not has_prefill,
-            **template_kwargs,
+            messages, tokenize=False, add_generation_prompt=True
         )
 
         # Treat the prompt as a suffix assuming that the prefix is in the
@@ -810,30 +208,15 @@ def main():
             test_prompt = tokenizer.apply_chat_template(
                 messages,
                 tokenize=False,
-                continue_final_message=has_prefill,
-                add_generation_prompt=not has_prefill,
+                add_generation_prompt=True,
             )
             prompt = prompt[test_prompt.index("<query>") :]
+
         prompt = tokenizer.encode(prompt, add_special_tokens=False)
     else:
         prompt = tokenizer.encode(prompt)
 
-    if args.draft_model is not None:
-        draft_model, draft_tokenizer = load(args.draft_model)
-        if draft_tokenizer.vocab_size != tokenizer.vocab_size:
-            raise ValueError("Draft model tokenizer does not match model tokenizer.")
-    else:
-        draft_model = None
-    sampler = make_sampler(
-        args.temp,
-        args.top_p,
-        args.min_p,
-        args.min_tokens_to_keep,
-        top_k=args.top_k,
-        xtc_probability=args.xtc_probability,
-        xtc_threshold=args.xtc_threshold,
-        xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids),
-    )
+    sampler = make_sampler(args.temp, args.top_p, args.min_p, args.min_tokens_to_keep)
     response = generate(
         model,
         tokenizer,
@@ -846,16 +229,10 @@ def main():
         kv_bits=args.kv_bits,
         kv_group_size=args.kv_group_size,
         quantized_kv_start=args.quantized_kv_start,
-        draft_model=draft_model,
-        num_draft_tokens=args.num_draft_tokens,
     )
     if not args.verbose:
         print(response)
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.generate...` directly is deprecated."
-        " Use `mlx_lm.generate...` or `python -m mlx_lm generate ...` instead."
-    )
     main()

mlx_lm/lora.py

@@ -7,13 +7,13 @@ import re
 import types
 from pathlib import Path
 
-import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
 import numpy as np
 import yaml
 
-from .tuner.datasets import CacheDataset, load_dataset
+from .tokenizer_utils import TokenizerWrapper
+from .tuner.datasets import load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
     build_schedule,
@@ -43,11 +43,6 @@ CONFIG_DEFAULTS = {
     "model": "mlx_model",
     "train": False,
     "fine_tune_type": "lora",
-    "optimizer": "adam",
-    "optimizer_config": {
-        "adam": {},
-        "adamw": {},
-    },
     "data": "data/",
     "seed": 0,
     "num_layers": 16,
@@ -63,11 +58,8 @@ CONFIG_DEFAULTS = {
     "test": False,
     "test_batches": 500,
     "max_seq_length": 2048,
-    "config": None,
-    "grad_checkpoint": False,
     "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 10.0},
-    "mask_prompt": False,
+    "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
 }
 
 
@@ -75,7 +67,6 @@ def build_parser():
     parser = argparse.ArgumentParser(description="LoRA or QLoRA finetuning.")
     parser.add_argument(
         "--model",
-        type=str,
         help="The path to the local model directory or Hugging Face repo.",
     )
 
@@ -98,21 +89,9 @@ def build_parser():
         "--fine-tune-type",
         type=str,
         choices=["lora", "dora", "full"],
+        default="lora",
         help="Type of fine-tuning to perform: lora, dora, or full.",
     )
-    parser.add_argument(
-        "--optimizer",
-        type=str,
-        choices=["adam", "adamw"],
-        default=None,
-        help="Optimizer to use for training: adam or adamw",
-    )
-    parser.add_argument(
-        "--mask-prompt",
-        action="store_true",
-        help="Mask the prompt in the loss when training",
-        default=None,
-    )
     parser.add_argument(
         "--num-layers",
         type=int,
@@ -170,7 +149,7 @@ def build_parser():
     parser.add_argument(
         "-c",
         "--config",
-        type=str,
+        default=None,
         help="A YAML configuration file with the training options",
     )
     parser.add_argument(
@@ -179,27 +158,21 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
-    parser.add_argument("--seed", type=int, help="The PRNG seed")
+    parser.add_argument("--seed", type=int, default=None, help="The PRNG seed")
     return parser
 
 
 def train_model(
     args,
     model: nn.Module,
+    tokenizer: TokenizerWrapper,
     train_set,
     valid_set,
     training_callback: TrainingCallback = None,
 ):
-    mx.random.seed(args.seed)
     model.freeze()
-    if args.num_layers > len(model.layers):
-        raise ValueError(
-            f"Requested to train {args.num_layers} layers "
-            f"but the model only has {len(model.layers)} layers."
-        )
-
     if args.fine_tune_type == "full":
-        for l in model.layers[-max(args.num_layers, 0) :]:
+        for l in model.layers[-min(args.num_layers, 0) :]:
             l.unfreeze()
     elif args.fine_tune_type in ["lora", "dora"]:
         # Convert linear layers to lora/dora layers and unfreeze in the process
@@ -238,36 +211,31 @@ def train_model(
         grad_checkpoint=args.grad_checkpoint,
     )
 
-    # Initialize the selected optimizer
-    lr = build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
-
-    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
-    else:
-        raise ValueError(f"Unsupported optimizer: {optimizer_name}")
-
-    opt = opt_class(learning_rate=lr, **optimizer_config)
-
+    model.train()
+    opt = optim.Adam(
+        learning_rate=(
+            build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
+        )
+    )
     # Train model
     train(
         model=model,
+        tokenizer=tokenizer,
         args=training_args,
         optimizer=opt,
-        train_dataset=CacheDataset(train_set),
-        val_dataset=CacheDataset(valid_set),
+        train_dataset=train_set,
+        val_dataset=valid_set,
         training_callback=training_callback,
     )
 
 
-def evaluate_model(args, model: nn.Module, test_set):
+def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
+    model.eval()
+
     test_loss = evaluate(
         model=model,
-        dataset=CacheDataset(test_set),
+        dataset=test_set,
+        tokenizer=tokenizer,
         batch_size=args.batch_size,
         num_batches=args.test_batches,
         max_seq_length=args.max_seq_length,
@@ -294,13 +262,13 @@ def run(args, training_callback: TrainingCallback = None):
 
     elif args.train:
         print("Training")
-        train_model(args, model, train_set, valid_set, training_callback)
+        train_model(args, model, tokenizer, train_set, valid_set, training_callback)
     else:
         raise ValueError("Must provide at least one of --train or --test")
 
     if args.test:
         print("Testing")
-        evaluate_model(args, model, test_set)
+        evaluate_model(args, model, tokenizer, test_set)
 
 
 def main():
@@ -326,8 +294,4 @@ def main():
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.lora...` directly is deprecated."
-        " Use `mlx_lm.lora...` or `python -m mlx_lm lora ...` instead."
-    )
     main()

mlx_lm/manage.py

@@ -2,22 +2,7 @@ import argparse
 from typing import List, Union
 
 from huggingface_hub import scan_cache_dir
-
-
-def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
-    """
-    Inspired by:
-    - stackoverflow.com/a/8356620/593036
-    - stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
-    """
-    col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
-    row_format = ("{{:{}}} " * len(headers)).format(*col_widths)
-    lines = []
-    lines.append(row_format.format(*headers))
-    lines.append(row_format.format(*["-" * w for w in col_widths]))
-    for row in rows:
-        lines.append(row_format.format(*row))
-    return "\n".join(lines)
+from transformers.commands.user import tabulate
 
 
 def ask_for_confirmation(message: str) -> bool:
@@ -136,8 +121,4 @@ def main():
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.manage...` directly is deprecated."
-        " Use `mlx_lm.manage...` or `python -m mlx_lm manage ...` instead."
-    )
     main()

mlx_lm/merge.py

@@ -169,8 +169,4 @@ def main():
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.merge...` directly is deprecated."
-        " Use `mlx_lm.merge...` or `python -m mlx_lm merge ...` instead."
-    )
     main()

mlx_lm/models/baichuan_m1.py

@@ -1,226 +0,0 @@
-# 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 CacheList, KVCache, MambaCache, RotatingKVCache
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    vocab_size: int
-    hidden_size: int
-    intermediate_size: int
-    num_hidden_layers: int
-    num_attention_heads: int
-    num_key_value_heads: int
-    rope_theta: float
-    sliding_window: int
-    sliding_window_layers: List[int]
-    conv_window: int
-    rms_norm_eps: float
-    model_type: str = "baichuan_m1"
-    num_swa_attention_heads: Optional[int] = None
-    num_swa_key_value_heads: Optional[int] = None
-    tie_word_embeddings: bool = False
-
-
-class Attention(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: Optional[int] = None):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        if layer_idx is None:
-            raise ValueError("Layer index must be provided to Attention module.")
-
-        self.is_swa = layer_idx in config.sliding_window_layers
-        self.num_heads = (
-            config.num_swa_attention_heads
-            if self.is_swa and config.num_swa_attention_heads
-            else config.num_attention_heads
-        )
-        self.num_kv_heads = (
-            config.num_swa_key_value_heads
-            if self.is_swa and config.num_swa_key_value_heads
-            else config.num_key_value_heads
-        )
-
-        self.hidden_size = config.hidden_size
-        self.head_dim = self.hidden_size // self.num_heads
-        assert self.head_dim * self.num_heads == self.hidden_size
-
-        self.scale = self.head_dim**-0.5
-
-        self.W_pack = nn.Linear(
-            config.hidden_size,
-            self.hidden_size + 2 * self.num_kv_heads * self.head_dim,
-            bias=False,
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, config.hidden_size, bias=False
-        )
-
-        self.rope = nn.RoPE(self.head_dim, traditional=False, base=config.rope_theta)
-
-        self.conv_window = config.conv_window
-        assert self.conv_window == 2
-        self.conv_k = mx.zeros((1, 1, self.num_kv_heads, 1, self.conv_window))
-        self.conv_v = mx.zeros((1, 1, self.num_kv_heads, 1, self.conv_window))
-
-    def _custom_convolution(self, u, weights, state=None):
-        B, H, L, D = u.shape
-        weights = weights.reshape((1, H, self.conv_window, 1, 1))
-        w0 = weights[:, :, 0]
-        w1 = weights[:, :, 1]
-        if state is None:
-            state = mx.zeros((B, H, 1, D), u.dtype)
-        if L > 1:
-            u_prev = mx.concatenate([state, u[:, :, :-1]], axis=2)
-        else:
-            u_prev = state
-        return u_prev * w0 + u * w1
-
-    def __call__(
-        self, x: mx.array, mask: mx.array = None, cache: Any = None
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        proj = self.W_pack(x)
-        q, k, v = mx.split(proj, (D, D + self.num_kv_heads * self.head_dim), axis=-1)
-
-        q = q.reshape(B, L, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
-        k = k.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
-        v = v.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
-
-        if cache is not None:
-            offset = cache[1].offset
-            last_k, last_v = cache[0][0], cache[0][1]
-        else:
-            offset = 0
-            last_k, last_v = None, None
-
-        k_init = k
-        v_init = v
-        k = self._custom_convolution(k, self.conv_k, state=last_k)
-        v = self._custom_convolution(v, self.conv_v, state=last_v)
-        q = self.rope(q, offset=offset)
-        k = self.rope(k, offset=offset)
-
-        if cache is not None:
-            k, v = cache[1].update_and_fetch(k, v)
-            if L > 0:
-                cache[0][0] = k_init[:, :, -1:, :]
-                cache[0][1] = v_init[:, :, -1:, :]
-
-        out = scaled_dot_product_attention(
-            q, k, v, cache=cache[1], scale=self.scale, mask=mask
-        )
-        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(out)
-
-
-class MLP(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.gate_proj = nn.Linear(
-            config.hidden_size, config.intermediate_size, bias=False
-        )
-        self.up_proj = nn.Linear(
-            config.hidden_size, config.intermediate_size, bias=False
-        )
-        self.down_proj = nn.Linear(
-            config.intermediate_size, config.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 DecoderLayer(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.self_attn = Attention(config, layer_idx)
-        self.mlp = 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: mx.array = None, cache: Any = None
-    ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), mask, cache)
-        x = x + r
-        r = self.mlp(self.post_attention_layernorm(x))
-        return x + r
-
-
-class BaichuanModel(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.args = config
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
-        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def __call__(
-        self, inputs: mx.array, mask: mx.array = None, cache: Any = None
-    ) -> mx.array:
-        x = self.embed_tokens(inputs)
-        if mask is None:
-            if cache is not None:
-                c = [cache[0][1]]
-            mask = create_attention_mask(x, c)
-        if cache is None:
-            cache = [None] * len(self.layers)
-        for layer, c in zip(self.layers, cache):
-            x = layer(x, mask, c)
-        return self.norm(x)
-
-
-class Model(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.model_type = config.model_type
-        self.model = BaichuanModel(config)
-        self.tie_word_embeddings = config.tie_word_embeddings
-        if not config.tie_word_embeddings:
-            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-    def make_cache(self) -> List[Any]:
-        caches = []
-        for i, layer in enumerate(self.model.layers):
-            is_swa = i in self.config.sliding_window_layers
-            conv_cache = MambaCache()
-            if is_swa:
-                kv_cache = RotatingKVCache(max_size=self.config.sliding_window)
-            else:
-                kv_cache = KVCache()
-            caches.append(CacheList(conv_cache, kv_cache))
-        return caches
-
-    def sanitize(self, weights: dict) -> dict:
-        is_quantized = "lm_head.scales" in weights
-        if not is_quantized and "lm_head.weight" in weights:
-            w = weights["lm_head.weight"]
-            dtype = w.dtype
-            w = w.astype(mx.float32)
-            norm = mx.linalg.norm(w, axis=-1, keepdims=True)
-            w = (w / (norm + 1e-7)).astype(dtype)
-            weights["lm_head.weight"] = w
-        return weights
-
-    def __call__(
-        self, inputs: mx.array, mask: mx.array = None, cache: Any = None
-    ) -> mx.array:
-        outputs = self.model(inputs, mask, cache)
-        return self.lm_head(outputs)
-
-    @property
-    def layers(self) -> List[nn.Module]:
-        return self.model.layers

mlx_lm/models/base.py

@@ -33,33 +33,29 @@ def create_causal_mask(
     linds = mx.arange(offset, offset + N) if offset else rinds
     linds = linds[:, None]
     rinds = rinds[None]
-    mask = linds >= rinds
+    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)
-    return mask
+        mask = mask | (rinds >= lengths)
+    return mask * -1e9
 
 
-def create_attention_mask(
-    h: mx.array, cache: Optional[Any] = None, return_array: bool = False
-):
+def create_attention_mask(h: mx.array, cache: Optional[Any] = None):
     T = h.shape[1]
     if T > 1:
-        offset = 0
         window_size = None
+        offset = 0
         if cache is not None and cache[0] is not None:
             c = cache[0]
-            offset = c.offset
             if hasattr(c, "max_size"):
+                offset = min(c.max_size, c.offset)
                 window_size = c.max_size
-                offset = min(window_size, offset)
-                return_array = return_array or offset + T > window_size
-        if return_array:
-            return create_causal_mask(T, offset, window_size=window_size)
-        else:
-            return "causal"
+            else:
+                offset = c.offset
+        mask = create_causal_mask(T, offset, window_size=window_size)
+        mask = mask.astype(h.dtype)
     else:
         mask = None
     return mask
@@ -89,15 +85,7 @@ def quantized_scaled_dot_product_attention(
         queries, *q_keys, transpose=True, group_size=group_size, bits=bits
     )
     if mask is not None:
-        if isinstance(mask, str):
-            qL, kL = scores.shape[-2:]
-            q_indices = mx.arange(kL - qL, kL)
-            k_indices = mx.arange(kL)
-            mask = q_indices[:, None] >= k_indices[None]
-        if mask.dtype == mx.bool_:
-            scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
-        else:
-            scores += mask
+        scores += mask
     scores = mx.softmax(scores, axis=-1, precise=True)
     out = mx.quantized_matmul(
         scores, *q_values, transpose=False, group_size=group_size, bits=bits

mlx_lm/models/cache.py

@@ -436,76 +436,3 @@ class MambaCache(_BaseCache):
     @state.setter
     def state(self, v):
         self.cache = v
-
-
-class ChunkedKVCache(KVCache):
-    def __init__(self, chunk_size=None):
-        super().__init__()
-        self.chunk_size = chunk_size
-        self.start_position = 0
-
-    def maybe_trim_front(self):
-        # Maintain the cache below the chunk size
-        if self.keys is not None and self.keys.shape[2] >= self.chunk_size:
-            self.start_position += self.keys.shape[2] - self.chunk_size
-            self.keys = self.keys[..., -self.chunk_size :, :]
-            self.values = self.values[..., -self.chunk_size :, :]
-
-    def update_and_fetch(self, keys, values):
-        prev = self.offset - self.start_position
-        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
-            B, n_kv_heads, _, k_head_dim = keys.shape
-            v_head_dim = values.shape[3]
-            n_steps = (self.step + keys.shape[2] - 1) // self.step
-            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
-            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
-            new_k = mx.zeros(k_shape, keys.dtype)
-            new_v = mx.zeros(v_shape, values.dtype)
-            if self.keys is not None:
-                if prev % self.step != 0:
-                    self.keys = self.keys[..., :prev, :]
-                    self.values = self.values[..., :prev, :]
-                self.keys = mx.concatenate([self.keys, new_k], axis=2)
-                self.values = mx.concatenate([self.values, new_v], axis=2)
-            else:
-                self.keys, self.values = new_k, new_v
-
-        self.offset += keys.shape[2]
-        end = self.offset - self.start_position
-        self.keys[..., prev:end, :] = keys
-        self.values[..., prev:end, :] = values
-        return self.keys[..., :end, :], self.values[..., :end, :]
-
-    def trim(self, n):
-        n = min(self.offset - self.start_position, n)
-        self.offset -= n
-        return n
-
-    @property
-    def meta_state(self):
-        return tuple(map(str, (self.chunk_size, self.start_position)))
-
-    @meta_state.setter
-    def meta_state(self, v):
-        self.chunk_size, self.start_position = map(int, v)
-
-
-class CacheList(KVCache):
-    def __init__(self, *caches):
-        self.caches = caches
-
-    def __getitem__(self, idx):
-        return self.caches[idx]
-
-    @property
-    def state(self):
-        return [s for c in self.caches for s in c.state]
-
-    @state.setter
-    def state(self, v):
-        state_lens = [len(c.state) for c in self.caches]
-        start = 0
-        for c in self.caches:
-            l = len(c.state)
-            c.state = v[start : start + l]
-            start += l

mlx_lm/models/cohere2.py

@@ -83,22 +83,15 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        if self.use_sliding_window and isinstance(mask, mx.array):
+        if self.use_sliding_window and mask is not None:
             key_len = keys.shape[-2]
             if mask.shape[-1] != key_len:
                 mask = mask[..., -key_len:]
 
-        # TODO: maybe remove cast once fused mask is supported since attention
-        # may be in higher precision
-        sdpa_type = mx.float32 if queries.dtype == mx.float16 else queries.dtype
         output = scaled_dot_product_attention(
-            queries.astype(sdpa_type),
-            keys,
-            values,
-            cache=cache,
-            scale=self.scale,
-            mask=mask,
-        ).astype(queries.dtype)
+            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)
 
@@ -133,11 +126,9 @@ class TransformerBlock(nn.Module):
         mask: Optional[mx.array] = None,
         cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
-
         h = self.input_layernorm(x)
         attn_h = self.self_attn(h, mask, cache)
         ff_h = self.mlp(h)
-
         return attn_h + ff_h + x
 
 
@@ -165,27 +156,14 @@ class CohereModel(nn.Module):
     ):
         h = self.embed_tokens(inputs)
 
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
         if cache is None:
             cache = [None] * len(self.layers)
 
-        if mask is None:
-            j = self.args.sliding_window_pattern
-            full_mask = create_attention_mask(h, cache[j - 1 : j])
-            sliding_window_mask = create_attention_mask(h, cache)
-
-        for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            is_global = (
-                i % self.args.sliding_window_pattern
-                == self.args.sliding_window_pattern - 1
-            )
-
-            local_mask = mask
-            if mask is None and is_global:
-                local_mask = full_mask
-            elif mask is None:
-                local_mask = sliding_window_mask
-
-            h = layer(h, local_mask, c)
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
 
         return self.norm(h)
 

mlx_lm/models/deepseek_v2.py

@@ -148,7 +148,7 @@ class DeepseekV2Attention(nn.Module):
             self.q_a_proj = nn.Linear(
                 self.hidden_size, self.q_lora_rank, bias=config.attention_bias
             )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
             self.q_b_proj = nn.Linear(
                 self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
             )
@@ -158,7 +158,7 @@ class DeepseekV2Attention(nn.Module):
             self.kv_lora_rank + self.qk_rope_head_dim,
             bias=config.attention_bias,
         )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
         self.kv_b_proj = nn.Linear(
             self.kv_lora_rank,
             self.num_heads
@@ -282,12 +282,12 @@ class MoEGate(nn.Module):
         if self.topk_method == "group_limited_greedy":
             bsz, seq_len = x.shape[:2]
             scores = scores.reshape(bsz, seq_len, self.n_group, -1)
-            group_scores = scores.max(axis=-1, keepdims=True)
+            group_scores = scores.max(axis=-1)
             k = self.n_group - self.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, scores.dtype), axis=-2
-            )
+            group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-1)[..., :k]
+            batch_idx = mx.expand_dims(mx.arange(bsz), (1, 2))
+            seq_idx = mx.expand_dims(mx.arange(seq_len), (0, 2))
+            scores[batch_idx, seq_idx, group_idx] = 0.0
             scores = scores.reshape(bsz, seq_len, -1)
 
         k = self.top_k
@@ -364,32 +364,8 @@ class DeepseekV2Model(nn.Module):
             DeepseekV2DecoderLayer(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)
 
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.num_layers = layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
-
     def __call__(
         self,
         x: mx.array,
@@ -398,27 +374,14 @@ class DeepseekV2Model(nn.Module):
     ) -> mx.array:
         h = self.embed_tokens(x)
 
-        pipeline_rank = self.pipeline_rank
-        pipeline_size = self.pipeline_size
         if mask is None:
             mask = create_attention_mask(h, cache)
 
         if cache is None:
-            cache = [None] * self.num_layers
+            cache = [None] * len(self.layers)
 
-        # Receive from the previous process in the pipeline
-        if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
-
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
-
-        # Send to the next process in the pipeline
-        if pipeline_rank != 0:
-            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
-
-        # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
 
         return self.norm(h)
 
@@ -455,4 +418,4 @@ class Model(nn.Module):
 
     @property
     def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
+        return self.model.layers

mlx_lm/models/deepseek_v3.py

@@ -1,536 +0,0 @@
-# Copyright © 2024 Apple Inc.
-
-import math
-from dataclasses import dataclass
-from functools import partial
-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 .switch_layers import SwitchGLU
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "deepseek_v3"
-    vocab_size: int = 102400
-    hidden_size: int = 4096
-    intermediate_size: int = 11008
-    moe_intermediate_size: int = 1407
-    num_hidden_layers: int = 30
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 32
-    n_shared_experts: Optional[int] = None
-    n_routed_experts: Optional[int] = None
-    routed_scaling_factor: float = 1.0
-    kv_lora_rank: int = 512
-    q_lora_rank: int = 1536
-    qk_rope_head_dim: int = 64
-    v_head_dim: int = 128
-    qk_nope_head_dim: int = 128
-    topk_method: str = "noaux_tc"
-    scoring_func: str = "sigmoid"
-    norm_topk_prob: bool = True
-    n_group: Optional[int] = None
-    topk_group: Optional[int] = None
-    num_experts_per_tok: Optional[int] = None
-    moe_layer_freq: int = 1
-    first_k_dense_replace: int = 0
-    max_position_embeddings: int = 2048
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 10000.0
-    rope_scaling: Dict = None
-    attention_bias: bool = False
-
-
-def yarn_find_correction_dim(
-    num_rotations, dim, base=10000, max_position_embeddings=2048
-):
-    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
-        2 * math.log(base)
-    )
-
-
-def yarn_find_correction_range(
-    low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
-):
-    low = math.floor(
-        yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
-    )
-    high = math.ceil(
-        yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
-    )
-    return max(low, 0), min(high, dim - 1)
-
-
-def yarn_get_mscale(scale=1, mscale=1):
-    if scale <= 1:
-        return 1.0
-    return 0.1 * mscale * math.log(scale) + 1.0
-
-
-def yarn_linear_ramp_mask(min_val, max_val, dim):
-    if min_val == max_val:
-        max_val += 0.001  # Prevent singularity
-
-    linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
-    return mx.clip(linear_func, 0, 1)
-
-
-class DeepseekV3YarnRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim,
-        max_position_embeddings=2048,
-        base=10000,
-        scaling_factor=1.0,
-        original_max_position_embeddings=4096,
-        beta_fast=32,
-        beta_slow=1,
-        mscale=1,
-        mscale_all_dim=0,
-    ):
-        super().__init__()
-        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
-            scaling_factor, mscale_all_dim
-        )
-        freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
-        freq_inter = scaling_factor * freq_extra
-        low, high = yarn_find_correction_range(
-            beta_fast,
-            beta_slow,
-            dim,
-            base,
-            original_max_position_embeddings,
-        )
-        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
-        self._freqs = (freq_inter * freq_extra) / (
-            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
-        )
-
-    def __call__(self, x, offset=0):
-        if self.mscale != 1.0:
-            x = self.mscale * x
-        return mx.fast.rope(
-            x,
-            x.shape[-1],
-            traditional=True,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
-# 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 DeepseekV3Attention(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_theta
-        self.q_lora_rank = config.q_lora_rank
-        self.qk_rope_head_dim = config.qk_rope_head_dim
-        self.kv_lora_rank = config.kv_lora_rank
-        self.v_head_dim = config.v_head_dim
-        self.qk_nope_head_dim = config.qk_nope_head_dim
-        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
-
-        self.scale = self.q_head_dim**-0.5
-
-        if self.q_lora_rank is None:
-            self.q_proj = nn.Linear(
-                self.hidden_size, self.num_heads * self.q_head_dim, bias=False
-            )
-        else:
-            self.q_a_proj = nn.Linear(
-                self.hidden_size, self.q_lora_rank, bias=config.attention_bias
-            )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
-            self.q_b_proj = nn.Linear(
-                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
-            )
-
-        self.kv_a_proj_with_mqa = nn.Linear(
-            self.hidden_size,
-            self.kv_lora_rank + self.qk_rope_head_dim,
-            bias=config.attention_bias,
-        )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
-        )
-
-        self.o_proj = nn.Linear(
-            self.num_heads * self.v_head_dim,
-            self.hidden_size,
-            bias=config.attention_bias,
-        )
-
-        if self.config.rope_scaling is not None:
-            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
-            scaling_factor = self.config.rope_scaling["factor"]
-            if mscale_all_dim:
-                mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
-                self.scale = self.scale * mscale * mscale
-
-            rope_kwargs = {
-                key: self.config.rope_scaling[key]
-                for key in [
-                    "original_max_position_embeddings",
-                    "beta_fast",
-                    "beta_slow",
-                    "mscale",
-                    "mscale_all_dim",
-                ]
-                if key in self.config.rope_scaling
-            }
-            self.rope = DeepseekV3YarnRotaryEmbedding(
-                dim=self.qk_rope_head_dim,
-                max_position_embeddings=self.max_position_embeddings,
-                scaling_factor=scaling_factor,
-                base=self.rope_theta,
-                **rope_kwargs,
-            )
-        else:
-            self.rope = nn.RoPE(
-                dims=self.qk_rope_head_dim,
-                base=self.rope_theta,
-                traditional=True,
-            )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        if self.q_lora_rank is None:
-            q = self.q_proj(x)
-        else:
-            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
-
-        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
-        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
-        compressed_kv = self.kv_a_proj_with_mqa(x)
-        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
-
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
-
-        if cache is not None:
-            q_pe = self.rope(q_pe, cache.offset)
-            k_pe = self.rope(k_pe, cache.offset)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys, values = cache.update_and_fetch(
-                mx.concatenate([k_nope, k_pe], axis=-1), values
-            )
-        else:
-            q_pe = self.rope(q_pe)
-            k_pe = self.rope(k_pe)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys = mx.concatenate([k_nope, k_pe], axis=-1)
-
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
-
-        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 DeepseekV3MLP(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,
-):
-
-    k = top_k
-    scores = mx.sigmoid(gates.astype(mx.float32))
-    orig_scores = scores
-    scores = scores + e_score_correction_bias
-    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, 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 DeepseekV3MoE(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,
-            activation=clipped_silu,
-        )
-
-        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 = DeepseekV3MLP(
-                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 DeepseekV3DecoderLayer(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.self_attn = DeepseekV3Attention(config)
-        self.mlp = (
-            DeepseekV3MoE(config)
-            if (
-                config.n_routed_experts is not None
-                and layer_idx >= config.first_k_dense_replace
-                and layer_idx % config.moe_layer_freq == 0
-            )
-            else DeepseekV3MLP(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 DeepseekV3Model(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 = [
-            DeepseekV3DecoderLayer(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)
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        h = self.embed_tokens(x)
-
-        pipeline_rank = self.pipeline_rank
-        pipeline_size = self.pipeline_size
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
-        if cache is None:
-            cache = [None] * self.num_layers
-
-        # Receive from the previous process in the pipeline
-
-        if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
-
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
-
-        # Send to the next process in the pipeline
-        if pipeline_rank != 0:
-            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
-
-        # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
-
-        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 = DeepseekV3Model(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):
-        def dequant(weight, scale_inv):
-            dtype = weight.dtype
-            bs = 128  # block size
-            m, n = weight.shape
-            pad_bottom = (-m) % bs
-            pad_side = (-n) % bs
-            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
-            weight = weight.reshape(
-                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
-            )
-            weight = (weight * scale_inv[:, None, :, None]).reshape(
-                m + pad_bottom, n + pad_side
-            )
-            return weight[:m, :n].astype(dtype)
-
-        # Dequantize
-        new_weights = {}
-        for k, v in weights.items():
-            if "weight_scale_inv" in k:
-                scale_inv = v
-                wk = k.replace("_scale_inv", "")
-                weight = weights[wk]
-                weight = dequant(weight, scale_inv)
-                new_weights[wk] = weight
-            elif k not in new_weights:
-                new_weights[k] = v
-        weights = new_weights
-
-        # 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 and any unused precomputed rotary freqs
-        return {
-            k: v
-            for k, v in weights.items()
-            if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
-        }
-
-    @property
-    def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
-
-    def cast_predicate(self):
-        def predicate(k):
-            return "e_score_correction_bias" not in k
-
-        return predicate

mlx_lm/models/gemma2.py

@@ -94,12 +94,7 @@ class Attention(nn.Module):
         scores *= self.attn_logit_softcapping
 
         if mask is not None:
-            if mask.dtype == mx.bool_:
-                scores = mx.where(
-                    mask, scores, mx.array(mx.finfo(scores.dtype).min, scores.dtype)
-                )
-            else:
-                scores = scores + mask
+            scores = scores + mask
         scores = mx.softmax(scores, precise=True, axis=-1)
         output = scores @ values
         if self.repeats > 1:
@@ -172,7 +167,7 @@ class GemmaModel(nn.Module):
         h = h * (self.args.hidden_size**0.5)
 
         if mask is None:
-            mask = create_attention_mask(h, cache, return_array=True)
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)

mlx_lm/models/gemma3.py

@@ -1,61 +0,0 @@
-# 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 gemma3_text
-from .base import BaseModelArgs
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    text_config: dict
-    vocab_size: int = 262208
-
-    def __post_init__(self):
-        self.text_config["vocab_size"] = self.vocab_size
-        self.text_config["num_attention_heads"] = self.text_config.get(
-            "num_attention_heads", 8
-        )
-        self.text_config["num_key_value_heads"] = self.text_config.get(
-            "num_key_value_heads", 4
-        )
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.language_model = gemma3_text.Model(
-            gemma3_text.ModelArgs.from_dict(args.text_config)
-        )
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        mask: Optional[mx.array] = None,
-    ):
-        return self.language_model(inputs, cache=cache, mask=mask)
-
-    def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights.pop("vision_tower", None)
-        weights.pop("multi_modal_projector", None)
-        lm_weights = dict(tree_flatten(weights["language_model"]))
-        lm_weights = self.language_model.sanitize(lm_weights)
-        weights["language_model"] = tree_unflatten(list(lm_weights.items()))
-        return dict(tree_flatten(weights))
-
-    @property
-    def layers(self):
-        return self.language_model.layers
-
-    def make_cache(self):
-        return self.language_model.make_cache()

mlx_lm/models/gemma3_text.py

@@ -1,248 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-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_attention_mask, scaled_dot_product_attention
-from .cache import KVCache, RotatingKVCache
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    hidden_size: int = 1152
-    num_hidden_layers: int = 26
-    intermediate_size: int = 6912
-    num_attention_heads: int = 4
-    head_dim: int = 256
-    rms_norm_eps: float = 1.0e-6
-    vocab_size: int = 262144
-    num_key_value_heads: int = 1
-    rope_global_base_freq: float = 1_000_000.0
-    rope_local_base_freq: float = 10_000.0
-    rope_traditional: bool = False
-    query_pre_attn_scalar: float = 256
-    sliding_window: int = 512
-    sliding_window_pattern: int = 6
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        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.repeats = n_heads // n_kv_heads
-        self.head_dim = head_dim = args.head_dim
-        self.layer_idx = layer_idx
-
-        self.scale = args.query_pre_attn_scalar**-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 = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
-        self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
-        self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
-
-        self.rope = nn.RoPE(
-            head_dim,
-            traditional=args.rope_traditional,
-            base=(
-                args.rope_local_base_freq
-                if self.is_sliding
-                else args.rope_global_base_freq
-            ),
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, _ = 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)
-
-        queries = self.q_norm(queries)
-        keys = self.k_norm(keys)
-
-        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)
-
-        # Sliding window
-        if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
-            mask = mask[..., -keys.shape[-2] :]
-        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 RMSNorm(nn.Module):
-    def __init__(self, dims: int, eps: float = 1e-5):
-        super().__init__()
-        self.weight = mx.ones((dims,))
-        self.eps = eps
-
-    def __call__(self, x):
-        return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
-
-
-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.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
-
-
-@partial(mx.compile, shapeless=True)
-def clip_residual(x, y):
-    if x.dtype != mx.float16:
-        return x + y
-    bound = mx.finfo(mx.float16).max
-    return mx.clip(x.astype(mx.float32) + y.astype(mx.float32), -bound, bound).astype(
-        mx.float16
-    )
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.num_attention_heads = args.num_attention_heads
-        self.hidden_size = args.hidden_size
-        self.self_attn = Attention(args, layer_idx)
-        self.mlp = MLP(args.hidden_size, args.intermediate_size)
-        self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.pre_feedforward_layernorm = RMSNorm(
-            args.hidden_size, eps=args.rms_norm_eps
-        )
-        self.post_feedforward_layernorm = 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 = clip_residual(x, self.post_attention_layernorm(r))
-        r = self.mlp(self.pre_feedforward_layernorm(h))
-        out = clip_residual(h, self.post_feedforward_layernorm(r))
-        return out
-
-
-class Gemma3Model(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)
-        self.layers = [
-            TransformerBlock(args=args, layer_idx=layer_idx)
-            for layer_idx in range(args.num_hidden_layers)
-        ]
-        self.norm = 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)
-        h *= mx.array(self.args.hidden_size**0.5, mx.bfloat16).astype(h.dtype)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        if mask is None:
-            j = self.args.sliding_window_pattern
-            full_mask = create_attention_mask(h, cache[j - 1 : j])
-            sliding_window_mask = create_attention_mask(h, cache)
-
-        for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            is_global = (
-                i % self.args.sliding_window_pattern
-                == self.args.sliding_window_pattern - 1
-            )
-
-            local_mask = mask
-            if mask is None and is_global:
-                local_mask = full_mask
-            elif mask is None:
-                local_mask = sliding_window_mask
-
-            h = layer(h, local_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 = Gemma3Model(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        mask: Optional[mx.array] = None,
-    ):
-        out = self.model(inputs, mask, cache)
-        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"]
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.layers
-
-    def make_cache(self):
-        caches = []
-        for i in range(self.args.num_hidden_layers):
-            if (
-                i % self.args.sliding_window_pattern
-                == self.args.sliding_window_pattern - 1
-            ):
-                caches.append(KVCache())
-            else:
-                caches.append(
-                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
-                )
-        return caches

mlx_lm/models/glm4.py

@@ -1,183 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from dataclasses import dataclass
-from typing import Any, Optional
-
-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
-    hidden_size: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    attention_bias: bool
-    head_dim: int
-    rms_norm_eps: float
-    vocab_size: int
-    num_key_value_heads: int
-    partial_rotary_factor: float
-    rope_theta: float
-    rope_traditional: bool = True
-    max_position_embeddings: int = 32768
-
-
-class Glm4MLP(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.gate_up_proj = nn.Linear(
-            args.hidden_size, 2 * args.intermediate_size, bias=False
-        )
-        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
-
-    def __call__(self, x) -> mx.array:
-        x = self.gate_up_proj(x)
-        gate, up_states = mx.split(x, 2, axis=-1)
-        return self.down_proj(nn.silu(gate) * up_states)
-
-
-class Glm4Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.head_dim = getattr(
-            args, "head_dim", args.hidden_size // args.num_attention_heads
-        )
-        self.n_heads = args.num_attention_heads
-        self.n_kv_heads = args.num_key_value_heads
-        self.scale = self.head_dim**-0.5
-
-        self.q_proj = nn.Linear(
-            args.hidden_size,
-            args.num_attention_heads * self.head_dim,
-            bias=args.attention_bias,
-        )
-        self.k_proj = nn.Linear(
-            args.hidden_size,
-            args.num_key_value_heads * self.head_dim,
-            bias=args.attention_bias,
-        )
-        self.v_proj = nn.Linear(
-            args.hidden_size,
-            args.num_key_value_heads * self.head_dim,
-            bias=args.attention_bias,
-        )
-        self.o_proj = nn.Linear(
-            args.num_attention_heads * self.head_dim, args.hidden_size, bias=False
-        )
-
-        self.rope = nn.RoPE(
-            dims=int(self.head_dim * args.partial_rotary_factor),
-            base=args.rope_theta,
-            traditional=args.rope_traditional,
-        )
-
-    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 Glm4DecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.self_attn = Glm4Attention(args=args)
-
-        self.mlp = Glm4MLP(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.post_self_attn_layernorm = nn.RMSNorm(
-            args.hidden_size, eps=args.rms_norm_eps
-        )
-        self.post_mlp_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:
-        x = x + self.post_self_attn_layernorm(
-            self.self_attn(self.input_layernorm(x), mask, cache)
-        )
-        residual = x
-        x = (
-            self.post_mlp_layernorm(self.mlp(self.post_attention_layernorm(x)))
-            + residual
-        )
-        return x
-
-
-class Glm4Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [
-            Glm4DecoderLayer(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,
-    ):
-        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 = Glm4Model(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

mlx_lm/models/gpt2.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any, Dict, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -133,15 +133,14 @@ class GPT2Model(nn.Module):
 
         hidden_states = self.wte(inputs)
 
-        offset = 0
-        if cache is not None and len(cache) > 0 and cache[0] is not None:
-            offset = cache[0].offset
+        mask = None
+        if hidden_states.shape[1] > 1:
 
-        position_ids = mx.arange(offset, offset + L)
-        hidden_states += self.wpe(position_ids)
+            position_ids = mx.array(np.arange(L))
+            hidden_states += self.wpe(position_ids)
 
-        if mask is None:
-            mask = create_attention_mask(hidden_states, cache)
+            if mask is None:
+                mask = create_attention_mask(hidden_states, cache)
 
         if cache is None:
             cache = [None] * len(self.h)

mlx_lm/models/gpt_bigcode.py

@@ -145,16 +145,16 @@ class GPTBigCodeModel(nn.Module):
         hidden_states = self.wte(inputs)
 
         mask = None
-        if mask is not None and hidden_states.shape[1] > 1:
-            mask = create_attention_mask(hidden_states, cache)
+        if hidden_states.shape[1] > 1:
+
+            position_ids = mx.array(np.arange(L))
+            hidden_states += self.wpe(position_ids)
+
+            if mask is None:
+                mask = create_attention_mask(hidden_states, cache)
 
         if cache is None:
             cache = [None] * len(self.h)
-            position_ids = mx.array(np.arange(L))
-        else:
-            position_ids = mx.array(np.arange(cache[0].offset, cache[0].offset + L))
-
-        hidden_states += self.wpe(position_ids)
 
         for layer, c in zip(self.h, cache):
             hidden_states = layer(hidden_states, mask, cache=c)

mlx_lm/models/granite.py

@@ -1,195 +0,0 @@
-# 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
-
-
-@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
-    mlp_bias: bool
-    rope_theta: float
-    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.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)
-
-        # 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, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_size
-        hidden_dim = args.intermediate_size
-        if hasattr(args, "mlp_bias"):
-            mlp_bias = args.mlp_bias
-        else:
-            mlp_bias = False
-
-        self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
-        self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias)
-        self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_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)
-        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.mlp(self.post_attention_layernorm(h))
-        out = h + r * self.residual_multiplier
-        return out
-
-
-class GraniteModel(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)
-        self.layers = [
-            TransformerBlock(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 = GraniteModel(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
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/helium.py

@@ -1,185 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from dataclasses import dataclass
-from typing import Any, Optional, Tuple
-
-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):
-    hidden_size: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    num_key_value_heads: int
-    rms_norm_eps: float
-    vocab_size: int
-    attention_bias: bool
-    head_dim: int
-    max_position_embeddings: int
-    mlp_bias: bool
-    model_type: str
-    rope_theta: float
-    tie_word_embeddings: bool
-
-
-class HeliumAttention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        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.hidden_size // n_heads
-        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.rope = nn.RoPE(head_dim, traditional=True, 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)
-
-        # 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 HeliumMLP(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.hidden_size = args.hidden_size
-        self.intermediate_size = args.intermediate_size
-
-        self.gate_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size, bias=args.mlp_bias
-        )
-        self.up_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size, bias=args.mlp_bias
-        )
-        self.down_proj = nn.Linear(
-            self.intermediate_size, self.hidden_size, bias=args.mlp_bias
-        )
-
-    def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class HeliumDecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.hidden_size = args.hidden_size
-
-        self.self_attn = HeliumAttention(args)
-        self.mlp = HeliumMLP(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.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 HeliumModel(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.num_hidden_layers = args.num_hidden_layers
-        self.vocab_size = args.vocab_size
-
-        assert self.vocab_size > 0
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-
-        self.layers = [HeliumDecoderLayer(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,
-    ) -> 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, 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 = HeliumModel(args)
-
-        self.vocab_size = args.vocab_size
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-
-        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,
-    ) -> mx.array:
-        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
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/hunyuan.py

@@ -76,6 +76,7 @@ class Attention(nn.Module):
 
         head_dim = args.hidden_size // n_heads
         self.scale = head_dim**-0.5
+
         self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
         if kv_proj:
             self.k_proj = nn.Linear(
@@ -106,6 +107,7 @@ class Attention(nn.Module):
         B, L, D = x.shape
 
         queries = self.q_proj(x)
+
         if kv_states is None:
             keys, values = self.k_proj(x), self.v_proj(x)
             kv_states = keys, values
@@ -196,10 +198,7 @@ class DecoderLayer(nn.Module):
         super().__init__()
         self.hidden_size = args.hidden_size
         self.self_attn = Attention(kv_proj, args)
-        if args.num_experts == 1:
-            self.mlp = MLP(args.hidden_size, args.intermediate_size)
-        else:
-            self.mlp = MoeBlock(args)
+        self.mlp = MoeBlock(args)
 
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
         self.post_attention_layernorm = nn.RMSNorm(
@@ -232,10 +231,7 @@ class HunYuanModel(nn.Module):
         assert self.vocab_size > 0
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
-            DecoderLayer(
-                args=args,
-                kv_proj=(not args.use_cla) or (i % args.cla_share_factor) == 0,
-            )
+            DecoderLayer(args=args, kv_proj=(i % args.cla_share_factor) == 0)
             for i in range(args.num_hidden_layers)
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
@@ -255,7 +251,7 @@ class HunYuanModel(nn.Module):
             cache = [None] * len(self.layers)
 
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            if (not self.args.use_cla) or i % self.args.cla_share_factor == 0:
+            if i % self.args.cla_share_factor == 0:
                 shared_kv_states = None
             h, shared_kv_states = layer(h, mask, c, shared_kv_states)
 
@@ -279,29 +275,6 @@ class Model(nn.Module):
         return self.model.embed_tokens.as_linear(out)
 
     def sanitize(self, weights):
-
-        if "model.layers.0.mlp.gate_and_up_proj.weight" in weights:
-            new_weights = {}
-            D = self.args.hidden_size
-            n_kv_heads = self.args.num_key_value_heads
-            n_kv_groups = self.args.num_attention_heads // n_kv_heads
-            head_dim = D // self.args.num_attention_heads
-            for k, v in weights.items():
-                if "qkv_proj" in k:
-                    v = v.reshape(n_kv_heads, n_kv_groups + 2, head_dim, -1)
-                    splits = v.split([n_kv_groups, n_kv_groups + 1], axis=1)
-                    for k_up, v_new in zip(["q_proj", "k_proj", "v_proj"], splits):
-                        k_new = k.replace("qkv_proj", k_up)
-                        new_weights[k_new] = mx.flatten(v_new, 0, 2)
-                elif "gate_and_up_proj" in k:
-                    splits = v.split(2, axis=0)
-                    for k_up, v_new in zip(["up_proj", "gate_proj"], splits):
-                        k_new = k.replace("gate_and_up_proj", k_up)
-                        new_weights[k_new] = v_new
-                else:
-                    new_weights[k] = v
-            weights = new_weights
-
         if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
             return weights
         for l in range(self.args.num_hidden_layers):

mlx_lm/models/internlm3.py

@@ -1,241 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, 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
-    hidden_size: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    rms_norm_eps: float
-    vocab_size: int
-    bias: bool = False
-    qkv_bias: bool = False
-    max_position_embeddings: int = 32768
-    num_key_value_heads: int = None
-    rope_theta: float = 10000
-    rope_traditional: bool = False
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-    tie_word_embeddings: bool = False
-
-    def __post_init__(self):
-        if self.num_key_value_heads is None:
-            self.num_key_value_heads = self.num_attention_heads
-
-        if self.rope_scaling:
-            required_keys = {"factor", "rope_type"}
-            if not all(key in self.rope_scaling for key in required_keys):
-                raise ValueError(f"rope_scaling must contain keys {required_keys}")
-
-            if self.rope_scaling["rope_type"] not in ["linear", "dynamic"]:
-                raise ValueError(
-                    "rope_scaling 'rope_type' currently only supports 'linear' or 'dynamic"
-                )
-
-
-class DynamicNTKScalingRoPE(nn.Module):
-    """Implements the rotary positional encoding with Dynamic NTK scaling."""
-
-    def __init__(
-        self,
-        dims: int,
-        max_position_embeddings: int = 2048,
-        traditional: bool = False,
-        base: float = 10000,
-        scale: float = 1.0,
-    ):
-        super().__init__()
-        self.max_position_embeddings = max_position_embeddings
-        self.original_base = base
-        self.dims = dims
-        self.traditional = traditional
-        self.scale = scale
-
-    def extra_repr(self):
-        return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, scaling_factor={self.scaling_factor}"
-
-    def __call__(self, x, offset: int = 0):
-        seq_len = x.shape[1] + offset
-        if seq_len > self.max_position_embeddings:
-            base = self.original_base * (
-                (self.scale * seq_len / self.max_position_embeddings) - (self.scale - 1)
-            ) ** (self.dims / (self.dims - 2))
-        else:
-            base = self.original_base
-
-        return mx.fast.rope(
-            x,
-            self.dims,
-            traditional=self.traditional,
-            base=base,
-            scale=self.scale,
-            offset=offset,
-        )
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_size
-        qkv_bias = args.qkv_bias
-        self.n_heads = n_heads = args.num_attention_heads
-        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
-        self.n_kv_groups = n_heads // args.num_key_value_heads
-
-        self.head_dim = head_dim = args.hidden_size // n_heads
-        self.scale = head_dim**-0.5
-
-        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=qkv_bias)
-        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
-        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
-        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=qkv_bias)
-
-        rope_scale = (
-            1 / args.rope_scaling["factor"]
-            if args.rope_scaling is not None
-            and args.rope_scaling["rope_type"] == "linear"
-            else 2.0
-        )
-
-        self.rope = DynamicNTKScalingRoPE(
-            head_dim,
-            max_position_embeddings=args.max_position_embeddings,
-            traditional=args.rope_traditional,
-            base=args.rope_theta,
-            scale=rope_scale,
-        )
-
-    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):
-        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.self_attn = Attention(args)
-        self.mlp = MLP(args.hidden_size, args.intermediate_size, args.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 InternLM2Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        assert args.vocab_size > 0
-        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 = InternLM2Model(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 sanitize(self, weights):
-        # Remove unused precomputed rotary freqs
-        return {k: v for k, v in weights.items() if "attention.rope.inv_freq" not in k}
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/kimi_vl.py

@@ -1,120 +0,0 @@
-# Copyright © 2024 Apple Inc.
-
-import math
-from dataclasses import dataclass
-from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .deepseek_v3 import DeepseekV3Model
-from .switch_layers import SwitchGLU
-
-
-@dataclass
-class TextArgs(BaseModelArgs):
-    vocab_size: int = 102400
-    hidden_size: int = 4096
-    intermediate_size: int = 11008
-    moe_intermediate_size: int = 1407
-    num_hidden_layers: int = 30
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 32
-    n_shared_experts: Optional[int] = None
-    n_routed_experts: Optional[int] = None
-    routed_scaling_factor: float = 1.0
-    kv_lora_rank: int = 512
-    q_lora_rank: int = 1536
-    qk_rope_head_dim: int = 64
-    v_head_dim: int = 128
-    qk_nope_head_dim: int = 128
-    topk_method: str = "noaux_tc"
-    scoring_func: str = "sigmoid"
-    norm_topk_prob: bool = True
-    n_group: Optional[int] = None
-    topk_group: Optional[int] = None
-    num_experts_per_tok: Optional[int] = None
-    moe_layer_freq: int = 1
-    first_k_dense_replace: int = 0
-    max_position_embeddings: int = 2048
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 10000.0
-    rope_scaling: Dict = None
-    attention_bias: bool = False
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    text_config: Union[TextArgs, dict]
-    model_type: str
-
-    def __post_init__(self):
-        self.text_config = TextArgs.from_dict(self.text_config)
-
-
-class LanguageModel(nn.Module):
-    def __init__(self, config: TextArgs):
-        super().__init__()
-        self.args = config
-        self.model = DeepseekV3Model(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)
-
-
-class Model(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.args = config
-        self.model_type = config.model_type
-        self.language_model = LanguageModel(config.text_config)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
-    ):
-        return self.language_model(inputs, cache, mask)
-
-    def sanitize(self, weights):
-        def keep(key):
-            return (
-                "vision_tower" not in key
-                and "rotary_emb" not in key
-                and "multi_modal_projector" not in key
-            )
-
-        weights = {k: v for k, v in weights.items() if keep(k)}
-        # Stack experts
-        for l in range(self.args.text_config.num_hidden_layers):
-            prefix = f"language_model.model.layers.{l}"
-            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.text_config.n_routed_experts)
-                        ]
-                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
-
-        return weights
-
-    @property
-    def layers(self):
-        return self.language_model.model.layers
-
-    def cast_predicate(self):
-        def predicate(k):
-            return "e_score_correction_bias" not in k
-
-        return predicate

mlx_lm/models/llama.py

@@ -69,14 +69,12 @@ class Attention(nn.Module):
         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)
+        queries = mx.unflatten(queries, -1, (self.n_heads, -1)).transpose(0, 2, 1, 3)
+        keys = mx.unflatten(keys, -1, (self.n_kv_heads, -1)).transpose(0, 2, 1, 3)
+        values = mx.unflatten(values, -1, (self.n_kv_heads, -1)).transpose(0, 2, 1, 3)
 
         if cache is not None:
             queries = self.rope(queries, offset=cache.offset)
@@ -90,7 +88,7 @@ class Attention(nn.Module):
             queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        output = output.transpose(0, 2, 1, 3).flatten(-2, -1)
         return self.o_proj(output)
 
 
@@ -196,12 +194,9 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         # Remove unused precomputed rotary freqs
-        weights = {
+        return {
             k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
         }
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-        return weights
 
     @property
     def layers(self):

mlx_lm/models/llama4.py

@@ -1,333 +0,0 @@
-# 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 ChunkedKVCache, KVCache
-from .rope_utils import initialize_rope
-from .switch_layers import SwitchGLU
-
-
-@dataclass
-class TextArgs(BaseModelArgs):
-    attention_bias: bool
-    attention_chunk_size: int
-    head_dim: int
-    hidden_act: str
-    hidden_size: int
-    interleave_moe_layer_step: int
-    intermediate_size: int
-    intermediate_size_mlp: int
-    max_position_embeddings: int
-    model_type: str
-    num_attention_heads: int
-    num_experts_per_tok: int
-    num_hidden_layers: int
-    num_key_value_heads: int
-    num_local_experts: int
-    rms_norm_eps: float
-    rope_scaling: Any
-    rope_theta: float
-    use_qk_norm: bool
-    vocab_size: int
-    attn_temperature_tuning: int = 4
-    floor_scale: int = 8192
-    attn_scale: float = 0.1
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    text_config: Union[TextArgs, dict]
-    model_type: str
-
-    def __post_init__(self):
-        self.text_config = TextArgs.from_dict(self.text_config)
-
-
-class Attention(nn.Module):
-    def __init__(self, args: TextArgs, layer_idx: int):
-        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.use_rope = int((layer_idx + 1) % 4 != 0)  # rope unused for dense layers
-        self.attn_temperature_tuning = args.attn_temperature_tuning
-        self.floor_scale = args.floor_scale
-        self.attn_scale = args.attn_scale
-
-        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
-
-        self.scale = head_dim**-0.5
-        if hasattr(args, "attention_bias"):
-            attention_bias = args.attention_bias
-        else:
-            attention_bias = False
-
-        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.use_qk_norm = args.use_qk_norm and self.use_rope
-
-        if self.use_rope:
-            self.rope = initialize_rope(
-                head_dim,
-                args.rope_theta,
-                traditional=True,
-                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 = 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:
-            offset = cache.offset
-        else:
-            offset = 0
-
-        if self.use_rope:
-            queries = self.rope(queries, offset=offset)
-            keys = self.rope(keys, offset=offset)
-
-        if self.use_qk_norm:
-            queries = mx.fast.rms_norm(queries, weight=None, eps=1e-6)
-            keys = mx.fast.rms_norm(keys, weight=None, eps=1e-6)
-
-        if self.attn_temperature_tuning and not self.use_rope:
-            attn_scales = (
-                mx.log(
-                    mx.floor(mx.arange(offset + 1, offset + L + 1) / self.floor_scale)
-                    + 1.0
-                )
-                * self.attn_scale
-                + 1.0
-            )
-            attn_scales = attn_scales[:, None]
-            queries = (queries * attn_scales).astype(queries.dtype)
-
-        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, intermediate_size: int = None):
-        super().__init__()
-
-        dim = args.hidden_size
-        hidden_dim = intermediate_size or 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 MoE(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.top_k = args.num_experts_per_tok
-        self.num_experts = args.num_local_experts
-        self.experts = SwitchGLU(
-            args.hidden_size, args.intermediate_size, self.num_experts
-        )
-        self.router = nn.Linear(args.hidden_size, args.num_local_experts, bias=False)
-        self.shared_expert = MLP(args)
-
-    def __call__(self, x) -> mx.array:
-        logits = self.router(x)
-        k = self.top_k
-        indices = mx.argpartition(-logits, kth=k - 1, axis=-1)[..., :k]
-        scores = mx.take_along_axis(logits, indices, axis=-1)
-        scores = mx.sigmoid(scores.astype(mx.float32)).astype(x.dtype)
-
-        out = self.experts(x * scores, indices).squeeze(2)
-        return out + self.shared_expert(x)
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: TextArgs, layer_idx: int):
-        super().__init__()
-        self.num_attention_heads = args.num_attention_heads
-        self.hidden_size = args.hidden_size
-        self.self_attn = Attention(args, layer_idx)
-        self.is_moe_layer = (layer_idx % args.interleave_moe_layer_step) == (
-            args.interleave_moe_layer_step - 1
-        )
-        if self.is_moe_layer:
-            self.feed_forward = MoE(args)
-        else:
-            self.feed_forward = MLP(args, args.intermediate_size_mlp)
-
-        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.feed_forward(self.post_attention_layernorm(h))
-        out = h + r
-        return out
-
-
-class LlamaModel(nn.Module):
-    def __init__(self, args: TextArgs):
-        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)
-        self.layers = [TransformerBlock(args, i) for i in range(args.num_hidden_layers)]
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.attention_chunk_size = args.attention_chunk_size
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        mask: mx.array = None,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs)
-
-        if cache is not None:
-            for idx, c in enumerate(cache):
-                if (idx + 1) % 4 != 0:
-                    c.maybe_trim_front()
-            start = cache[0].start_position
-            offset = cache[0].offset
-        else:
-            start = 0
-            offset = 0
-        end = offset + h.shape[1]
-        linds = mx.arange(start, end)
-        rinds = mx.arange(offset, end)[:, None]
-        block_pos = mx.abs(
-            (linds // self.attention_chunk_size) - (rinds // self.attention_chunk_size)
-        )
-        token_pos = linds <= rinds
-        chunk_mask = (block_pos == 0) & token_pos
-
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-        else:
-            chunk_mask &= mask
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        for idx, (layer, c) in enumerate(zip(self.layers, cache)):
-            use_chunked_attention = (idx + 1) % 4 != 0
-            if use_chunked_attention:
-                local_mask = chunk_mask
-            else:
-                local_mask = mask
-            h = layer(h, local_mask, cache=c)
-
-        return self.norm(h)
-
-
-class LanguageModel(nn.Module):
-    def __init__(self, args: TextArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = LlamaModel(self.args)
-        self.lm_head = nn.Linear(
-            self.args.hidden_size, self.args.vocab_size, bias=False
-        )
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        mask: mx.array = None,
-        cache=None,
-    ):
-        out = self.model(inputs, mask, cache)
-        return self.lm_head(out)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.language_model = LanguageModel(args.text_config)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        mask: mx.array = None,
-        cache=None,
-    ):
-        return self.language_model(inputs, mask, cache)
-
-    def sanitize(self, weights):
-        def to_remove(k):
-            return "vision_model" in k or "multi_modal_projector" in k
-
-        # Remove vision weights
-        weights = {k: v for k, v in weights.items() if not to_remove(k)}
-
-        # Rename expert weights for SwitchGLU
-        for l in range(self.args.text_config.num_hidden_layers):
-            prefix = f"language_model.model.layers.{l}.feed_forward.experts"
-            if f"{prefix}.gate_up_proj" in weights:
-                v = weights.pop(f"{prefix}.gate_up_proj")
-                gate_k = f"{prefix}.gate_proj.weight"
-                up_k = f"{prefix}.up_proj.weight"
-                gate_proj, up_proj = mx.split(v, 2, axis=-1)
-                weights[gate_k] = mx.swapaxes(gate_proj, 1, 2)
-                weights[up_k] = mx.swapaxes(up_proj, 1, 2)
-            if f"{prefix}.down_proj" in weights:
-                down_proj = weights.pop(f"{prefix}.down_proj")
-                weights[f"{prefix}.down_proj.weight"] = mx.swapaxes(down_proj, 1, 2)
-        return weights
-
-    @property
-    def layers(self):
-        return self.language_model.model.layers
-
-    def make_cache(self):
-        chunk_size = self.args.text_config.attention_chunk_size
-        caches = []
-        for i in range(len(self.layers)):
-            if (i + 1) % 4 != 0:
-                caches.append(ChunkedKVCache(chunk_size))
-            else:
-                caches.append(KVCache())
-        return caches

mlx_lm/models/mamba.py

@@ -1,4 +1,4 @@
-# Copyright © 2024-2025 Apple Inc.
+# Copyright © 2024 Apple Inc.
 
 import math
 from dataclasses import dataclass
@@ -123,16 +123,17 @@ class MambaBlock(nn.Module):
             self.intermediate_size, self.hidden_size, bias=args.use_bias
         )
 
-    def ssm_step(self, x, A, state=None):
+    def ssm_step(self, x, state=None):
+        A = -mx.exp(self.A_log)
         D = self.D
         deltaBC = self.x_proj(x)
-        delta, B, C = map(
-            self.mixer_norm if self.use_bcdt_rms else lambda x: x,
-            mx.split(
-                deltaBC,
-                [self.time_step_rank, self.time_step_rank + self.ssm_state_size],
-                axis=-1,
-            ),
+        delta, B, C = mx.split(
+            deltaBC,
+            indices_or_sections=[
+                self.time_step_rank,
+                self.time_step_rank + self.ssm_state_size,
+            ],
+            axis=-1,
         )
         if self.use_bcdt_rms:
             delta, B, C = map(self.mixer_norm, (delta, B, C))
@@ -144,40 +145,25 @@ class MambaBlock(nn.Module):
         y = y + D * x
         return y, new_state
 
-    def _process_sequence(self, x, conv_cache, state_cache):
+    def __call__(self, x, cache):
         B, T, D = x.shape
-        xz = self.in_proj(x)
-        x, z = xz.split(indices_or_sections=2, axis=-1)
-
-        conv_out, new_conv_cache = self.conv1d(x, conv_cache)
-        x = nn.silu(conv_out)
-
-        A = -mx.exp(self.A_log)
+        if cache is None:
+            cache = [None, None]
 
         outputs = []
-        current_state = state_cache
-        y = []
         for t in range(T):
-            y_t, current_state = self.ssm_step(x[:, t], A, current_state)
-            y.append(y_t)
-        y = mx.stack(y, axis=1)
-        z = self.out_proj(nn.silu(z) * y)
-        return z, (new_conv_cache, current_state)
-
-    def __call__(self, x, cache):
-        if cache is None:
-            conv_cache, state_cache = None, None
-        else:
-            conv_cache, state_cache = cache[0], cache[1]
-
-        output, (new_conv_cache, new_state_cache) = self._process_sequence(
-            x, conv_cache, state_cache
-        )
-
-        if isinstance(cache, MambaCache):
-            cache[0] = new_conv_cache
-            cache[1] = new_state_cache
-
+            xt = x[:, t, :]
+            xz = self.in_proj(xt)
+            x_t, z_t = xz.split(indices_or_sections=2, axis=1)
+            conv_out, cache[0] = self.conv1d(mx.expand_dims(x_t, 1), cache[0])
+            x_t = conv_out.squeeze(1)
+            x_t = nn.silu(x_t)
+            y_t, cache[1] = self.ssm_step(x_t, cache[1])
+            z_t = nn.silu(z_t)
+            output_t = y_t * z_t
+            output_t = self.out_proj(output_t)
+            outputs.append(output_t)
+        output = mx.stack(outputs, axis=1)
         return output
 
 

mlx_lm/models/mimo.py

@@ -1,194 +0,0 @@
-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
-    max_position_embeddings: int = 32768
-    rope_theta: float = 10000.0
-    rope_traditional: bool = False
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-    tie_word_embeddings: bool = False
-    num_nextn_predict_layers: int = 2
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        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.hidden_size // n_heads
-        self.scale = head_dim**-0.5
-
-        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True)
-        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
-        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
-        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
-
-        self.rope = initialize_rope(
-            head_dim,
-            base=args.rope_theta,
-            traditional=args.rope_traditional,
-            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 = 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):
-        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):
-        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)
-        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 MiMoModel(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.num_nextn_predict_layers = args.num_nextn_predict_layers
-
-        assert self.vocab_size > 0
-        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, c)
-
-        h = self.norm(h)
-
-        return h
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = MiMoModel(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 sanitize(self, weights):
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-
-        return {
-            k: v
-            for k, v in weights.items()
-            if "self_attn.rotary_emb.inv_freq" not in k
-            and not k.startswith("model.mtp_layers.")
-        }
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/minicpm.py

@@ -1,10 +1,11 @@
-# Copyright © 2023-2025 Apple Inc.
+# Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+import numpy as np
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
@@ -137,9 +138,9 @@ class DecoderLayer(nn.Module):
         cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
-        h = x + r * (self.scale_depth / self.num_hidden_layers**0.5)
+        h = x + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
         r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r * (self.scale_depth / self.num_hidden_layers**0.5)
+        out = h + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
         return out
 
 

mlx_lm/models/minicpm3.py

@@ -1,250 +0,0 @@
-# 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
-from .su_rope import SuScaledRotaryEmbedding
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    hidden_size: int
-    dim_model_base: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    rms_norm_eps: float
-    vocab_size: int
-    num_key_value_heads: int
-    q_lora_rank: int
-    qk_nope_head_dim: int
-    qk_rope_head_dim: int
-    kv_lora_rank: int
-    scale_depth: float
-    scale_emb: float
-    max_position_embeddings: int
-    attention_bias: bool = False
-    rope_theta: float = 1000000.0
-    rope_traditional: bool = False
-    rope_scaling: Optional[Dict[str, Union[str, float]]] = None
-    tie_word_embeddings: bool = False
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-
-        self.qk_rope_head_dim = self.args.qk_rope_head_dim
-        self.qk_nope_head_dim = self.args.qk_nope_head_dim
-        self.attention_bias = self.args.attention_bias
-        self.kv_lora_rank = self.args.kv_lora_rank
-        self.num_heads = self.args.num_attention_heads
-        self.q_lora_rank = self.args.q_lora_rank
-        self.hidden_size = self.args.hidden_size
-
-        self.v_head_dim = self.hidden_size // self.args.num_attention_heads
-        self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
-        self.softmax_scale = self.q_head_dim ** (-0.5)
-
-        self.q_a_proj = nn.Linear(
-            self.hidden_size, self.q_lora_rank, bias=self.attention_bias
-        )
-        self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
-
-        self.q_b_proj = nn.Linear(
-            self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
-        )
-
-        self.kv_a_proj_with_mqa = nn.Linear(
-            self.hidden_size,
-            self.kv_lora_rank + self.qk_rope_head_dim,
-            bias=self.attention_bias,
-        )
-
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
-
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
-        )
-
-        self.o_proj = nn.Linear(
-            self.num_heads * self.v_head_dim,
-            self.hidden_size,
-            bias=self.attention_bias,
-        )
-
-        self.rope = SuScaledRotaryEmbedding(
-            dims=args.qk_rope_head_dim,
-            base=args.rope_theta,
-            max_position_embeddings=args.max_position_embeddings,
-            original_max_position_embeddings=args.rope_scaling.get(
-                "original_max_position_embeddings", 4096
-            ),
-            short_factor=args.rope_scaling.get("short_factor", 1.0),
-            long_factor=args.rope_scaling.get("long_factor", 1.0),
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Dict[str, mx.array]] = None,
-    ):
-        B, L, _ = x.shape
-
-        # Project query
-        q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
-        q = q.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
-        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
-
-        # Project key and value
-        compressed_kv = self.kv_a_proj_with_mqa(x)
-        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
-
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
-
-        # Apply RoPE to the query and key parts that need position embedding
-        if cache is not None:
-            q_pe = self.rope(q_pe, offset=cache.offset)
-            k_pe = self.rope(k_pe, offset=cache.offset)
-        else:
-            q_pe = self.rope(q_pe)
-            k_pe = self.rope(k_pe)
-
-        # Create the full query and key tensors by combining the parts
-        # Broadcast k_pe to all heads
-        k_pe_broadcasted = mx.broadcast_to(
-            k_pe, (B, self.num_heads, L, self.qk_rope_head_dim)
-        )
-
-        # Use concatenate for queries
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
-
-        # Use concatenate for keys
-        keys = mx.concatenate([k_nope, k_pe_broadcasted], axis=-1)
-
-        # Update cache if needed
-        if cache is not None:
-            keys, values = cache.update_and_fetch(keys, values)
-
-        # Perform attention
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.softmax_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):
-        super().__init__()
-        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
-
-    def __call__(self, x):
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class DecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.hidden_size = args.hidden_size
-        self.num_hidden_layers = args.num_hidden_layers
-
-        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.scale_depth = args.scale_depth
-        self.num_hidden_layers = args.num_hidden_layers
-
-    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.scale_depth / (self.num_hidden_layers**0.5))
-        r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r * (self.scale_depth / (self.num_hidden_layers**0.5))
-        return out
-
-
-class MiniCPM3Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.vocab_size = args.vocab_size
-        assert self.vocab_size > 0
-
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [DecoderLayer(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) * self.args.scale_emb
-
-        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 = MiniCPM3Model(args)
-
-        if not self.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 not self.args.tie_word_embeddings:
-            out = self.lm_head(out / (self.args.hidden_size / self.args.dim_model_base))
-        else:
-            out = self.model.embed_tokens.as_linear(out)
-
-        return out
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/nemotron-nas.py

@@ -1,386 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from dataclasses import dataclass, field
-from dataclasses import fields as dataclass_fields
-from typing import Any, Dict, List, 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(frozen=True)
-class AttentionConfig:
-    no_op: bool = False
-    replace_with_linear: bool = False
-    sparsify: Optional[list[str]] = None
-    n_heads_in_group: Optional[int] = None  # GQA group size
-    window_length: Optional[int] = None  # Not directly used here, placeholder
-    num_sink_tokens: Optional[int] = None  # Not directly used here, placeholder
-    use_prefill_window_in_sink_attention: bool = (
-        False  # Not directly used here, placeholder
-    )
-    unshifted_sink: bool = False  # Not directly used here, placeholder
-
-    def __post_init__(self):
-        # Ensure consistency: If no-op or linear, other attn params are irrelevant
-        if self.no_op or self.replace_with_linear:
-            # Use object.__setattr__ because the dataclass is frozen
-            object.__setattr__(self, "n_heads_in_group", None)
-            object.__setattr__(self, "window_length", None)
-            object.__setattr__(self, "num_sink_tokens", None)
-        # If it's a standard attention block, n_heads_in_group must be provided
-        elif not self.no_op:
-            if self.n_heads_in_group is None:
-                raise ValueError(
-                    "n_heads_in_group must be specified for active attention blocks"
-                )
-            if self.n_heads_in_group <= 0:
-                raise ValueError(
-                    f"n_heads_in_group must be positive, got {self.n_heads_in_group}"
-                )
-
-
-@dataclass(frozen=True)
-class FFNConfig:
-    no_op: bool = False
-    replace_with_linear: bool = False
-    sparsify: Optional[list[str]] = None
-    ffn_mult: Optional[float] = None
-
-    def __post_init__(self):
-        # Ensure consistency: If no-op or linear, ffn_mult is irrelevant
-        if self.no_op or self.replace_with_linear:
-            object.__setattr__(self, "ffn_mult", None)
-        # If it's a standard FFN block, ffn_mult must be provided
-        elif not self.no_op:
-            if self.ffn_mult is None:
-                raise ValueError("ffn_mult must be specified for active FFN blocks")
-            # Round to prevent potential floating point inconsistencies if needed
-            object.__setattr__(self, "ffn_mult", round(self.ffn_mult, 6))
-
-
-@dataclass(frozen=True)
-class BlockConfig:
-    attention: AttentionConfig
-    ffn: FFNConfig
-
-    @classmethod
-    def from_dict(cls, data: dict):
-        # Helper to create BlockConfig from a dictionary (e.g., loaded from JSON)
-        attn_conf = AttentionConfig(**data.get("attention", {}))
-        ffn_conf = FFNConfig(**data.get("ffn", {}))
-        return cls(attention=attn_conf, ffn=ffn_conf)
-
-
-def _find_multiple(n: int, k: int) -> int:
-    """Finds the smallest multiple of k greater than or equal to n."""
-    if n % k == 0:
-        return n
-    return n + k - (n % k)
-
-
-def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int:
-    """Calculates intermediate size based on multiplier, rounding up to multiple of 256."""
-    intermediate_size = int(2 * ffn_mult * n_embd / 3)
-    return _find_multiple(intermediate_size, 256)
-
-
-# Activation function mapping
-_ACT2FN = {
-    "silu": nn.silu,
-    "relu": nn.relu,
-    "gelu": nn.gelu,
-    "gelu_new": nn.gelu_approx,
-    "gelu_fast": nn.gelu_approx,
-}
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "nemotron-nas"
-    hidden_size: int = 8192
-    num_hidden_layers: int = 80
-    num_attention_heads: int = 64
-    rms_norm_eps: float = 1e-5
-    vocab_size: int = 128256
-    block_configs: list = field(default_factory=list)  # List of BlockConfig or dicts
-    hidden_act: str = "silu"
-    attention_bias: bool = False
-    mlp_bias: bool = False
-    rope_theta: float = 500000.0
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-    max_position_embeddings: int = 131072
-    tie_word_embeddings: bool = False
-
-    def __post_init__(self):
-        # Automatically parse block_configs if they are loaded as dicts
-        if self.block_configs and isinstance(self.block_configs[0], dict):
-            self.block_configs = [
-                BlockConfig.from_dict(conf) for conf in self.block_configs
-            ]
-
-        if len(self.block_configs) != self.num_hidden_layers:
-            raise ValueError(
-                f"Number of block_configs ({len(self.block_configs)}) must match "
-                f"num_hidden_layers ({self.num_hidden_layers})"
-            )
-
-        # Basic validation for RoPE scaling if provided
-        if self.rope_scaling:
-            if "factor" not in self.rope_scaling:
-                raise ValueError("rope_scaling must contain 'factor'")
-            rope_type = self.rope_scaling.get("rope_type")
-            if rope_type is None:
-                raise ValueError("rope_scaling must contain 'rope_type'")
-
-        # Validate individual block configs (post_init in dataclasses already does some)
-        for i, block_conf in enumerate(self.block_configs):
-            attn_conf = block_conf.attention
-            if not attn_conf.no_op and not attn_conf.replace_with_linear:
-                if self.num_attention_heads % attn_conf.n_heads_in_group != 0:
-                    raise ValueError(
-                        f"Layer {i}: num_attention_heads ({self.num_attention_heads}) "
-                        f"must be divisible by n_heads_in_group ({attn_conf.n_heads_in_group})"
-                    )
-
-
-class Attention(nn.Module):
-    """Standard GQA Attention mechanism for layers that use it."""
-
-    def __init__(self, args: ModelArgs, attention_config: AttentionConfig):
-        super().__init__()
-
-        dim = args.hidden_size
-        self.n_heads = n_heads = args.num_attention_heads
-        self.n_kv_heads = n_kv_heads = n_heads // attention_config.n_heads_in_group
-
-        self.head_dim = head_dim = args.hidden_size // n_heads
-        if (self.head_dim * n_heads) != dim:
-            raise ValueError(
-                f"hidden_size ({dim}) must be divisible by num_attention_heads ({n_heads})"
-            )
-
-        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)
-
-        # Initialize RoPE based on global config
-        self.rope = initialize_rope(
-            self.head_dim,
-            args.rope_theta,
-            False,  # Llama uses traditional=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, 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)
-            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):
-    """Standard Feed-Forward Network for layers that use it."""
-
-    def __init__(self, args: ModelArgs, ffn_config: FFNConfig):
-        super().__init__()
-
-        dim = args.hidden_size
-        # Calculate intermediate dim based on layer's specific config
-        hidden_dim = _ffn_mult_to_intermediate_size(ffn_config.ffn_mult, dim)
-
-        self.gate_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
-        self.down_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
-        self.up_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
-
-        try:
-            self.act_fn = _ACT2FN[args.hidden_act]
-        except KeyError:
-            raise ValueError(f"Unknown activation function: {args.hidden_act}")
-
-    def __call__(self, x) -> mx.array:
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-
-
-class LinearSubblockReplacement(nn.Module):
-    """A simple linear layer used to replace Attention or MLP blocks."""
-
-    def __init__(self, hidden_size: int, bias: bool):
-        super().__init__()
-        self.linear = nn.Linear(hidden_size, hidden_size, bias=bias)
-
-    def __call__(self, x: mx.array, *args, **kwargs) -> mx.array:
-        # Accepts potential extra args (like mask, cache) but ignores them
-        return self.linear(x)
-
-
-class TransformerBlock(nn.Module):
-    """A single transformer block, potentially heterogeneous based on config."""
-
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.hidden_size = args.hidden_size
-        # Get the specific configuration for this layer
-        block_config = args.block_configs[layer_idx]
-        self.attention_config = block_config.attention
-        self.ffn_config = block_config.ffn
-
-        # Conditionally initialize Input LayerNorm (needed unless Attention is no-op)
-        if not self.attention_config.no_op:
-            self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        else:
-            self.input_layernorm = None
-
-        # Conditionally initialize Attention block
-        if self.attention_config.no_op:
-            self.self_attn = None
-        elif self.attention_config.replace_with_linear:
-            self.self_attn = LinearSubblockReplacement(
-                args.hidden_size, args.attention_bias
-            )
-        else:
-            # Standard attention for this layer
-            self.self_attn = Attention(args, self.attention_config)
-
-        # Conditionally initialize Post-Attention LayerNorm (needed unless FFN is no-op)
-        if not self.ffn_config.no_op:
-            self.post_attention_layernorm = nn.RMSNorm(
-                args.hidden_size, eps=args.rms_norm_eps
-            )
-        else:
-            self.post_attention_layernorm = None
-
-        # Conditionally initialize MLP block
-        if self.ffn_config.no_op:
-            self.mlp = None
-        elif self.ffn_config.replace_with_linear:
-            self.mlp = LinearSubblockReplacement(args.hidden_size, args.mlp_bias)
-        else:
-            # Standard MLP for this layer
-            self.mlp = MLP(args, self.ffn_config)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-
-        # Attention part (Input Norm -> Attention -> Residual)
-        if self.self_attn is not None:
-            residual = x
-            h = self.input_layernorm(x)
-            attn_out = self.self_attn(h, mask=mask, cache=cache)
-            x = residual + attn_out
-
-        # MLP part (Post-Attention Norm -> MLP -> Residual)
-        if self.mlp is not None:
-            residual = x
-            h = self.post_attention_layernorm(x)
-            mlp_out = self.mlp(h)
-            x = residual + mlp_out
-
-        return x
-
-
-class NemotronNASModel(nn.Module):
-    """The core Nemotron-NAS style transformer model."""
-
-    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=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[List[Any]] = 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 i, layer in enumerate(self.layers):
-            h = layer(h, mask, cache=cache[i])
-
-        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 = NemotronNASModel(args)
-        if not args.tie_word_embeddings:
-            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-        else:
-            self.lm_head = None
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        mask=None,
-        cache=None,
-    ):
-        out = self.model(inputs, mask=mask, cache=cache)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    def sanitize(self, weights):
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/olmoe.py

@@ -1,217 +0,0 @@
-# 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
-    num_experts: int
-    num_experts_per_tok: int
-    norm_topk_prob: bool = False
-    head_dim: Optional[int] = None
-    max_position_embeddings: Optional[int] = None
-    num_key_value_heads: Optional[int] = None
-    attention_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
-
-    def __post_init__(self):
-        if self.num_key_value_heads is None:
-            self.num_key_value_heads = self.num_attention_heads
-
-
-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 or args.hidden_size // n_heads
-
-        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.rope = initialize_rope(
-            self.head_dim,
-            args.rope_theta,
-            args.rope_traditional,
-            args.rope_scaling,
-            args.max_position_embeddings,
-        )
-
-        self.q_norm = nn.RMSNorm(n_heads * head_dim, args.rms_norm_eps)
-        self.k_norm = nn.RMSNorm(n_kv_heads * head_dim, args.rms_norm_eps)
-
-    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)
-        keys = self.k_norm(keys)
-        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 OlmoeSparseMoeBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.num_experts = args.num_experts
-        self.top_k = args.num_experts_per_tok
-        self.norm_topk_prob = args.norm_topk_prob
-
-        self.gate = nn.Linear(args.hidden_size, self.num_experts, bias=False)
-        self.switch_mlp = SwitchGLU(
-            args.hidden_size,
-            args.intermediate_size,
-            self.num_experts,
-            bias=args.mlp_bias,
-        )
-
-    def __call__(self, x: mx.array) -> mx.array:
-        B, L, D = x.shape
-        x_flat = x.reshape(-1, D)
-        router_logits = self.gate(x_flat)
-        routing_weights = mx.softmax(router_logits, axis=1, precise=True)
-        k = self.top_k
-        indices = mx.stop_gradient(
-            mx.argpartition(-routing_weights, kth=k - 1, axis=-1)[..., :k]
-        )
-        scores = mx.take_along_axis(routing_weights, indices, axis=-1)
-        if self.norm_topk_prob:
-            scores = scores / scores.sum(axis=-1, keepdims=True)
-        y = self.switch_mlp(x_flat, indices)
-        y = (y * scores[..., None]).sum(axis=-2)
-        return y.reshape(B, L, D)
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.self_attn = Attention(args)
-        self.mlp = OlmoeSparseMoeBlock(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:
-        x = x + self.self_attn(self.input_layernorm(x), mask, cache)
-        x = x + self.mlp(self.post_attention_layernorm(x))
-        return x
-
-
-class OlmoeModel(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)
-        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,
-        cache=None,
-        mask=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 = OlmoeModel(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,
-        cache=None,
-        mask=None,
-    ):
-        out = self.model(inputs, cache, mask)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    def sanitize(self, weights):
-        if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
-            return weights
-        for l in range(self.args.num_hidden_layers):
-            prefix = f"model.layers.{l}"
-            for n in ["up_proj", "down_proj", "gate_proj"]:
-                for k in ["weight", "scales", "biases"]:
-                    if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
-                        to_join = [
-                            weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
-                            for e in range(self.args.num_experts)
-                        ]
-                        weights[f"{prefix}.mlp.switch_mlp.{n}.{k}"] = mx.stack(to_join)
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/phi3.py

@@ -23,10 +23,8 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float = 10000
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, List[float]]]] = None
-    partial_rotary_factor: float = 1.0
     max_position_embeddings: int = 131072
     original_max_position_embeddings: int = 4096
-    tie_word_embeddings: bool = False
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
@@ -61,10 +59,9 @@ class Attention(nn.Module):
         self.qkv_proj = nn.Linear(dim, op_size, bias=False)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
 
-        rope_dim = int(head_dim * args.partial_rotary_factor)
         if args.rope_scaling and args.rope_scaling["type"] in ["longrope", "su"]:
             self.rope = SuScaledRotaryEmbedding(
-                rope_dim,
+                head_dim,
                 base=args.rope_theta,
                 max_position_embeddings=args.max_position_embeddings,
                 original_max_position_embeddings=args.original_max_position_embeddings,
@@ -77,7 +74,7 @@ class Attention(nn.Module):
                 assert isinstance(args.rope_scaling["factor"], float)
                 rope_scale = 1 / args.rope_scaling["factor"]
             self.rope = nn.RoPE(
-                rope_dim,
+                head_dim,
                 traditional=args.rope_traditional,
                 base=args.rope_theta,
                 scale=rope_scale,
@@ -193,8 +190,7 @@ class Model(nn.Module):
         super().__init__()
         self.model_type = args.model_type
         self.model = Phi3Model(args)
-        if not args.tie_word_embeddings:
-            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
         self.args = args
 
     def __call__(
@@ -204,11 +200,7 @@ class Model(nn.Module):
         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
+        return self.lm_head(out)
 
     @property
     def layers(self):

mlx_lm/models/phi3small.py

@@ -266,7 +266,7 @@ class Phi3Model(nn.Module):
             h = self.mup_embedding_multiplier * h
 
         if mask is None:
-            mask = create_attention_mask(h, cache, return_array=True)
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)

mlx_lm/models/plamo2.py

@@ -1,599 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import math
-from dataclasses import dataclass
-from typing import Any, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from mlx_lm.models.base import BaseModelArgs, create_attention_mask
-
-from .cache import KVCache, MambaCache
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "plamo2"
-    hidden_size: int = 4096
-    num_hidden_layers: int = 32
-    rms_norm_eps: float = 1e-6
-    tie_word_embeddings: bool = True
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 4
-    hidden_size_per_head: int = 128
-    max_position_embeddings: int = 2048
-    attention_window_size: int = 2048
-    full_attention_idx: Optional[list[int]] = None
-    mamba_d_state: int = 64
-    mamba_d_conv: int = 4
-    mamba_num_heads: int = 64
-    mamba_step: int = 2
-    mamba_chunk_size: int = 256
-    mamba_enabled: bool = True
-    intermediate_size: int = 13312
-    vocab_size: int = 32000
-
-
-class RMSNorm(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        eps: float = 1e-6,
-        offset: float = 1.0,
-    ) -> None:
-        super().__init__()
-        self.weight = mx.zeros(hidden_size)
-        self.variance_epsilon = eps
-        self.offset = offset
-
-    def __call__(self, hidden_states: mx.array) -> mx.array:
-        return mx.fast.rms_norm(
-            hidden_states, self.weight + self.offset, self.variance_epsilon
-        )
-
-
-def get_initial_dt_bias(num_heads: int) -> mx.array:
-    dt_min = 0.001
-    dt_max = 0.1
-    dt = mx.exp(
-        mx.random.uniform(shape=(num_heads,)) * (math.log(dt_max) - math.log(dt_min))
-        + math.log(dt_min)
-    )
-    dt = mx.clip(dt, a_min=1e-4, a_max=None)
-    inv_dt = dt + mx.log(-mx.expm1(-dt))
-    return inv_dt
-
-
-def get_initial_A(num_heads: int) -> mx.array:
-    A = mx.arange(1, num_heads + 1, dtype=mx.float32)
-    return mx.log(A)
-
-
-# From: https://github.com/state-spaces/mamba/blob/0cce0fa645f100f00620ddf2333c2b7712abfdec/mamba_ssm/ops/triton/selective_state_update.py#L219
-def selective_state_update_ref(
-    state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
-) -> tuple[mx.array, mx.array]:
-    """
-    Argument:
-        state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
-        x: (batch, dim) or (batch, nheads, dim)
-        dt: (batch, dim) or (batch, nheads, dim)
-        A: (dim, dstate) or (nheads, dim, dstate)
-        B: (batch, dstate) or (batch, ngroups, dstate)
-        C: (batch, dstate) or (batch, ngroups, dstate)
-        D: (dim,) or (nheads, dim)
-        z: (batch, dim) or (batch, nheads, dim)
-        dt_bias: (dim,) or (nheads, dim)
-    Return:
-        out: (batch, dim) or (batch, nheads, dim)
-    """
-    has_heads = state.ndim > 3
-    if state.ndim == 3:
-        state = mx.expand_dims(state, 1)
-    if x.ndim == 2:
-        x = mx.expand_dims(x, 1)
-    if dt.ndim == 2:
-        dt = mx.expand_dims(dt, 1)
-    if A.ndim == 2:
-        A = mx.expand_dims(A, 0)
-    if B.ndim == 2:
-        B = mx.expand_dims(B, 1)
-    if C.ndim == 2:
-        C = mx.expand_dims(C, 1)
-    if D is not None and D.ndim == 1:
-        D = mx.expand_dims(D, 0)
-    if z is not None and z.ndim == 2:
-        z = mx.expand_dims(z, 1)
-    if dt_bias is not None and dt_bias.ndim == 1:
-        dt_bias = mx.expand_dims(dt_bias, 0)
-    batch, nheads, dim, dstate = state.shape
-    assert x.shape == (batch, nheads, dim)
-    assert dt.shape == x.shape
-    assert A.shape == (nheads, dim, dstate)
-    ngroups = B.shape[1]
-    assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
-    assert B.shape == (batch, ngroups, dstate)
-    assert C.shape == B.shape
-    if D is not None:
-        assert D.shape == (nheads, dim)
-    if z is not None:
-        assert z.shape == x.shape
-    if dt_bias is not None:
-        assert dt_bias.shape == (nheads, dim)
-        dt = dt + dt_bias
-    dt = nn.softplus(dt) if dt_softplus else dt
-    dA = mx.exp(mx.expand_dims(dt, axis=-1) * A)  # (batch, nheads, dim, dstate)
-    B = mx.reshape(
-        mx.repeat(mx.expand_dims(B, axis=2), nheads // ngroups, 2),
-        (batch, nheads, dstate),
-    )  # (batch, nheads, dstate)
-    C = mx.reshape(
-        mx.repeat(mx.expand_dims(C, axis=2), nheads // ngroups, 2),
-        (batch, nheads, dstate),
-    )  # (batch, nheads, dstate)
-    dB = mx.expand_dims(dt, axis=-1) * mx.expand_dims(
-        B, axis=-2
-    )  # (batch, nheads, dim, dstate)
-    state = state * dA + dB * mx.expand_dims(x, axis=-1)  # (batch, dim, dstate)
-    out = mx.einsum("bhdn,bhn->bhd", state.astype(C.dtype), C)
-    if D is not None:
-        out += (x * D).astype(out.dtype)
-    out = (out if z is None else out * nn.silu(z)).astype(x.dtype)
-    if not has_heads:
-        out = out.squeeze(1)
-    return out, state
-
-
-def ssd_update_state(
-    ssm_state: mx.array,
-    x: mx.array,
-    dt: mx.array,
-    A: mx.array,
-    B: mx.array,
-    C: mx.array,
-    D: mx.array,
-    z: mx.array,
-    dt_bias: mx.array,
-    dt_softplus: bool,
-) -> tuple[mx.array, mx.array]:
-    assert ssm_state.dtype == mx.float32
-    dtype = x.dtype
-
-    hidden_size_per_head = x.shape[-1]
-    d_state = B.shape[-1]
-    A = mx.broadcast_to(
-        A[:, None, None], (A.shape[0], hidden_size_per_head, d_state)
-    ).astype(mx.float32)
-    dt = mx.broadcast_to(
-        dt[..., None], (dt.shape[0], dt.shape[1], hidden_size_per_head)
-    )
-    dt_bias = mx.broadcast_to(
-        dt_bias[:, None], (dt_bias.shape[0], hidden_size_per_head)
-    )
-    D = mx.broadcast_to(D[:, None], (D.shape[0], hidden_size_per_head))
-    out, ssm_state = selective_state_update_ref(
-        ssm_state,
-        x.astype(dtype),
-        dt.astype(dtype),
-        A.astype(mx.float32),
-        B.astype(dtype),
-        C.astype(dtype),
-        D.astype(mx.float32),
-        z.astype(dtype),
-        dt_bias.astype(mx.float32),
-        dt_softplus=dt_softplus,
-    )
-    return out[:, None], ssm_state
-
-
-def ssd_chunk_scan_combined(
-    x: mx.array,
-    dt: mx.array,
-    A: mx.array,
-    B: mx.array,
-    C: mx.array,
-    D: mx.array,
-    z: mx.array,
-    dt_bias: mx.array,
-    dt_softplus: bool,
-    ssm_state: mx.array,
-) -> tuple[mx.array, mx.array]:
-    assert ssm_state.dtype == mx.float32
-    length = x.shape[1]
-    ys = []
-    for i in range(length):
-        y, ssm_state = ssd_update_state(
-            ssm_state,
-            x[:, i],
-            dt[:, i],
-            A,
-            B[:, i],
-            C[:, i],
-            D if D.ndim == 1 else D[:, i],
-            z=z[:, i],
-            dt_bias=dt_bias,
-            dt_softplus=dt_softplus,
-        )
-        ys.append(y)
-    return mx.concatenate(ys, axis=1), ssm_state
-
-
-def causal_conv1d_update(conv_state, x, weight) -> tuple[mx.array, mx.array]:
-    _, seqlen, dim = x.shape
-    state_len = conv_state.shape[-2]
-    x = mx.concatenate([conv_state, x], axis=-2)
-    conv_state = x[:, -state_len:]
-    out = mx.conv1d(
-        x,
-        weight,
-        padding=0,
-        groups=dim,
-    )[:, -seqlen:]
-    return nn.silu(out), conv_state
-
-
-class Mamba(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.d_state = config.mamba_d_state
-        self.d_conv = config.mamba_d_conv
-        self.chunk_size = config.mamba_chunk_size
-        self.num_heads = config.mamba_num_heads
-        self.hidden_size_per_head = config.hidden_size_per_head
-
-        self.intermediate_size = self.num_heads * self.hidden_size_per_head
-
-        self.in_proj = nn.Linear(
-            self.hidden_size, 2 * self.intermediate_size, bias=False
-        )
-        self.conv1d = nn.Conv1d(
-            in_channels=self.intermediate_size,
-            out_channels=self.intermediate_size,
-            bias=False,
-            kernel_size=self.d_conv,
-            groups=self.intermediate_size,
-            padding=0,
-        )
-        self.dt_dim = max(64, self.hidden_size // 16)
-        self.bcdt_proj = nn.Linear(
-            self.intermediate_size,
-            self.dt_dim + 2 * self.d_state,
-            bias=False,
-        )
-        self.dt_proj = nn.Linear(self.dt_dim, self.num_heads, bias=False)
-
-        self.dt_bias = get_initial_dt_bias(self.num_heads)
-        self.A_log = get_initial_A(self.num_heads)
-        self.D = mx.ones(self.num_heads, dtype=mx.float32)
-
-        self.dt_norm_weight = mx.ones(self.dt_dim)
-        self.B_norm_weight = mx.ones(self.d_state)
-        self.C_norm_weight = mx.ones(self.d_state)
-
-        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        mask: Optional[mx.array] = None,
-        cache=None,
-    ):
-        bsize, length, _ = hidden_states.shape
-
-        if cache is not None and cache[0] is not None:
-            conv_state = cache[0]
-            ssm_state = cache[1]
-        else:
-            conv_state = mx.zeros(
-                (bsize, self.d_conv - 1, self.intermediate_size),
-                dtype=hidden_states.dtype,
-            )
-            ssm_state = mx.zeros(
-                (bsize, self.num_heads, self.hidden_size_per_head, self.d_state),
-                dtype=mx.float32,
-            )
-
-        zx = self.in_proj(hidden_states)
-        zx = zx.reshape(bsize, length, self.num_heads, -1)
-        # z: (bsize, length, num_heads, hidden_size_per_head)
-        # x: (bsize, length, num_heads, hidden_size_per_head)
-        z, x = mx.split(
-            zx,
-            [
-                self.hidden_size_per_head,
-            ],
-            axis=-1,
-        )
-
-        x = x.reshape(bsize, -1, self.num_heads * self.hidden_size_per_head)
-        x, conv_state = causal_conv1d_update(conv_state, x, self.conv1d.weight)
-        BCdt = self.bcdt_proj(x)
-        x = x.reshape(bsize, length, self.num_heads, -1)
-        B, C, dt = mx.split(BCdt, [self.d_state, self.d_state * 2], axis=-1)
-
-        A = -mx.exp(self.A_log.astype(mx.float32))  # (num_heads,)
-        dt = mx.fast.rms_norm(dt, self.dt_norm_weight, self.config.rms_norm_eps)
-        B = mx.fast.rms_norm(B, self.B_norm_weight, self.config.rms_norm_eps)
-        C = mx.fast.rms_norm(C, self.C_norm_weight, self.config.rms_norm_eps)
-
-        # (bsize, length, num_heads, 1)
-        dt = self.dt_proj(dt)[..., None]
-
-        out, ssm_state = ssd_chunk_scan_combined(
-            x,
-            dt.reshape(bsize, length, -1),
-            A,
-            B,
-            C,
-            D=self.D,
-            z=z,
-            dt_bias=self.dt_bias,
-            dt_softplus=True,
-            ssm_state=ssm_state,
-        )
-
-        if cache is not None:
-            cache[0] = conv_state
-            cache[1] = ssm_state
-        y = self.out_proj(out.reshape(bsize, length, -1))
-
-        return y
-
-
-class Attention(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        head_dim = config.hidden_size_per_head
-        self.max_position_embeddings = config.max_position_embeddings
-        self.scale = head_dim**-0.5
-
-        self.q_num_heads = config.num_attention_heads
-        self.qk_dim = self.v_dim = head_dim
-        self.k_num_heads = self.v_num_heads = config.num_key_value_heads
-        assert self.q_num_heads % self.k_num_heads == 0
-        self.n_group = self.q_num_heads // self.k_num_heads
-
-        self.q_proj_dim = self.q_num_heads * self.qk_dim
-        self.k_proj_dim = self.k_num_heads * self.qk_dim
-        self.v_proj_dim = self.k_num_heads * self.v_dim
-        self.qkv_proj = nn.Linear(
-            self.hidden_size,
-            self.q_proj_dim + self.k_proj_dim + self.v_proj_dim,
-            bias=False,
-        )
-        self.o_proj = nn.Linear(
-            self.q_num_heads * self.v_dim, self.hidden_size, bias=False
-        )
-
-        self.q_weight = mx.ones((self.q_num_heads, self.qk_dim))
-        self.k_weight = mx.ones((self.k_num_heads, self.qk_dim))
-
-        self.rope = nn.RoPE(self.qk_dim)
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        mask: Optional[mx.array] = None,
-        cache=None,
-    ):
-        B, T, _ = hidden_states.shape
-
-        qkv = self.qkv_proj(hidden_states)
-        q, k, v = mx.split(
-            qkv, [self.q_proj_dim, self.q_proj_dim + self.k_proj_dim], axis=-1
-        )
-        q = q.reshape(B, T, self.q_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
-        k = k.reshape(B, T, self.k_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
-        v = v.reshape(B, T, self.v_num_heads, self.v_dim).transpose(0, 2, 1, 3)
-
-        q = mx.fast.rms_norm(q, weight=None, eps=1e-6) * self.q_weight[:, None]
-        k = mx.fast.rms_norm(k, weight=None, eps=1e-6) * self.k_weight[:, None]
-
-        if cache is not None:
-            q = self.rope(q, offset=cache.offset)
-            k = self.rope(k, offset=cache.offset)
-            k, v = cache.update_and_fetch(k, v)
-        else:
-            q = self.rope(q)
-            k = self.rope(k)
-
-        output = mx.fast.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            scale=self.scale,
-            mask=mask,
-        )
-        output = output.transpose(0, 2, 1, 3).reshape(
-            B, T, self.q_num_heads * self.v_dim
-        )
-        return self.o_proj(output)
-
-
-class MLP(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_up_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size * 2, bias=False
-        )
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-
-    def __call__(self, x: mx.array) -> mx.array:
-        h = self.gate_up_proj(x)
-        hs = mx.split(h, 2, axis=-1)
-        return self.down_proj(nn.silu(hs[0]) * hs[1])
-
-
-class PlamoDecoderLayer(nn.Module):
-    def __init__(self, config: ModelArgs, is_mamba: bool) -> None:
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.is_mamba = is_mamba
-        self.mixer: nn.Module
-        if is_mamba:
-            self.mixer = Mamba(config)
-        else:
-            self.mixer = Attention(config)
-        self.mlp = MLP(config)
-        self.pre_mixer_norm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, offset=1.0
-        )
-        self.post_mixer_norm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / 5
-        )
-        self.pre_mlp_norm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, offset=1.0
-        )
-        self.post_mlp_norm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / (5**1.5)
-        )
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        mask: Optional[mx.array] = None,
-        cache=None,
-    ):
-        residual = hidden_states
-        hidden_states = self.pre_mixer_norm(hidden_states)
-
-        hidden_states_sa = self.mixer(
-            hidden_states=hidden_states,
-            mask=mask,
-            cache=cache,
-        )
-
-        hidden_states_sa = self.post_mixer_norm(hidden_states_sa)
-        hidden_states = residual + hidden_states_sa
-
-        residual = hidden_states
-        hidden_states = self.pre_mlp_norm(hidden_states)
-
-        # Fully Connected
-        hidden_states_mlp = self.mlp(hidden_states)
-
-        # Residual
-        hidden_states_mlp = self.post_mlp_norm(hidden_states_mlp)
-        return residual + hidden_states_mlp
-
-
-def is_mamba(config: ModelArgs, i: int) -> bool:
-    if not config.mamba_enabled:
-        return False
-    assert config.mamba_step > 1
-    assert i < config.num_hidden_layers
-
-    if config.num_hidden_layers <= (config.mamba_step // 2):
-        # use attention in last layer
-        return i != config.num_hidden_layers - 1
-    return (i % config.mamba_step) != (config.mamba_step // 2)
-
-
-class PlamoDecoder(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-
-        self.layers = [
-            PlamoDecoderLayer(config, is_mamba=is_mamba(config, i))
-            for i in range(config.num_hidden_layers)
-        ]
-
-    def __call__(self, x: mx.array, mask: mx.array, cache):
-        for i, decoder_layer in enumerate(self.layers):
-            x = decoder_layer(
-                x,
-                mask=mask,
-                cache=cache[i],
-            )
-        return x
-
-
-class PlamoModel(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-
-        self.config = config
-        self.vocab_size = config.vocab_size
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = PlamoDecoder(config)  # type: ignore
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        mask: Optional[mx.array] = None,
-        cache=None,
-    ):
-        batch_size, seq_length = inputs.shape
-
-        h = self.embed_tokens(inputs)
-
-        if mask is None:
-            mask = create_attention_mask(h, [cache[1]] if cache is not None else None)
-
-        if cache is None:
-            cache = [None] * len(self.layers.layers)
-
-        # decoder layers
-        out = self.layers(
-            h,
-            mask,
-            cache,
-        )
-
-        return self.norm(out)
-
-
-class Model(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.config = config
-        self.model_type = config.model_type
-        self.model = PlamoModel(config)
-
-        self.vocab_size = config.vocab_size
-
-        if not config.tie_word_embeddings:
-            self.lm_head: nn.Module = nn.Linear(
-                config.hidden_size, self.vocab_size, bias=False
-            )
-
-    def sanitize(self, weights: dict[Any, Any]) -> dict[Any, Any]:
-        for k, v in weights.items():
-            if "conv1d.weight" in k and v.shape[-1] != 1:
-                weights[k] = v.moveaxis(2, 1)
-        return weights
-
-    def make_cache(self):
-        # TODO use RotatingKVCache is not full_attn
-        # full_attn = self.layer_idx in self.config.full_attention_idx
-        return [MambaCache() if l.is_mamba else KVCache() for l in self.layers]
-
-    def __call__(
-        self, inputs: mx.array, mask: Optional[mx.array] = None, cache=None
-    ) -> mx.array:
-        outputs = self.model(
-            inputs=inputs,
-            mask=None,
-            cache=cache,
-        )
-        if self.config.tie_word_embeddings:
-            logits = self.model.embed_tokens.as_linear(outputs)
-        else:
-            logits = self.lm_head(outputs)
-
-        return logits
-
-    @property
-    def layers(self):
-        return self.model.layers.layers

mlx_lm/models/qwen2.py

@@ -7,7 +7,6 @@ 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
@@ -19,13 +18,24 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     rms_norm_eps: float
     vocab_size: int
-    num_key_value_heads: int
-    max_position_embeddings: int = 32768
+    num_key_value_heads: Optional[int] = None
     rope_theta: float = 1000000
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = True
 
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        if self.rope_scaling:
+            required_keys = {"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}")
+
+            if self.rope_scaling["type"] != "linear":
+                raise ValueError("rope_scaling 'type' currently only supports 'linear'")
+
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -44,12 +54,16 @@ class Attention(nn.Module):
         self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
 
-        self.rope = initialize_rope(
+        rope_scale = (
+            1 / args.rope_scaling["factor"]
+            if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
+            else 1
+        )
+        self.rope = nn.RoPE(
             head_dim,
-            base=args.rope_theta,
             traditional=args.rope_traditional,
-            scaling_config=args.rope_scaling,
-            max_position_embeddings=args.max_position_embeddings,
+            base=args.rope_theta,
+            scale=rope_scale,
         )
 
     def __call__(

mlx_lm/models/qwen3.py

@@ -1,189 +0,0 @@
-# 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
-
-
-@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: Optional[Dict[str, Union[float, str]]] = None
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        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.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:
-            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):
-        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):
-        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)
-        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 Qwen3Model(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)
-        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, 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 = Qwen3Model(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 sanitize(self, weights):
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/qwen3_moe.py

@@ -1,243 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import math
-from dataclasses import dataclass
-from typing import Any, Dict, List, 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
-    num_experts_per_tok: int
-    num_experts: int
-    num_experts_per_tok: int
-    decoder_sparse_step: int
-    mlp_only_layers: List[int]
-    moe_intermediate_size: int
-    rms_norm_eps: float
-    vocab_size: int
-    num_key_value_heads: int
-    head_dim: int
-    rope_theta: float
-    tie_word_embeddings: bool
-    max_position_embeddings: int
-    norm_topk_prob: bool
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        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 = getattr(
-            args, "head_dim", args.hidden_size // args.num_attention_heads
-        )
-        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.rope = nn.RoPE(
-            head_dim,
-            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)
-
-        # Prepare the queries, keys and values for the attention computation
-        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:
-            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):
-        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 Qwen3MoeSparseMoeBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        dim = args.hidden_size
-        intermediate_size = args.moe_intermediate_size
-
-        self.num_experts = num_experts = args.num_experts
-        self.top_k = args.num_experts_per_tok
-        self.norm_topk_prob = args.norm_topk_prob
-
-        self.gate = nn.Linear(dim, num_experts, bias=False)
-        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
-
-    def __call__(
-        self,
-        x: mx.array,
-    ):
-        gates = self.gate(x)
-        gates = mx.softmax(gates, axis=-1, precise=True)
-
-        k = self.top_k
-        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
-        scores = mx.take_along_axis(gates, inds, axis=-1)
-        if self.norm_topk_prob:
-            scores /= mx.sum(scores, axis=-1, keepdims=True)
-
-        y = self.switch_mlp(x, inds)
-        y = (y * scores[..., None]).sum(axis=-2)
-
-        return y
-
-
-class Qwen3MoeDecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.hidden_size = args.hidden_size
-        self.self_attn = Attention(args, layer_idx)
-
-        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
-
-        if (layer_idx not in args.mlp_only_layers) and (
-            args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0
-        ):
-            self.mlp = Qwen3MoeSparseMoeBlock(args)
-        else:
-            self.mlp = MLP(args.hidden_size, args.intermediate_size)
-
-    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 Qwen3MoeModel(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)
-        self.layers = [
-            Qwen3MoeDecoderLayer(args=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: 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 = Qwen3MoeModel(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,
-    ):
-        out = self.model(inputs, mask, cache)
-        return self.lm_head(out)
-
-    def sanitize(self, weights):
-        if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
-            return weights
-        for l in range(self.args.num_hidden_layers):
-            prefix = f"model.layers.{l}"
-            for n in ["up_proj", "down_proj", "gate_proj"]:
-                if f"{prefix}.mlp.experts.0.{n}.weight" in weights:
-                    to_join = [
-                        weights.pop(f"{prefix}.mlp.experts.{e}.{n}.weight")
-                        for e in range(self.args.num_experts)
-                    ]
-                    weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(to_join)
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/rope_utils.py

@@ -1,6 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import math
 from typing import Optional
 
 import mlx.core as mx
@@ -62,78 +61,6 @@ class Llama3RoPE(nn.Module):
         )
 
 
-class YarnRoPE(nn.Module):
-    def __init__(
-        self,
-        dims,
-        traditional=False,
-        max_position_embeddings=2048,
-        base=10000,
-        scaling_factor=1.0,
-        original_max_position_embeddings=4096,
-        beta_fast=32,
-        beta_slow=1,
-        mscale=1,
-        mscale_all_dim=0,
-    ):
-        super().__init__()
-
-        def yarn_find_correction_dim(num_rotations):
-            return (
-                dims
-                * math.log(
-                    original_max_position_embeddings / (num_rotations * 2 * math.pi)
-                )
-            ) / (2 * math.log(base))
-
-        def yarn_find_correction_range():
-            low = math.floor(yarn_find_correction_dim(beta_fast))
-            high = math.ceil(yarn_find_correction_dim(beta_slow))
-            return max(low, 0), min(high, dims - 1)
-
-        def yarn_get_mscale(scale=1, mscale=1):
-            if scale <= 1:
-                return 1.0
-            return 0.1 * mscale * math.log(scale) + 1.0
-
-        def yarn_linear_ramp_mask(min_val, max_val, dim):
-            if min_val == max_val:
-                max_val += 0.001  # Prevent singularity
-
-            linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (
-                max_val - min_val
-            )
-            return mx.clip(linear_func, 0, 1)
-
-        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
-            scaling_factor, mscale_all_dim
-        )
-        freq_extra = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        freq_inter = scaling_factor * base ** (
-            mx.arange(0, dims, 2, dtype=mx.float32) / dims
-        )
-        low, high = yarn_find_correction_range()
-        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
-        self._freqs = (freq_inter * freq_extra) / (
-            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
-        )
-        self.dims = dims
-        self.traditional = traditional
-
-    def __call__(self, x, offset=0):
-        if self.mscale != 1.0:
-            x[..., : self.dims] = self.mscale * x[..., : self.dims]
-        return mx.fast.rope(
-            x,
-            self.dims,
-            traditional=self.traditional,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
 def initialize_rope(
     dims,
     base,
@@ -160,25 +87,5 @@ def initialize_rope(
             base=base,
             scaling_config=scaling_config,
         )
-    elif rope_type == "yarn":
-        scaling_factor = scaling_config["factor"]
-        rope_kwargs = {
-            key: scaling_config[key]
-            for key in [
-                "original_max_position_embeddings",
-                "beta_fast",
-                "beta_slow",
-                "mscale",
-                "mscale_all_dim",
-            ]
-            if key in scaling_config
-        }
-        return YarnRoPE(
-            dims=dims,
-            max_position_embeddings=max_position_embeddings,
-            traditional=traditional,
-            base=base,
-            **rope_kwargs,
-        )
     else:
         raise ValueError(f"Unsupported RoPE type {rope_type}")

mlx_lm/models/su_rope.py

@@ -20,7 +20,7 @@ class SuScaledRotaryEmbedding(nn.Module):
         long_mscale: float = None,
     ):
         """
-        Su Scaled Rotary Embedding layer.
+        Phi3Su Scaled Rotary Embedding layer for Phi-3 models.
 
         Args:
             dims (int): The feature dimensions to be rotated.
@@ -51,13 +51,11 @@ class SuScaledRotaryEmbedding(nn.Module):
             + math.log(max_position_embeddings / original_max_position_embeddings)
             / math.log(original_max_position_embeddings)
         )
-        self.dim = dims
 
     def __call__(self, x, offset: int = 0):
-        x[..., : self.dim] = self.scale * x[..., : self.dim]
         return mx.fast.rope(
-            x,
-            self.dim,
+            self.scale * x,
+            x.shape[-1],
             traditional=False,
             base=None,
             scale=1.0,

mlx_lm/models/switch_layers.py

@@ -6,21 +6,6 @@ import mlx.core as mx
 import mlx.nn as nn
 
 
-def _gather_sort(x, indices):
-    *_, M = indices.shape
-    indices = indices.flatten()
-    order = mx.argsort(indices)
-    inv_order = mx.argsort(order)
-    return x.flatten(0, -3)[order // M], indices[order], inv_order
-
-
-def _scatter_unsort(x, inv_order, shape=None):
-    x = x[inv_order]
-    if shape is not None:
-        x = mx.unflatten(x, 0, shape)
-    return x
-
-
 class QuantizedSwitchLinear(nn.Module):
     def __init__(
         self,
@@ -71,7 +56,7 @@ class QuantizedSwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices, sorted_indices=False):
+    def __call__(self, x, indices):
         x = mx.gather_qmm(
             x,
             self["weight"],
@@ -81,7 +66,6 @@ class QuantizedSwitchLinear(nn.Module):
             transpose=True,
             group_size=self.group_size,
             bits=self.bits,
-            sorted_indices=sorted_indices,
         )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
@@ -115,13 +99,8 @@ class SwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices, sorted_indices=False):
-        x = mx.gather_mm(
-            x,
-            self["weight"].swapaxes(-1, -2),
-            rhs_indices=indices,
-            sorted_indices=sorted_indices,
-        )
+    def __call__(self, x, indices):
+        x = mx.gather_mm(x, self["weight"].swapaxes(-1, -2), rhs_indices=indices)
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
         return x
@@ -143,7 +122,7 @@ class SwitchGLU(nn.Module):
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.SiLU(),
+        activation=nn.silu,
         bias: bool = False,
     ):
         super().__init__()
@@ -156,24 +135,9 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        # When we have many tokens, then sort them to make sure that the access
-        # of different experts is in order.
-        do_sort = indices.size >= 64
-        idx = indices
-        inv_order = None
-        if do_sort:
-            x, idx, inv_order = _gather_sort(x, indices)
-
-        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,
-            idx,
-            sorted_indices=do_sort,
-        )
-
-        if do_sort:
-            x = _scatter_unsort(x, inv_order, indices.shape)
+        x_up = self.up_proj(x, indices)
+        x_gate = self.gate_proj(x, indices)
+        x = self.down_proj(self.activation(x_gate) * x_up, indices)
 
         return x.squeeze(-2)
 
@@ -184,7 +148,7 @@ class SwitchMLP(nn.Module):
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.GELU(approx="precise"),
+        activation=nn.gelu_approx,
         bias: bool = False,
     ):
         super().__init__()
@@ -196,19 +160,8 @@ class SwitchMLP(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        # When we have many tokens, then sort them to make sure that the access
-        # of different experts is in order.
-        do_sort = indices.size >= 64
-        idx = indices
-        inv_order = None
-        if do_sort:
-            x, idx, inv_order = _gather_sort(x, indices)
-
-        x = self.fc1(x, idx, sorted_indices=do_sort)
+        x = self.fc1(x, indices)
         x = self.activation(x)
-        x = self.fc2(x, idx, sorted_indices=do_sort)
-
-        if do_sort:
-            x = _scatter_unsort(x, inv_order, indices.shape)
+        x = self.fc2(x, indices)
 
         return x.squeeze(-2)

mlx_lm/qat.py

@@ -1,339 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-import copy
-import glob
-import shutil
-import time
-import types
-from pathlib import Path
-
-import mlx.core as mx
-import mlx.nn as nn
-import mlx.optimizers as optimizers
-import numpy as np
-from mlx.utils import tree_flatten, tree_map, tree_map_with_path
-
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-from mlx_lm.tuner.datasets import load_dataset
-from mlx_lm.tuner.trainer import iterate_batches
-from mlx_lm.tuner.utils import print_trainable_parameters
-from mlx_lm.utils import (
-    create_model_card,
-    fetch_from_hub,
-    get_model_path,
-    quantize_model,
-    save_config,
-    save_weights,
-)
-
-
-class StraightThroughQuantizedEmbedding(nn.Module):
-    def __init__(
-        self,
-        num_embeddings: int,
-        dims: int,
-        group_size: int = 64,
-        bits: int = 4,
-    ):
-        super().__init__()
-
-        # Quantization config
-        self.group_size = group_size
-        self.bits = bits
-
-        # Initialize the quantized weight
-        self.weight = mx.zeros(shape=(num_embeddings, dims))
-        self.num_embeddings = num_embeddings
-        self.dims = dims
-
-    def __call__(self, x):
-        w, s, b = mx.quantize(self.weight, self.group_size, self.bits)
-        y = self.weight[x]
-        yq = mx.dequantize(
-            w[x],
-            scales=s[x],
-            biases=b[x],
-            group_size=self.group_size,
-            bits=self.bits,
-        )
-        return (y - mx.stop_gradient(y)) + mx.stop_gradient(yq)
-
-    def as_linear(self, x):
-        # Quantize and then matmul
-        w, s, b = mx.quantize(self.weight, self.group_size, self.bits)
-        y = x @ self.weight.T
-        yq = mx.quantized_matmul(
-            x,
-            w,
-            scales=s,
-            biases=b,
-            transpose=True,
-            group_size=self.group_size,
-            bits=self.bits,
-        )
-        return (y - mx.stop_gradient(y)) + mx.stop_gradient(yq)
-
-    @classmethod
-    def from_embedding(
-        cls, embedding_layer: nn.Module, group_size: int = 64, bits: int = 4
-    ):
-        embedding_dims, dims = embedding_layer.weight.shape
-        ql = cls(embedding_dims, dims, group_size, bits)
-        ql.weight = embedding_layer.weight
-        return ql
-
-
-class StraightThroughQuantizedLinear(nn.Module):
-    def __init__(
-        self,
-        input_dims: int,
-        output_dims: int,
-        bias: bool = True,
-        group_size: int = 64,
-        bits: int = 4,
-    ):
-        super().__init__()
-
-        # Quantization config
-        self.group_size = group_size
-        self.bits = bits
-
-        self.weight = mx.zeros(shape=(output_dims, input_dims))
-        if bias:
-            self.bias = mx.zeros((output_dims,))
-
-    def __call__(self, x):
-        # Quantize and then matmul
-        w, s, b = mx.quantize(self.weight, self.group_size, self.bits)
-        y = x @ self.weight.T
-        yq = mx.quantized_matmul(
-            x,
-            w,
-            scales=s,
-            biases=b,
-            transpose=True,
-            group_size=self.group_size,
-            bits=self.bits,
-        )
-        x = (y - mx.stop_gradient(y)) + mx.stop_gradient(yq)
-        if "bias" in self:
-            x = x + self["bias"]
-        return x
-
-    @classmethod
-    def from_linear(cls, linear_layer: nn.Module, group_size: int = 64, bits: int = 4):
-        output_dims, input_dims = linear_layer.weight.shape
-        ql = cls(input_dims, output_dims, False, group_size, bits)
-        if "bias" in linear_layer:
-            ql.bias = linear_layer.bias
-        return ql
-
-
-def quantize(
-    model: nn.Module,
-    group_size: int = 64,
-    bits: int = 4,
-):
-    def _maybe_quantize(path, m):
-        if isinstance(m, nn.Linear):
-            return StraightThroughQuantizedLinear.from_linear(
-                m, group_size=group_size, bits=bits
-            )
-        elif isinstance(m, nn.Embedding):
-            return StraightThroughQuantizedEmbedding.from_embedding(
-                m, group_size=group_size, bits=bits
-            )
-        else:
-            return m
-
-    leaves = tree_map_with_path(
-        _maybe_quantize, model.leaf_modules(), is_leaf=nn.Module.is_module
-    )
-    model.update_modules(leaves)
-
-
-def qat(
-    model,
-    opt,
-    data,
-    group_size: int = 64,
-    bits: int = 3,
-    batch_size: int = 2,
-    max_seq_length: int = 2048,
-    temperature: float = 0.5,
-    dtype: mx.Dtype = mx.bfloat16,
-):
-    group = mx.distributed.init()
-    world_size = group.size()
-    rank = group.rank()
-
-    def log_norm(x):
-        x = x * (1 / temperature)
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
-
-    q_model = copy.deepcopy(model)
-    quantize(q_model, bits=bits, group_size=group_size)
-
-    def loss_fn(params, x, targets, lengths):
-        q_model.update(tree_map(lambda x: x.astype(dtype), params))
-        logits = q_model(x).astype(mx.float32)
-        losses = nn.losses.kl_div_loss(log_norm(logits), targets, reduction="none")
-        mask = mx.arange(targets.shape[1]) < lengths[:, 1:]
-        ntoks = mask.sum()
-        loss = (mask * losses).sum() / ntoks
-        return loss, ntoks
-
-    def step(inputs, targets, lengths, params):
-        (loss, ntoks), grads = mx.value_and_grad(loss_fn)(
-            params, inputs, targets, lengths
-        )
-        grads = nn.average_gradients(grads)
-        params = opt.apply_gradients(grads, params)
-        return loss, ntoks, params
-
-    # Accumulate learned weights in higher precision
-    params = tree_map(
-        lambda x: x.astype(mx.float32),
-        model.trainable_parameters(),
-    )
-
-    avg_loss = None
-    tokens = 0
-    tic = time.time()
-    for it, (batch, lengths) in enumerate(
-        iterate_batches(data, batch_size, max_seq_length)
-    ):
-        targets = log_norm(model(batch).astype(mx.float32))
-        mx.eval(targets)
-        loss, ntoks, params = step(batch, targets, lengths, params)
-        mx.eval(loss, params)
-        loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
-        ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
-        tokens += ntoks
-        toks_per_sec = tokens / (time.time() - tic)
-        avg_loss = 0.95 * (avg_loss or loss) + 0.05 * loss
-        if rank == 0:
-            print(
-                f"{it=}, {loss=:.3f}, {avg_loss=:.4f}, {tokens=}, {toks_per_sec=:.3f}",
-                flush=True,
-            )
-    model.update(tree_map(lambda x: x.astype(dtype), params))
-
-
-def save_model(
-    model: nn.Module,
-    tokenizer: TokenizerWrapper,
-    config,
-    model_path: Path,
-    mlx_path: str,
-    hf_path: str,
-):
-    weights = dict(tree_flatten(model.parameters()))
-
-    mlx_path = Path(mlx_path)
-    save_weights(mlx_path, weights, donate_weights=True)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, mlx_path)
-
-    tokenizer.save_pretrained(mlx_path)
-
-    save_config(config, config_path=mlx_path / "config.json")
-    create_model_card(mlx_path, hf_path)
-
-
-def load_data(tokenizer, data_path: str, num_samples: int):
-    args = types.SimpleNamespace(
-        hf_dataset={
-            "path": data_path,
-            "train_split": f"train[:{num_samples}]",
-            "valid_split": "train[:1]",
-        },
-        train=True,
-        test=False,
-    )
-    dataset = load_dataset(args, tokenizer)[0]
-    return [dataset.process(d) for d in dataset]
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", "-m", default="Qwen/Qwen3-1.7B")
-    parser.add_argument(
-        "--mlx-path", default="mlx_model", help="Path to save the quantized model."
-    )
-    parser.add_argument(
-        "--bits",
-        type=int,
-        default=4,
-        help="Bits per weight for quantization.",
-    )
-    parser.add_argument(
-        "--group-size", type=int, default=64, help="Group size for quantization."
-    )
-    parser.add_argument(
-        "--num-samples",
-        type=int,
-        default=1024,
-        help="Number of samples to use for training.",
-    )
-    parser.add_argument("--max-seq-length", type=int, default=2048)
-    parser.add_argument("--seed", type=int, default=123)
-    parser.add_argument("--learning-rate", type=float, default=1e-5)
-    parser.add_argument("--batch-size", type=int, default=8)
-    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(
-        "--temperature",
-        type=float,
-        default=0.5,
-        help="Temperature scaling for the loss.",
-    )
-    args = parser.parse_args()
-
-    group = mx.distributed.init()
-
-    num_samples = args.num_samples
-    if num_samples % group.size() > 0:
-        num_samples += group.size() - num_samples % group.size()
-
-    np.random.seed(args.seed)
-    mx.random.seed(args.seed)
-
-    model_path = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
-    if "quantization" in config:
-        raise ValueError("Teacher model for QAT training should not be quantized")
-
-    calibration_data = load_data(tokenizer, args.data_path, args.num_samples)
-
-    q_model = copy.deepcopy(model)
-    tree_flatten(q_model.parameters())
-    opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
-    qat(
-        model,
-        opt,
-        calibration_data,
-        bits=args.bits,
-        group_size=args.group_size,
-        batch_size=args.batch_size,
-        max_seq_length=args.max_seq_length,
-        temperature=args.temperature,
-    )
-    _, config = quantize_model(
-        model,
-        config,
-        q_group_size=args.group_size,
-        q_bits=args.bits,
-    )
-    save_model(model, tokenizer, config, model_path, args.mlx_path, args.model)
-
-
-if __name__ == "__main__":
-    main()

mlx_lm/requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.25.0
+mlx>=0.19.2
 numpy
 transformers[sentencepiece]>=4.39.3
 protobuf

mlx_lm/sample_utils.py

@@ -2,7 +2,7 @@
 
 import math
 from functools import partial
-from typing import Callable, Dict, List, Optional
+from typing import Callable, Dict, Optional
 
 import mlx.core as mx
 
@@ -12,10 +12,7 @@ def make_sampler(
     top_p: float = 0.0,
     min_p: float = 0.0,
     min_tokens_to_keep: int = 1,
-    top_k: int = 0,
-    xtc_probability: float = 0.0,
-    xtc_threshold: float = 0.0,
-    xtc_special_tokens: List[int] = [],
+    top_k: int = -1,
 ) -> Callable[mx.array, mx.array]:
     """
     Make a sampler function for use with ``generate_step``.
@@ -31,13 +28,6 @@ def make_sampler(
           be filtered by min_p sampling.
         top_k (int, optional): The top k tokens ranked by probability to constrain
           the sampling to.
-        xtc_probability (float, optional): The probability of applying XTC
-            sampling.
-        xtc_threshold (float, optional): The threshold the probs need to reach
-            for being sampled.
-        xtc_special_tokens (list(int), optional): List of special tokens IDs to
-            be excluded from XTC sampling.
-
 
     Returns:
         Callable[mx.array, mx.array]:
@@ -45,29 +35,14 @@ def make_sampler(
     """
     if temp == 0:
         return lambda x: mx.argmax(x, axis=-1)
-
-    # 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:
-        sampling_methods.append(lambda x: apply_min_p(x, min_p, min_tokens_to_keep))
-    if xtc_probability > 0.0:
-        sampling_methods.append(
-            lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens)
-        )
-
-    # Apply the sampling methods
-    def sampler(logits):
-        for method in sampling_methods:
-            logits = method(logits)
-
-        # Return the sampled token
-        return categorical_sampling(logits, temp)
-
-    return sampler
+    elif top_p > 0 and top_p < 1.0:
+        return lambda x: top_p_sampling(x, top_p, temp)
+    elif min_p != 0.0:
+        return lambda x: min_p_sampling(x, min_p, min_tokens_to_keep, temp)
+    elif top_k > 0:
+        return lambda x: top_k_sampling(x, top_k, temp)
+    else:
+        return lambda x: categorical_sampling(x, temp)
 
 
 def make_logits_processors(
@@ -110,9 +85,10 @@ def make_logits_processors(
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_top_k(
+def top_k_sampling(
     logprobs: mx.array,
     top_k: int,
+    temperature=1.0,
 ) -> mx.array:
     """
     Sample from only the top K tokens ranked by probability.
@@ -127,18 +103,20 @@ def apply_top_k(
             f"`top_k` has to be an integer in the (0, {vocab_size}] interval,"
             f" but is {top_k}."
         )
+    logprobs = logprobs * (1 / temperature)
     mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]
     masked_logprobs = mx.put_along_axis(
         logprobs, mask_idx, mx.array(-float("inf"), logprobs.dtype), axis=-1
     )
-    return masked_logprobs
+    return mx.random.categorical(masked_logprobs, axis=-1)
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_min_p(
+def min_p_sampling(
     logprobs: mx.array,
     min_p: float,
     min_tokens_to_keep: int = 1,
+    temperature=1.0,
 ) -> mx.array:
     """
     Apply min-p sampling to the logprobs.
@@ -166,12 +144,14 @@ def apply_min_p(
         )
     # reference implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L531-L605
 
+    logprobs = logprobs * (1 / temperature)
+
     # Indices sorted in decreasing order
-    sorted_indices = mx.argsort(-logprobs, axis=-1)
-    sorted_logprobs = mx.take_along_axis(logprobs, sorted_indices, axis=-1)
+    sorted_indices = mx.argsort(-logprobs).squeeze(0)
+    sorted_logprobs = logprobs[..., sorted_indices]
 
     # Top probability
-    top_logprobs = sorted_logprobs[:, 0:1]
+    top_logprobs = logprobs[..., sorted_indices[0]]
 
     # Calculate the min_p threshold
     scaled_min_p = top_logprobs + math.log(min_p)
@@ -183,34 +163,29 @@ def apply_min_p(
     # Create pool of tokens with probability less than scaled min_p
     selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)
 
-    # Create a mapping to rearrange back to original indices
-    # Use argsort of sorted_indices to get the inverse permutation
-    inverse_indices = mx.argsort(sorted_indices, axis=-1)
-
-    # Rearrange selected_logprobs back to original order
-    original_order_logprobs = mx.take_along_axis(
-        selected_logprobs, inverse_indices, axis=-1
-    )
-
-    return original_order_logprobs
+    # Return sampled token
+    sorted_token = mx.random.categorical(selected_logprobs)
+    return sorted_indices[sorted_token]
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_top_p(logits: mx.array, top_p: float) -> mx.array:
+def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.array:
     """
     Apply top-p (nucleus) sampling to logits.
 
     Args:
         logits: The logits from the model's output.
         top_p: The cumulative probability threshold for top-p filtering.
+        temperature: Temperature parameter for softmax distribution reshaping.
     Returns:
         token selected based on the top-p criterion.
     """
     # referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460
-    probs = mx.softmax(logits, axis=-1)
+    probs = mx.softmax(logits * (1 / temperature), axis=-1)
+
     # sort probs in ascending order
     sorted_indices = mx.argsort(probs, axis=-1)
-    sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
+    sorted_probs = probs[..., sorted_indices.squeeze(0)]
 
     cumulative_probs = mx.cumsum(sorted_probs, axis=-1)
 
@@ -221,52 +196,10 @@ def apply_top_p(logits: mx.array, top_p: float) -> mx.array:
         0,
     )
 
-    # Create a mapping to rearrange back to original indices
-    # Use argsort of sorted_indices to get the inverse permutation
-    inverse_indices = mx.argsort(sorted_indices, axis=-1)
+    sorted_token = mx.random.categorical(mx.log(top_probs))
+    token = sorted_indices.squeeze(0)[sorted_token]
 
-    # Rearrange top_probs back to original order
-    original_order_probs = mx.take_along_axis(top_probs, inverse_indices, axis=-1)
-
-    # Convert back to logits and return
-    return mx.log(original_order_probs)
-
-
-@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_xtc(
-    logits: mx.array,
-    xtc_probability: float,
-    xtc_threshold: float,
-    xtc_special_tokens: List[int],
-) -> mx.array:
-    """
-    Apply XTC sampling to the logits.
-
-    Args:
-        logits: The logits from the model's output.
-        xtc_probability (float): Probability of XTC sampling to happen for each token
-        xtc_threshold (float): The threshold the probs need to reach for being sampled.
-        special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling.
-    """
-    if not (0 <= xtc_threshold <= 0.5):
-        raise ValueError(
-            f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}"
-        )
-    if not (0 <= xtc_probability <= 1.0):
-        raise ValueError(
-            f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}"
-        )
-
-    probs = mx.softmax(logits, -1)
-    mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
-    if xtc_special_tokens:
-        mask[..., xtc_special_tokens] = False
-
-    return mx.where(
-        mx.random.uniform(0, 1) > xtc_probability,
-        logits,
-        mx.where(mask, -mx.inf, logits),
-    )
+    return token
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)

mlx_lm/server.py

@@ -4,7 +4,6 @@ import argparse
 import json
 import logging
 import platform
-import socket
 import time
 import uuid
 import warnings
@@ -27,10 +26,9 @@ import mlx.core as mx
 from huggingface_hub import scan_cache_dir
 
 from ._version import __version__
-from .generate import stream_generate
-from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
+from .models.cache import make_prompt_cache
 from .sample_utils import make_logits_processors, make_sampler
-from .utils import common_prefix_len, load
+from .utils import load, stream_generate
 
 
 def get_system_fingerprint():
@@ -116,37 +114,10 @@ def convert_chat(messages: List[dict], role_mapping: Optional[dict] = None):
     return prompt.rstrip()
 
 
-def process_message_content(messages):
-    """
-    Convert message content to a format suitable for `apply_chat_template`.
-
-    The function operates on messages in place. It converts the 'content' field
-    to a string instead of a list of text fragments.
-
-    Args:
-        message_list (list): A list of dictionaries, where each dictionary may
-          have a 'content' key containing a list of dictionaries with 'type' and
-          'text' keys.
-
-    Raises:
-        ValueError: If the 'content' type is not supported or if 'text' is missing.
-
-    """
-    for message in messages:
-        content = message["content"]
-        if isinstance(content, list):
-            text_fragments = [
-                fragment["text"] for fragment in content if fragment["type"] == "text"
-            ]
-            if len(text_fragments) != len(content):
-                raise ValueError("Only 'text' content type is supported.")
-            message["content"] = "".join(text_fragments)
-
-
 @dataclass
 class PromptCache:
     cache: List[Any] = field(default_factory=list)
-    model_key: Tuple[str, Optional[str]] = ("", None, None)
+    model_key: Tuple[str, Optional[str]] = ("", None)
     tokens: List[int] = field(default_factory=list)
 
 
@@ -157,11 +128,10 @@ class ModelProvider:
         self.model_key = None
         self.model = None
         self.tokenizer = None
-        self.draft_model = None
 
         # Preload the default model if it is provided
         if self.cli_args.model is not None:
-            self.load("default_model", draft_model_path="default_model")
+            self.load("default_model")
 
     def _validate_model_path(self, model_path: str):
         model_path = Path(model_path)
@@ -171,15 +141,14 @@ class ModelProvider:
             )
 
     # Added in adapter_path to load dynamically
-    def load(self, model_path, adapter_path=None, draft_model_path=None):
-        if self.model_key == (model_path, adapter_path, draft_model_path):
+    def load(self, model_path, adapter_path=None):
+        if self.model_key == (model_path, adapter_path):
             return self.model, self.tokenizer
 
         # Remove the old model if it exists.
         self.model = None
         self.tokenizer = None
         self.model_key = None
-        self.draft_model = None
 
         # Building tokenizer_config
         tokenizer_config = {
@@ -188,12 +157,7 @@ class ModelProvider:
         if self.cli_args.chat_template:
             tokenizer_config["chat_template"] = self.cli_args.chat_template
 
-        if model_path == "default_model":
-            if self.cli_args.model is None:
-                raise ValueError(
-                    "A model path has to be given as a CLI "
-                    "argument or in the HTTP request"
-                )
+        if model_path == "default_model" and self.cli_args.model is not None:
             model, tokenizer = load(
                 self.cli_args.model,
                 adapter_path=(
@@ -211,30 +175,10 @@ class ModelProvider:
             if tokenizer.chat_template is None:
                 tokenizer.chat_template = tokenizer.default_chat_template
 
-        self.model_key = (model_path, adapter_path, draft_model_path)
+        self.model_key = (model_path, adapter_path)
         self.model = model
         self.tokenizer = tokenizer
 
-        def validate_draft_tokenizer(draft_tokenizer):
-            # Check if tokenizers are compatible
-            if draft_tokenizer.vocab_size != tokenizer.vocab_size:
-                logging.warning(
-                    "Draft model tokenizer does not match model tokenizer. "
-                    "Speculative decoding may not work as expected."
-                )
-
-        # Load draft model if specified
-        if (
-            draft_model_path == "default_model"
-            and self.cli_args.draft_model is not None
-        ):
-            self.draft_model, draft_tokenizer = load(self.cli_args.draft_model)
-            validate_draft_tokenizer(draft_tokenizer)
-
-        elif draft_model_path is not None and draft_model_path != "default_model":
-            self._validate_model_path(draft_model_path)
-            self.draft_model, draft_tokenizer = load(draft_model_path)
-            validate_draft_tokenizer(draft_tokenizer)
         return self.model, self.tokenizer
 
 
@@ -306,8 +250,6 @@ class APIHandler(BaseHTTPRequestHandler):
         self.stream = self.body.get("stream", False)
         self.stream_options = self.body.get("stream_options", None)
         self.requested_model = self.body.get("model", "default_model")
-        self.requested_draft_model = self.body.get("draft_model", "default_model")
-        self.num_draft_tokens = self.body.get("num_draft_tokens", 3)
         self.adapter = self.body.get("adapters", None)
         self.max_tokens = self.body.get("max_completion_tokens", None)
         if self.max_tokens is None:
@@ -316,17 +258,14 @@ class APIHandler(BaseHTTPRequestHandler):
         self.top_p = self.body.get("top_p", 1.0)
         self.repetition_penalty = self.body.get("repetition_penalty", 1.0)
         self.repetition_context_size = self.body.get("repetition_context_size", 20)
-        self.xtc_probability = self.body.get("xtc_probability", 0.0)
-        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.validate_model_parameters()
+
         # Load the model if needed
         try:
             self.model, self.tokenizer = self.model_provider.load(
-                self.requested_model,
-                self.adapter,
-                self.requested_draft_model,
+                self.requested_model, self.adapter
             )
         except:
             self._set_completion_headers(404)
@@ -395,15 +334,7 @@ class APIHandler(BaseHTTPRequestHandler):
                 self.logit_bias = {int(k): v for k, v in self.logit_bias.items()}
             except ValueError:
                 raise ValueError("logit_bias must be a dict of int to float")
-        if not (
-            isinstance(self.xtc_probability, float)
-            and 0.00 <= self.xtc_probability <= 1.00
-        ):
-            raise ValueError(f"xtc_probability must be a float between 0.00 and 1.00")
-        if not (
-            isinstance(self.xtc_threshold, float) and 0.00 <= self.xtc_threshold <= 0.50
-        ):
-            raise ValueError(f"xtc_threshold must be a float between 0.00 and 0.5")
+
         if not isinstance(self.requested_model, str):
             raise ValueError("model must be a string")
         if self.adapter is not None and not isinstance(self.adapter, str):
@@ -492,84 +423,18 @@ class APIHandler(BaseHTTPRequestHandler):
 
         return response
 
-    def reset_prompt_cache(self, prompt):
-        """Resets the prompt cache and associated state.
-
-        Args:
-            prompt (List[int]): The tokenized new prompt which will populate the
-                reset cache.
-        """
-        logging.debug(f"*** Resetting cache. ***")
-        self.prompt_cache.model_key = self.model_provider.model_key
-        self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
-        if self.model_provider.draft_model is not None:
-            self.prompt_cache.cache += make_prompt_cache(
-                self.model_provider.draft_model
-            )
-        self.prompt_cache.tokens = list(prompt)  # Cache the new prompt fully
-
     def get_prompt_cache(self, prompt):
-        """
-        Determines the portion of the prompt that needs processing by comparing
-        it to the cached prompt and attempting to reuse the common prefix.
-
-        This function updates the internal prompt cache state (tokens and model cache)
-        based on the comparison. If a common prefix exists, it attempts to trim
-        the model cache (if supported) to match the common prefix length, avoiding
-        recomputation.
-
-        Args:
-            prompt (List[int]): The tokenized new prompt.
-
-        Returns:
-            List[int]: The suffix of the prompt that actually needs to be processed
-                       by the model. This will be the full prompt if the cache is
-                       reset or cannot be effectively used.
-        """
         cache_len = len(self.prompt_cache.tokens)
-        prompt_len = len(prompt)
-        com_prefix_len = common_prefix_len(self.prompt_cache.tokens, prompt)
-
-        # Condition 1: Model changed or no common prefix at all. Reset cache.
         if (
             self.prompt_cache.model_key != self.model_provider.model_key
-            or com_prefix_len == 0
+            or cache_len >= len(prompt)
+            or self.prompt_cache.tokens != prompt[:cache_len]
         ):
-            self.reset_prompt_cache(prompt)
-
-        # Condition 2: Common prefix exists and matches cache length. Process suffix.
-        elif com_prefix_len == cache_len:
-            logging.debug(
-                f"*** Cache is prefix of prompt (cache_len: {cache_len}, prompt_len: {prompt_len}). Processing suffix. ***"
-            )
-            prompt = prompt[com_prefix_len:]
-            self.prompt_cache.tokens.extend(prompt)
-
-        # Condition 3: Common prefix exists but is shorter than cache length. Attempt trim.
-        elif com_prefix_len < cache_len:
-            logging.debug(
-                f"*** Common prefix ({com_prefix_len}) shorter than cache ({cache_len}). Attempting trim. ***"
-            )
-
-            if can_trim_prompt_cache(self.prompt_cache.cache):
-                num_to_trim = cache_len - com_prefix_len
-                logging.debug(f"    Trimming {num_to_trim} tokens from cache.")
-                trim_prompt_cache(self.prompt_cache.cache, num_to_trim)
-                self.prompt_cache.tokens = self.prompt_cache.tokens[:com_prefix_len]
-                prompt = prompt[com_prefix_len:]
-                self.prompt_cache.tokens.extend(prompt)
-            else:
-                logging.debug(f"    Cache cannot be trimmed. Resetting cache.")
-                self.reset_prompt_cache(prompt)
-
-        # This case should logically not be reached if com_prefix_len <= cache_len
+            self.prompt_cache.model_key = self.model_provider.model_key
+            self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
         else:
-            logging.error(
-                f"Unexpected cache state: com_prefix_len ({com_prefix_len}) > cache_len ({cache_len}). Resetting cache."
-            )
-            self.reset_prompt_cache(prompt)
-
-        logging.debug(f"Returning {len(prompt)} tokens for processing.")
+            prompt = prompt[cache_len:]
+        self.prompt_cache.tokens.extend(prompt)
         return prompt
 
     def handle_completion(
@@ -600,22 +465,10 @@ class APIHandler(BaseHTTPRequestHandler):
 
         text = ""
         tic = time.perf_counter()
-        sampler = make_sampler(
-            self.temperature,
-            top_p=self.top_p,
-            xtc_probability=self.xtc_probability,
-            xtc_threshold=self.xtc_threshold,
-            xtc_special_tokens=[
-                self.tokenizer.eos_token_id,
-                self.tokenizer.encode("\n"),
-            ],
-        )
+        sampler = make_sampler(self.temperature, top_p=self.top_p)
         logits_processors = make_logits_processors(
-            self.logit_bias,
-            self.repetition_penalty,
-            self.repetition_context_size,
+            self.logit_bias, self.repetition_penalty, self.repetition_context_size
         )
-
         for gen_response in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
@@ -624,8 +477,6 @@ class APIHandler(BaseHTTPRequestHandler):
             sampler=sampler,
             logits_processors=logits_processors,
             prompt_cache=self.prompt_cache.cache,
-            draft_model=self.model_provider.draft_model,
-            num_draft_tokens=self.num_draft_tokens,
         ):
             segment = gen_response.text
             text += segment
@@ -740,10 +591,8 @@ class APIHandler(BaseHTTPRequestHandler):
         self.request_id = f"chatcmpl-{uuid.uuid4()}"
         self.object_type = "chat.completion.chunk" if self.stream else "chat.completion"
         if self.tokenizer.chat_template:
-            messages = body["messages"]
-            process_message_content(messages)
             prompt = self.tokenizer.apply_chat_template(
-                messages,
+                body["messages"],
                 body.get("tools", None),
                 add_generation_prompt=True,
             )
@@ -784,20 +633,10 @@ class APIHandler(BaseHTTPRequestHandler):
         self._set_completion_headers(200)
         self.end_headers()
 
-        files = ["config.json", "model.safetensors.index.json", "tokenizer_config.json"]
-
-        def probably_mlx_lm(repo):
-            if repo.repo_type != "model":
-                return False
-            if "main" not in repo.refs:
-                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)
-
         # Scan the cache directory for downloaded mlx models
         hf_cache_info = scan_cache_dir()
         downloaded_models = [
-            repo for repo in hf_cache_info.repos if probably_mlx_lm(repo)
+            repo for repo in hf_cache_info.repos if "mlx" in repo.repo_id
         ]
 
         # Create a list of available models
@@ -826,10 +665,6 @@ def run(
 ):
     server_address = (host, port)
     prompt_cache = PromptCache()
-    infos = socket.getaddrinfo(
-        *server_address, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE
-    )
-    server_class.address_family, _, _, _, server_address = next(iter(infos))
     httpd = server_class(
         server_address,
         lambda *args, **kwargs: handler_class(
@@ -872,12 +707,6 @@ def main():
         default=8080,
         help="Port for the HTTP server (default: 8080)",
     )
-    parser.add_argument(
-        "--draft-model",
-        type=str,
-        help="A model to be used for speculative decoding.",
-        default=None,
-    )
     parser.add_argument(
         "--trust-remote-code",
         action="store_true",
@@ -890,6 +719,13 @@ def main():
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Set the logging level (default: INFO)",
     )
+    parser.add_argument(
+        "--cache-limit-gb",
+        type=int,
+        default=None,
+        help="Set the MLX cache limit in GB",
+        required=False,
+    )
     parser.add_argument(
         "--chat-template",
         type=str,
@@ -908,12 +744,13 @@ def main():
         level=getattr(logging, args.log_level.upper(), None),
         format="%(asctime)s - %(levelname)s - %(message)s",
     )
+
+    if args.cache_limit_gb is not None:
+        logging.debug(f"Setting cache limit to {args.cache_limit_gb} GB")
+        mx.metal.set_cache_limit(args.cache_limit_gb * 1024 * 1024 * 1024)
+
     run(args.host, args.port, ModelProvider(args))
 
 
 if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.server...` directly is deprecated."
-        " Use `mlx_lm.server...` or `python -m mlx_lm server ...` instead."
-    )
     main()

mlx_lm/tokenizer_utils.py

@@ -1,7 +1,5 @@
 import json
 from functools import partial
-from json import JSONDecodeError
-from typing import List
 
 from transformers import AutoTokenizer
 
@@ -342,9 +340,7 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
-def load_tokenizer(
-    model_path, tokenizer_config_extra={}, return_tokenizer=True, eos_token_ids=None
-):
+def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -355,12 +351,8 @@ def load_tokenizer(
 
     tokenizer_file = model_path / "tokenizer.json"
     if tokenizer_file.exists():
-        with open(tokenizer_file, "r", encoding="utf-8") as fid:
-            try:
-                tokenizer_content = json.load(fid)
-            except JSONDecodeError as e:
-                raise JSONDecodeError("Failed to parse tokenizer.json", e.doc, e.pos)
-
+        with open(tokenizer_file, "r") as fid:
+            tokenizer_content = json.load(fid)
         if "decoder" in tokenizer_content:
             if _is_spm_decoder(tokenizer_content["decoder"]):
                 detokenizer_class = SPMStreamingDetokenizer
@@ -371,17 +363,8 @@ def load_tokenizer(
 
     if isinstance(eos_token_ids, int):
         eos_token_ids = [eos_token_ids]
-
-    if return_tokenizer:
-        return TokenizerWrapper(
-            AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
-            detokenizer_class,
-            eos_token_ids=eos_token_ids,
-        )
-    else:
-        return detokenizer_class
-
-
-def no_bos_or_eos(sequence: List, bos: int, eos: int) -> List:
-    removed_bos = sequence if sequence[0] != bos else sequence[1:]
-    return removed_bos[:-1] if removed_bos[-1] == eos else removed_bos
+    return TokenizerWrapper(
+        AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
+        detokenizer_class,
+        eos_token_ids=eos_token_ids,
+    )

mlx_lm/tuner/datasets.py

@@ -1,12 +1,11 @@
 import json
-import types
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Dict, List
 
 from transformers import PreTrainedTokenizer
 
 
-class TextDataset:
+class Dataset:
     """
     Light-weight wrapper to hold a dataset.
     """
@@ -17,15 +16,10 @@ class TextDataset:
         tokenizer: PreTrainedTokenizer,
         text_key: str = "text",
     ):
-        self._data = data
-        self.tokenizer = tokenizer
-        self.text_key = text_key
-
-    def process(self, d):
-        d = self.tokenizer.encode(d[self.text_key])
-        if d[-1] != self.tokenizer.eos_token_id:
-            d.append(self.tokenizer.eos_token_id)
-        return d
+        self._data = [tokenizer.encode(d[text_key]) for d in data]
+        for d in self._data:
+            if d[-1] != tokenizer.eos_token_id:
+                d.append(tokenizer.eos_token_id)
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -40,28 +34,14 @@ class ChatDataset:
     https://platform.openai.com/docs/guides/fine-tuning/example-format
     """
 
-    def __init__(
-        self,
-        data: List[Dict[str, str]],
-        tokenizer: PreTrainedTokenizer,
-        chat_key: str = "messages",
-        mask_prompt: bool = False,
-    ):
-        self._data = data
-        self.chat_key = chat_key
-        self.mask_prompt = mask_prompt
-        self.tokenizer = tokenizer
-
-    def process(self, d):
-        messages = d[self.chat_key]
-        tools = d.get("tools", None)
-        tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
-        if self.mask_prompt:
-            messages = messages[:-1]
-            offset = len(self.tokenizer.apply_chat_template(messages, tools=tools))
-            return (tokens, offset)
-        else:
-            return tokens
+    def __init__(self, data: List[Dict[str, str]], tokenizer: PreTrainedTokenizer):
+        self._data = [
+            tokenizer.apply_chat_template(
+                d["messages"],
+                tools=d.get("tools", None),
+            )
+            for d in data
+        ]
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -81,32 +61,18 @@ class CompletionsDataset:
         self,
         data: List[Dict[str, str]],
         tokenizer: PreTrainedTokenizer,
-        prompt_key: str,
-        completion_key: str,
-        mask_prompt: bool,
+        prompt_key: str = "prompt",
+        completion_key: str = "completion",
     ):
-        self._data = data
-        self.prompt_key = prompt_key
-        self.completion_key = completion_key
-        self.mask_prompt = mask_prompt
-        self.tokenizer = tokenizer
-
-    def process(self, d):
-        tokens = self.tokenizer.apply_chat_template(
-            [
-                {"role": "user", "content": d[self.prompt_key]},
-                {"role": "assistant", "content": d[self.completion_key]},
-            ],
-        )
-        if self.mask_prompt:
-            offset = len(
-                self.tokenizer.apply_chat_template(
-                    [{"role": "user", "content": d[self.prompt_key]}]
-                )
+        self._data = [
+            tokenizer.apply_chat_template(
+                [
+                    {"role": "user", "content": d[prompt_key]},
+                    {"role": "assistant", "content": d[completion_key]},
+                ],
             )
-            return (tokens, offset)
-
-        return tokens
+            for d in data
+        ]
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -115,68 +81,15 @@ class CompletionsDataset:
         return len(self._data)
 
 
-class ConcatenatedDataset:
-    def __init__(self, data: List[Any]):
-        self._data = data
-        self._len = sum(len(d) for d in self._data)
-
-    def __getitem__(self, idx: int):
-        for data_idx, data in enumerate(self._data):
-            j = idx - len(data)
-            if j < 0:
-                break
-            idx = j
-        datum = data[idx]
-        datum["_dataset"] = data_idx
-        return datum
-
-    def process(self, d):
-        return self._data[d["_dataset"]].process(d)
-
-    def __len__(self):
-        return self._len
-
-
-class CacheDataset:
-    def __init__(self, data: Any):
-        self._data = data
-        self._proc_data = [None] * len(data)
-
-    def itemlen(self, idx: int):
-        return len(self._data[idx])
-
-    def __getitem__(self, idx: int):
-        if self._proc_data[idx] is None:
-            self._proc_data[idx] = self._data.process(self._data[idx])
-        return self._proc_data[idx]
-
-    def __len__(self):
-        return len(self._data)
-
-
-def create_dataset(
-    data,
-    tokenizer: PreTrainedTokenizer,
-    config,
-):
-    mask_prompt = getattr(config, "mask_prompt", False)
-    prompt_feature = getattr(config, "prompt_feature", "prompt")
-    text_feature = getattr(config, "text_feature", "text")
-    completion_feature = getattr(config, "completion_feature", "completion")
-    chat_feature = getattr(config, "chat_feature", "messages")
+def create_dataset(data, tokenizer: PreTrainedTokenizer):
     sample = data[0]
-    if prompt_feature in sample and completion_feature in sample:
-        return CompletionsDataset(
-            data, tokenizer, prompt_feature, completion_feature, mask_prompt
-        )
-    elif chat_feature in sample:
-        return ChatDataset(
-            data, tokenizer, chat_key=chat_feature, mask_prompt=mask_prompt
-        )
-    elif text_feature in sample:
-        if mask_prompt:
-            raise ValueError("Prompt masking not supported for text dataset.")
-        return TextDataset(data, tokenizer, text_key=text_feature)
+
+    if "messages" in sample:
+        return ChatDataset(data, tokenizer)
+    elif "prompt" in sample and "completion" in sample:
+        return CompletionsDataset(data, tokenizer)
+    elif "text" in sample:
+        return Dataset(data, tokenizer)
     else:
         raise ValueError(
             "Unsupported data format, check the supported formats here:\n"
@@ -184,28 +97,20 @@ def create_dataset(
         )
 
 
-def load_local_dataset(
-    data_path: Path,
-    tokenizer: PreTrainedTokenizer,
-    config,
-):
+def load_local_dataset(data_path: Path, tokenizer: PreTrainedTokenizer):
     def load_subset(path):
         if not path.exists():
             return []
         with open(path, "r") as fid:
             data = [json.loads(l) for l in fid]
-        return create_dataset(data, tokenizer, config)
+        return create_dataset(data, tokenizer)
 
     names = ("train", "valid", "test")
     train, valid, test = [load_subset(data_path / f"{n}.jsonl") for n in names]
     return train, valid, test
 
 
-def load_hf_dataset(
-    data_id: str,
-    tokenizer: PreTrainedTokenizer,
-    config,
-):
+def load_hf_dataset(data_id: str, tokenizer: PreTrainedTokenizer):
     from datasets import exceptions, load_dataset
 
     try:
@@ -214,11 +119,7 @@ def load_hf_dataset(
         names = ("train", "valid", "test")
 
         train, valid, test = [
-            (
-                create_dataset(dataset[n], tokenizer, config)
-                if n in dataset.keys()
-                else []
-            )
+            create_dataset(dataset[n], tokenizer) if n in dataset.keys() else []
             for n in names
         ]
 
@@ -231,61 +132,42 @@ def load_hf_dataset(
 def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
     import datasets
 
-    def create_hf_dataset(dataset_name, config, split, hf_config):
+    hf_args = args.hf_dataset
+    dataset_name = hf_args["name"]
+    print(f"Loading Hugging Face dataset {dataset_name}.")
+    text_feature = hf_args.get("text_feature")
+    prompt_feature = hf_args.get("prompt_feature")
+    completion_feature = hf_args.get("completion_feature")
+
+    def create_hf_dataset(split: str = None):
         ds = datasets.load_dataset(
             dataset_name,
             split=split,
-            **hf_config,
+            **hf_args.get("config", {}),
         )
-        return create_dataset(ds, tokenizer, config)
-
-    dataset_collection = args.hf_dataset
-    if isinstance(dataset_collection, dict):
-        dataset_collection = [dataset_collection]
-
-    collection = []
-    for ds in dataset_collection:
-        ds_path = ds["path"]
-        print(f"Loading Hugging Face dataset {ds_path}.")
-        ds["mask_prompt"] = getattr(args, "mask_prompt", False)
-        config = types.SimpleNamespace(**ds)
-        hf_config = ds.get("config", {})
-        if args.train:
-            train_split = ds.get("train_split", "train[:80%]")
-            valid_split = ds.get("valid_split", "train[-10%:]")
-            train = create_hf_dataset(
-                ds_path,
-                config,
-                train_split,
-                hf_config,
-            )
-            valid = create_hf_dataset(
-                ds_path,
-                config,
-                valid_split,
-                hf_config,
-            )
+        if prompt_feature and completion_feature:
+            return CompletionsDataset(ds, tokenizer, prompt_feature, completion_feature)
+        elif text_feature:
+            return Dataset(train_ds, tokenizer, text_key=text_feature)
         else:
-            train, valid = [], []
-
-        if args.test:
-            test_split = ds.get("test_split")
-            test = create_hf_dataset(
-                ds_path,
-                config,
-                test_split,
-                hf_config,
+            raise ValueError(
+                "Specify either a prompt and completion feature or a text "
+                "feature for the Hugging Face dataset."
             )
-        else:
-            test = []
 
-        collection.append((train, valid, test))
+    if args.train:
+        train_split = hf_args.get("train_split", "train[:80%]")
+        valid_split = hf_args.get("valid_split", "train[-10%:]")
+        train = create_hf_dataset(split=train_split)
+        valid = create_hf_dataset(split=valid_split)
+    else:
+        train, valid = [], []
+    if args.test:
+        test = create_hf_dataset(split=hf_args.get("test_split"))
+    else:
+        test = []
 
-    if len(collection) == 1:
-        return collection[0]
-
-    # Otherwise concatenate them
-    return tuple(map(ConcatenatedDataset, zip(*collection)))
+    return train, valid, test
 
 
 def load_dataset(args, tokenizer: PreTrainedTokenizer):
@@ -294,10 +176,10 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
     else:
         data_path = Path(args.data)
         if data_path.exists():
-            train, valid, test = load_local_dataset(data_path, tokenizer, args)
+            train, valid, test = load_local_dataset(data_path, tokenizer)
         else:
             print(f"Loading Hugging Face dataset {args.data}.")
-            train, valid, test = load_hf_dataset(args.data, tokenizer, args)
+            train, valid, test = load_hf_dataset(args.data, tokenizer)
 
     if args.train and len(train) == 0:
         raise ValueError(

mlx_lm/tuner/trainer.py

@@ -4,18 +4,14 @@ import glob
 import shutil
 import time
 from dataclasses import dataclass, field
-from functools import partial
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import Union
 
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten
-from transformers import PreTrainedTokenizer
-
-from .datasets import CacheDataset
 
 
 def grad_checkpoint(layer):
@@ -67,34 +63,22 @@ class TrainingArgs:
     )
 
 
-def default_loss(model, batch, lengths):
-    inputs = batch[:, :-1]
-    targets = batch[:, 1:]
-
+def default_loss(model, inputs, targets, lengths):
     logits = model(inputs)
+    logits = logits.astype(mx.float32)
 
-    steps = mx.arange(1, targets.shape[1] + 1)
-    mask = mx.logical_and(steps >= lengths[:, 0:1], steps <= lengths[:, 1:])
+    length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]
 
-    ce = nn.losses.cross_entropy(logits, targets) * mask
-    ntoks = mask.sum()
-    ce = ce.astype(mx.float32).sum() / ntoks
+    ce = nn.losses.cross_entropy(logits, targets) * length_mask
+    ntoks = length_mask.sum()
+    ce = ce.sum() / ntoks
 
     return ce, ntoks
 
 
-def iterate_batches(
-    dataset,
-    batch_size,
-    max_seq_length,
-    train=False,
-):
+def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False):
     # Sort by length:
-    if isinstance(dataset, CacheDataset):
-        len_fn = lambda idx: dataset.itemlen(idx)
-    else:
-        len_fn = lambda idx: len(dataset[idx])
-    idx = sorted(range(len(dataset)), key=len_fn)
+    idx = sorted(range(len(dataset)), key=lambda idx: len(dataset[idx]))
     if len(dataset) < batch_size:
         raise ValueError(
             f"Dataset must have at least batch_size={batch_size}"
@@ -117,10 +101,6 @@ def iterate_batches(
         indices = np.random.permutation(len(batch_idx))
         for i in indices:
             batch = [dataset[j] for j in batch_idx[i]]
-            if len(batch[0]) == 2:
-                batch, offsets = zip(*batch)
-            else:
-                offsets = [0] * len(batch)
             lengths = [len(x) for x in batch]
             if max(lengths) > max_seq_length:
                 print(
@@ -129,9 +109,9 @@ def iterate_batches(
                     "Consider pre-splitting your data to save memory."
                 )
 
-            # Pad to one plus nearest multiple of pad_to or the maximum length
-            pad_to = 32
-            max_length_in_batch = 1 + pad_to * ((max(lengths) + pad_to - 1) // pad_to)
+            # Pad to the nearest multiple of 8 or the maximum length
+            pad_to = 8
+            max_length_in_batch = pad_to * ((max(lengths) + pad_to - 1) // pad_to)
             max_length_in_batch = min(max_length_in_batch, max_seq_length)
 
             batch_arr = np.zeros((batch_size // step, max_length_in_batch), np.int32)
@@ -143,7 +123,8 @@ def iterate_batches(
                     truncated_length  # Update lengths to match truncated lengths
                 )
             batch = mx.array(batch_arr)
-            yield batch, mx.array(list(zip(offsets, lengths)))
+
+            yield batch[:, :-1], batch[:, 1:], mx.array(lengths)
 
         if not train:
             break
@@ -152,15 +133,15 @@ def iterate_batches(
 def evaluate(
     model,
     dataset,
+    tokenizer,
     batch_size,
     num_batches,
     max_seq_length=2048,
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
 ):
-    model.eval()
-    all_losses = mx.array(0.0)
-    ntokens = mx.array(0)
+    all_losses = 0
+    ntokens = 0
 
     index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
 
@@ -168,6 +149,7 @@ def evaluate(
         index_iterator,
         iterate_batches(
             dataset=dataset,
+            tokenizer=tokenizer,
             batch_size=batch_size,
             max_seq_length=max_seq_length,
         ),
@@ -177,8 +159,8 @@ def evaluate(
         ntokens += toks
         mx.eval(all_losses, ntokens)
 
-    all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
-    ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
+    all_losses = mx.distributed.all_sum(all_losses)
+    ntokens = mx.distributed.all_sum(ntokens)
 
     return (all_losses / ntokens).item()
 
@@ -196,6 +178,7 @@ class TrainingCallback:
 
 def train(
     model,
+    tokenizer,
     optimizer,
     train_dataset,
     val_dataset,
@@ -204,7 +187,6 @@ def train(
     iterate_batches: callable = iterate_batches,
     training_callback: TrainingCallback = None,
 ):
-    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()
@@ -215,9 +197,8 @@ def train(
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
 
-    state = [model.state, optimizer.state, mx.random.state]
+    state = [model.state, optimizer.state]
 
-    @partial(mx.compile, inputs=state, outputs=state)
     def step(batch):
         # Forward and backward pass
         (lvalue, toks), grad = loss_value_and_grad(model, *batch)
@@ -232,38 +213,37 @@ def train(
 
     loss_value_and_grad = nn.value_and_grad(model, loss)
 
-    model.train()
     losses = 0
     n_tokens = 0
     steps = 0
     trained_tokens = 0
-    train_time = 0
     # Main training loop
+    start = time.perf_counter()
     for it, batch in zip(
         range(1, args.iters + 1),
         iterate_batches(
             dataset=train_dataset,
+            tokenizer=tokenizer,
             batch_size=args.batch_size,
             max_seq_length=args.max_seq_length,
             train=True,
         ),
     ):
-        tic = time.perf_counter()
         # Report validation loss if needed, the first validation loss
         # is always measured before any training.
         if it == 1 or it % args.steps_per_eval == 0 or it == args.iters:
-            tic = time.perf_counter()
+            stop = time.perf_counter()
             val_loss = evaluate(
                 model=model,
                 dataset=val_dataset,
                 loss=loss,
+                tokenizer=tokenizer,
                 batch_size=args.batch_size,
                 num_batches=args.val_batches,
                 max_seq_length=args.max_seq_length,
                 iterate_batches=iterate_batches,
             )
-            model.train()
-            val_time = time.perf_counter() - tic
+            val_time = time.perf_counter() - stop
             if rank == 0:
                 print(
                     f"Iter {it}: "
@@ -280,25 +260,26 @@ def train(
                 }
                 training_callback.on_val_loss_report(val_info)
 
-            tic = time.perf_counter()
+            start = time.perf_counter()
 
         lvalue, toks = step(batch)
         losses += lvalue
         n_tokens += toks
         steps += 1
         mx.eval(state, losses, n_tokens)
-        train_time += time.perf_counter() - tic
 
         # Report training loss if needed
         if it % args.steps_per_report == 0 or it == args.iters:
-            train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
-            train_loss /= steps * world_size
-            n_tokens = mx.distributed.all_sum(n_tokens, stream=mx.cpu).item()
+            stop = time.perf_counter()
+
+            train_loss = mx.distributed.all_sum(losses).item()
+            train_loss /= steps * mx.distributed.init().size()
+            n_tokens = mx.distributed.all_sum(n_tokens).item()
             learning_rate = optimizer.learning_rate.item()
-            it_sec = args.steps_per_report / train_time
-            tokens_sec = float(n_tokens) / train_time
+            it_sec = args.steps_per_report / (stop - start)
+            tokens_sec = float(n_tokens) / (stop - start)
             trained_tokens += n_tokens
-            peak_mem = mx.get_peak_memory() / 1e9
+            peak_mem = mx.metal.get_peak_memory() / 1e9
             if rank == 0:
                 print(
                     f"Iter {it}: Train loss {train_loss:.3f}, "
@@ -325,10 +306,10 @@ def train(
             losses = 0
             n_tokens = 0
             steps = 0
-            train_time = 0
+            start = time.perf_counter()
 
         # Save adapter weights
-        if it % args.steps_per_save == 0 and rank == 0:
+        if it % args.steps_per_save == 0:
             adapter_weights = dict(tree_flatten(model.trainable_parameters()))
             mx.save_safetensors(str(args.adapter_file), adapter_weights)
             checkpoint = (
@@ -341,7 +322,6 @@ def train(
             )
 
     # Save final weights
-    if rank == 0:
-        adapter_weights = dict(tree_flatten(model.trainable_parameters()))
-        mx.save_safetensors(str(args.adapter_file), adapter_weights)
-        print(f"Saved final weights to {args.adapter_file}.")
+    adapter_weights = dict(tree_flatten(model.trainable_parameters()))
+    mx.save_safetensors(str(args.adapter_file), adapter_weights)
+    print(f"Saved final weights to {args.adapter_file}.")

mlx_lm/tuner/utils.py

@@ -52,6 +52,11 @@ def linear_to_lora_layers(
         use_dora (bool): If True, uses DoRA instead of LoRA.
           Default: ``False``
     """
+    if num_layers > len(model.layers):
+        raise ValueError(
+            f"Requested {num_layers} LoRA layers "
+            f"but the model only has {len(model.layers)} layers."
+        )
 
     def to_lora(layer):
         if isinstance(layer, (nn.Linear, nn.QuantizedLinear)):
@@ -84,29 +89,17 @@ def linear_to_lora_layers(
         "mixtral",
         "nemotron",
         "stablelm",
-        "hunyuan",
         "qwen2",
         "qwen2_moe",
-        "qwen3",
-        "qwen3_moe",
         "phimoe",
         "gemma",
         "gemma2",
-        "gemma3",
-        "gemma3_text",
-        "granite",
-        "helium",
         "starcoder2",
         "cohere",
         "cohere2",
         "minicpm",
-        "minicpm3",
         "deepseek",
         "olmo2",
-        "olmoe",
-        "internlm3",
-        "glm4",
-        "mimo",
     ]:
         keys = set(["self_attn.q_proj", "self_attn.v_proj"])
         if model.model_type in ["mixtral", "phimoe"]:
@@ -114,8 +107,6 @@ def linear_to_lora_layers(
         if model.model_type == "qwen2_moe":
             keys.add("mlp.gate")
             keys.add("mlp.shared_expert_gate")
-        if model.model_type in ["olmoe", "qwen3_moe"]:
-            keys.add("mlp.gate")
 
     elif model.model_type == "gpt_bigcode":
         keys = set(["attn.c_attn"])
@@ -135,7 +126,7 @@ def linear_to_lora_layers(
         keys = set(["norm_attn_norm.attn.Wqkv", "ffn.router.layer"])
     elif model.model_type == "internlm2":
         keys = set(["attention.wqkv", "attention.wo"])
-    elif model.model_type == "deepseek_v2" or model.model_type == "minicpm3":
+    elif model.model_type == "deepseek_v2":
         keys = set(
             [
                 "self_attn.q_proj",
@@ -159,7 +150,7 @@ def linear_to_lora_layers(
     else:
         raise ValueError(f"Lora does not support {model.model_type}")
 
-    for l in model.layers[-max(num_layers, 0) :]:
+    for l in model.layers[-min(num_layers, 0) :]:
         lora_layers = [(k, to_lora(m)) for k, m in l.named_modules() if k in keys]
         if lora_layers:
             l.update_modules(tree_unflatten(lora_layers))

mlx_lm/upload.py

@@ -1,22 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-
-from .utils import upload_to_hub
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Upload a model to the Hugging Face Hub"
-    )
-
-    parser.add_argument(
-        "--path", type=str, default="mlx_model", help="Path to the MLX model."
-    )
-    parser.add_argument(
-        "--upload-repo",
-        help="The Hugging Face repo to upload the model to.",
-        type=str,
-    )
-    args = parser.parse_args()
-    upload_to_hub(args.path, args.upload_repo)

mlx_lm/utils.py

@@ -1,41 +1,27 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import contextlib
 import copy
 import glob
 import importlib
 import json
 import logging
-import os
 import shutil
+import time
+from dataclasses import dataclass
 from pathlib import Path
 from textwrap import dedent
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Optional,
-    Tuple,
-    Type,
-    Union,
-)
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type, Union
 
 import mlx.core as mx
 import mlx.nn as nn
-
-if os.getenv("MLXLM_USE_MODELSCOPE", "False").lower() == "true":
-    try:
-        from modelscope import snapshot_download
-    except ImportError:
-        raise ImportError(
-            "Please run `pip install modelscope` to activate the ModelScope."
-        )
-else:
-    from huggingface_hub import snapshot_download
-
+from huggingface_hub import snapshot_download
 from mlx.utils import tree_flatten, tree_reduce
 from transformers import PreTrainedTokenizer
 
 # Local imports
+from .models import cache
+from .sample_utils import make_logits_processors, make_sampler
 from .tokenizer_utils import TokenizerWrapper, load_tokenizer
 from .tuner.utils import dequantize as dequantize_model
 from .tuner.utils import load_adapters, nparams
@@ -49,6 +35,9 @@ MODEL_REMAPPING = {
 
 MAX_FILE_SIZE_GB = 5
 
+# A stream on the default device just for generation
+generation_stream = mx.new_stream(mx.default_device())
+
 
 class ModelNotFoundError(Exception):
     def __init__(self, message):
@@ -56,6 +45,68 @@ class ModelNotFoundError(Exception):
         super().__init__(self.message)
 
 
+@dataclass
+class GenerationResponse:
+    """
+    The output of :func:`stream_generate`.
+
+    Args:
+        text (str): The next segment of decoded text. This can be an empty string.
+        token (int): The next token.
+        logprobs (mx.array): A vector of log probabilities.
+        prompt_tokens (int): The number of tokens in the prompt.
+        prompt_tps (float): The prompt processing tokens-per-second.
+        generation_tokens (int): The number of generated tokens.
+        generation_tps (float): The tokens-per-second for generation.
+        peak_memory (float): The peak memory used so far in GB.
+        finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
+    """
+
+    text: str
+    token: int
+    logprobs: mx.array
+    prompt_tokens: int
+    prompt_tps: float
+    generation_tokens: int
+    generation_tps: float
+    peak_memory: float
+    finish_reason: Optional[str] = None
+
+
+@contextlib.contextmanager
+def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
+    """
+    A context manager to temporarily change the wired limit.
+
+    Note, the wired limit should not be changed during an async eval.  If an
+    async eval could be running pass in the streams to synchronize with prior
+    to exiting the context manager.
+    """
+    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-examples/tree/main/llms#large-models"
+        )
+    old_limit = mx.metal.set_wired_limit(max_rec_size)
+    try:
+        yield None
+    finally:
+        if streams is not None:
+            for s in streams:
+                mx.synchronize(s)
+        else:
+            mx.synchronize()
+        mx.metal.set_wired_limit(old_limit)
+
+
 def _get_classes(config: dict):
     """
     Retrieve the model and model args classes based on the configuration.
@@ -102,12 +153,11 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
         Path: The path to the model.
     """
     model_path = Path(path_or_hf_repo)
-
     if not model_path.exists():
         try:
             model_path = Path(
                 snapshot_download(
-                    path_or_hf_repo,
+                    repo_id=path_or_hf_repo,
                     revision=revision,
                     allow_patterns=[
                         "*.json",
@@ -115,9 +165,7 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
                         "*.py",
                         "tokenizer.model",
                         "*.tiktoken",
-                        "tiktoken.model",
                         "*.txt",
-                        "*.jsonl",
                     ],
                 )
             )
@@ -132,6 +180,280 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
     return model_path
 
 
+def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
+    if (
+        kv_bits is not None
+        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
+        and prompt_cache[0].offset > quantized_kv_start
+    ):
+        for i in range(len(prompt_cache)):
+            if isinstance(prompt_cache[i], cache.KVCache):
+                prompt_cache[i] = prompt_cache[i].to_quantized(
+                    group_size=kv_group_size, bits=kv_bits
+                )
+
+
+def generate_step(
+    prompt: mx.array,
+    model: nn.Module,
+    *,
+    max_tokens: int = 256,
+    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
+    max_kv_size: Optional[int] = None,
+    prompt_cache: Optional[Any] = None,
+    prefill_step_size: int = 512,
+    kv_bits: Optional[int] = None,
+    kv_group_size: int = 64,
+    quantized_kv_start: int = 0,
+    prompt_progress_callback: Optional[Callable[int, int]] = None,
+    temp: Optional[float] = None,
+    repetition_penalty: Optional[float] = None,
+    repetition_context_size: Optional[int] = None,
+    top_p: Optional[float] = None,
+    min_p: Optional[float] = None,
+    min_tokens_to_keep: Optional[int] = None,
+) -> Generator[Tuple[mx.array, mx.array], None, None]:
+    """
+    A generator producing token ids based on the given prompt from the model.
+
+    Args:
+        prompt (mx.array): The input prompt.
+        model (nn.Module): The model to use for generation.
+        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
+          generator. Default: ``256``.
+        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
+          token from a vector of log probabilities. Default: ``None``.
+        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed
+          logits. Default: ``None``.
+        max_kv_size (int, optional): Maximum size of the key-value cache. Old
+          entries (except the first 4 tokens) will be overwritten.
+        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
+          provided, the cache will be updated in place.
+        prefill_step_size (int): Step size for processing the prompt.
+        kv_bits (int, optional): Number of bits to use for KV cache quantization.
+          None implies no cache quantization. Default: ``None``.
+        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
+        quantized_kv_start (int): Step to begin using a quantized KV cache.
+           when ``kv_bits`` is non-None. Default: ``0``.
+        prompt_prorgress_callback (Callable[int, int]): A call-back which takes the
+           prompt tokens processed so far and the total number of prompt tokens.
+
+    Yields:
+        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
+    """
+
+    y = prompt
+    tokens = None
+
+    # Create the KV cache for generation
+    if prompt_cache is None:
+        prompt_cache = cache.make_prompt_cache(
+            model,
+            max_kv_size=max_kv_size,
+        )
+    elif len(prompt_cache) != len(model.layers):
+        raise ValueError("Wrong number of layers in the prompt cache.")
+
+    if temp is not None or top_p is not None or min_tokens_to_keep is not None:
+        print(
+            "[Warning] Specifying sampling arguments to ``generate_step`` is "
+            "deprecated. Pass in a ``sampler`` instead."
+        )
+    if repetition_penalty is not None:
+        print(
+            "[Warning] Specifying ``repetition_penalty`` is deprecated. "
+            "Pass in ``logits_processors`` instead."
+        )
+
+    sampler = sampler or make_sampler(
+        temp or 0.0, top_p or 0.0, min_p or 0.0, min_tokens_to_keep or 1
+    )
+    logits_processors = logits_processors or make_logits_processors(
+        None, repetition_penalty, repetition_context_size or 20
+    )
+    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
+
+    def _step(y):
+        with mx.stream(generation_stream):
+            logits = model(y[None], cache=prompt_cache)
+            logits = logits[:, -1, :]
+
+            if logits_processors:
+                nonlocal tokens
+                tokens = mx.concat([tokens, y]) if tokens is not None else y
+
+                for processor in logits_processors:
+                    logits = processor(tokens, logits)
+
+            maybe_quantize_kv_cache(
+                prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+            )
+
+            logprobs = logits - mx.logsumexp(logits, keepdims=True)
+            y = sampler(logprobs)
+            return y, logprobs.squeeze(0)
+
+    with mx.stream(generation_stream):
+        total_prompt_tokens = y.size
+        prompt_processed_tokens = 0
+        while y.size > prefill_step_size:
+            model(y[:prefill_step_size][None], cache=prompt_cache)
+            maybe_quantize_kv_cache(
+                prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+            )
+            mx.eval([c.state for c in prompt_cache])
+            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
+            prompt_processed_tokens += prefill_step_size
+            y = y[prefill_step_size:]
+            mx.metal.clear_cache()
+
+        y, logprobs = _step(y)
+
+    mx.async_eval(y, logprobs)
+    n = 0
+    while True:
+        if n != max_tokens:
+            next_y, next_logprobs = _step(y)
+            mx.async_eval(next_y, next_logprobs)
+        if n == 0:
+            mx.eval(y)
+            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
+        if n == max_tokens:
+            break
+        yield y.item(), logprobs
+        if n % 256 == 0:
+            mx.metal.clear_cache()
+        y, logprobs = next_y, next_logprobs
+        n += 1
+
+
+def stream_generate(
+    model: nn.Module,
+    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
+    prompt: Union[str, mx.array, List[int]],
+    **kwargs,
+) -> Generator[GenerationResponse, None, None]:
+    """
+    A generator producing text based on the given prompt from the model.
+
+    Args:
+        model (nn.Module): The model to use for generation.
+        tokenizer (PreTrainedTokenizer): The tokenizer.
+        prompt (Union[str, mx.array, List[int]]): The input prompt string or integer tokens.
+        kwargs: The remaining options get passed to :func:`generate_step`.
+          See :func:`generate_step` for more details.
+
+    Yields:
+        GenerationResponse: An instance containing the generated text segment and
+            associated metadata. See :class:`GenerationResponse` for details.
+    """
+    if not isinstance(tokenizer, TokenizerWrapper):
+        tokenizer = TokenizerWrapper(tokenizer)
+
+    if not isinstance(prompt, mx.array):
+        if isinstance(prompt, str):
+            # Try to infer if special tokens are needed
+            add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
+                tokenizer.bos_token
+            )
+            prompt = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
+        prompt = mx.array(prompt)
+
+    detokenizer = tokenizer.detokenizer
+
+    with wired_limit(model, [generation_stream]):
+        detokenizer.reset()
+        tic = time.perf_counter()
+        for n, (token, logprobs) in enumerate(generate_step(prompt, model, **kwargs)):
+            if n == 0:
+                prompt_time = time.perf_counter() - tic
+                prompt_tps = prompt.size / prompt_time
+                tic = time.perf_counter()
+            if token in tokenizer.eos_token_ids:
+                break
+
+            detokenizer.add_token(token)
+
+            yield GenerationResponse(
+                text=detokenizer.last_segment,
+                token=token,
+                logprobs=logprobs,
+                prompt_tokens=prompt.size,
+                prompt_tps=prompt_tps,
+                generation_tokens=n + 1,
+                generation_tps=(n + 1) / (time.perf_counter() - tic),
+                peak_memory=mx.metal.get_peak_memory() / 1e9,
+                finish_reason=None,
+            )
+
+        detokenizer.finalize()
+        yield GenerationResponse(
+            text=detokenizer.last_segment,
+            token=token,
+            logprobs=logprobs,
+            prompt_tokens=prompt.size,
+            prompt_tps=prompt_tps,
+            generation_tokens=n + 1,
+            generation_tps=(n + 1) / (time.perf_counter() - tic),
+            peak_memory=mx.metal.get_peak_memory() / 1e9,
+            finish_reason="stop" if token in tokenizer.eos_token_ids else "length",
+        )
+
+
+def generate(
+    model: nn.Module,
+    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
+    prompt: Union[str, List[int]],
+    verbose: bool = False,
+    formatter: Optional[Callable] = None,
+    **kwargs,
+) -> str:
+    """
+    Generate a complete response from the model.
+
+    Args:
+       model (nn.Module): The language model.
+       tokenizer (PreTrainedTokenizer): The tokenizer.
+       prompt (Union[str, List[int]]): The input prompt string or integer tokens.
+       verbose (bool): If ``True``, print tokens and timing information.
+           Default: ``False``.
+       kwargs: The remaining options get passed to :func:`stream_generate`.
+          See :func:`stream_generate` for more details.
+    """
+    if formatter is not None:
+        print(
+            "[Warning] Text formatting is deprecated and no longer used. "
+            "The argument will be removed in a future version."
+        )
+    if verbose:
+        print("=" * 10)
+
+    text = ""
+    for response in stream_generate(model, tokenizer, prompt, **kwargs):
+        if verbose:
+            print(response.text, end="", flush=True)
+        text += response.text
+
+    if verbose:
+        print()
+        print("=" * 10)
+        if len(text) == 0:
+            print("No text generated for this prompt")
+            return
+        print(
+            f"Prompt: {response.prompt_tokens} tokens, "
+            f"{response.prompt_tps:.3f} tokens-per-sec"
+        )
+        print(
+            f"Generation: {response.generation_tokens} tokens, "
+            f"{response.generation_tps:.3f} tokens-per-sec"
+        )
+        print(f"Peak memory: {response.peak_memory:.3f} GB")
+    return text
+
+
 def load_config(model_path: Path) -> dict:
     try:
         with open(model_path / "config.json", "r") as f:
@@ -145,7 +467,6 @@ def load_config(model_path: Path) -> dict:
 def load_model(
     model_path: Path,
     lazy: bool = False,
-    strict: bool = True,
     model_config: dict = {},
     get_model_classes: Callable[[dict], Tuple[Type[nn.Module], Type]] = _get_classes,
 ) -> nn.Module:
@@ -157,8 +478,6 @@ def load_model(
         lazy (bool): If False eval the model parameters to make sure they are
             loaded in memory before returning, otherwise they will be loaded
             when needed. Default: ``False``
-        strict (bool): Whether or not to raise an exception if weights don't
-            match. Default: ``True``
         model_config (dict, optional): Optional configuration parameters for the
             model. Defaults to an empty dictionary.
         get_model_classes (Callable[[dict], Tuple[Type[nn.Module], Type]], optional):
@@ -181,7 +500,7 @@ def load_model(
         # Try weight for back-compat
         weight_files = glob.glob(str(model_path / "weight*.safetensors"))
 
-    if not weight_files and strict:
+    if not weight_files:
         logging.error(f"No safetensors found in {model_path}")
         raise FileNotFoundError(f"No safetensors found in {model_path}")
 
@@ -215,7 +534,7 @@ def load_model(
             class_predicate=class_predicate,
         )
 
-    model.load_weights(list(weights.items()), strict=strict)
+    model.load_weights(list(weights.items()))
 
     if not lazy:
         mx.eval(model.parameters())
@@ -242,7 +561,7 @@ def load(
             Defaults to an empty dictionary.
         adapter_path (str, optional): Path to the LoRA adapters. If provided, applies LoRA layers
             to the model. Default: ``None``.
-        lazy (bool): If ``False`` eval the model parameters to make sure they are
+        lazy (bool): If False eval the model parameters to make sure they are
             loaded in memory before returning, otherwise they will be loaded
             when needed. Default: ``False``
     Returns:
@@ -299,49 +618,29 @@ def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list
     return shards
 
 
-def create_model_card(path: Union[str, Path], hf_path: Union[str, Path]):
-    """
-    Uploads the model to Hugging Face hub.
-
-    Args:
-        path (Union[str, Path]): Local path to the model.
-        hf_path (Union[str, Path]): Path to the original Hugging Face model.
-    """
-    from huggingface_hub import ModelCard
-
-    card = ModelCard.load(hf_path)
-    card.data.library_name = "mlx"
-    card.data.pipeline_tag = "text-generation"
-    if card.data.tags is None:
-        card.data.tags = ["mlx"]
-    elif "mlx" not in card.data.tags:
-        card.data.tags += ["mlx"]
-    card.data.base_model = str(hf_path)
-    card.text = ""
-    card.save(os.path.join(path, "README.md"))
-
-
-def upload_to_hub(path: str, upload_repo: str):
+def upload_to_hub(path: str, upload_repo: str, hf_path: str):
     """
     Uploads the model to Hugging Face hub.
 
     Args:
         path (str): Local path to the model.
         upload_repo (str): Name of the HF repo to upload to.
+        hf_path (str): Path to the original Hugging Face model.
     """
+    import os
+
     from huggingface_hub import HfApi, ModelCard, logging
 
     from . import __version__
 
-    logging.set_verbosity_info()
-    card_path = Path(path) / "README.md"
-    card = ModelCard.load(card_path)
-    hf_path = card.data.base_model
+    card = ModelCard.load(hf_path)
+    card.data.tags = ["mlx"] if card.data.tags is None else card.data.tags + ["mlx"]
+    card.data.base_model = hf_path
     card.text = dedent(
         f"""
         # {upload_repo}
 
-        This model [{upload_repo}](https://huggingface.co/{upload_repo}) was
+        The Model [{upload_repo}](https://huggingface.co/{upload_repo}) was
         converted to MLX format from [{hf_path}](https://huggingface.co/{hf_path})
         using mlx-lm version **{__version__}**.
 
@@ -356,7 +655,7 @@ def upload_to_hub(path: str, upload_repo: str):
 
         model, tokenizer = load("{upload_repo}")
 
-        prompt = "hello"
+        prompt="hello"
 
         if tokenizer.chat_template is not None:
             messages = [{{"role": "user", "content": prompt}}]
@@ -368,7 +667,9 @@ def upload_to_hub(path: str, upload_repo: str):
         ```
         """
     )
-    card.save(card_path)
+    card.save(os.path.join(path, "README.md"))
+
+    logging.set_verbosity_info()
 
     api = HfApi()
     api.create_repo(repo_id=upload_repo, exist_ok=True)
@@ -497,7 +798,6 @@ def save_config(
     """
     # Clean unused keys
     config.pop("_name_or_path", None)
-    config.pop("vision_config", None)
 
     # sort the config for better readability
     config = dict(sorted(config.items()))
@@ -507,50 +807,63 @@ def save_config(
         json.dump(config, fid, indent=4)
 
 
-def save(
-    dst_path: Union[str, Path],
-    src_path: Union[str, Path],
-    weights: Dict[str, mx.array],
-    tokenizer: TokenizerWrapper,
-    config: Dict[str, Any],
-    hf_repo: Optional[str] = None,
-    donate_weights: bool = True,
+def convert(
+    hf_path: str,
+    mlx_path: str = "mlx_model",
+    quantize: bool = False,
+    q_group_size: int = 64,
+    q_bits: int = 4,
+    dtype: str = "float16",
+    upload_repo: str = None,
+    revision: Optional[str] = None,
+    dequantize: bool = False,
+    quant_predicate: Optional[
+        Callable[[str, nn.Module, dict], Union[bool, dict]]
+    ] = None,
 ):
-    src_path = Path(src_path)
-    dst_path = Path(dst_path)
-    save_weights(dst_path, weights, donate_weights=True)
-    save_config(config, config_path=dst_path / "config.json")
-    tokenizer.save_pretrained(dst_path)
+    # Check the save path is empty
+    if isinstance(mlx_path, str):
+        mlx_path = Path(mlx_path)
 
-    for p in ["*.py", "generation_config.json"]:
-        for file in glob.glob(str(src_path / p)):
-            shutil.copy(file, dst_path)
+    if mlx_path.exists():
+        raise ValueError(
+            f"Cannot save to the path {mlx_path} as it already exists."
+            " Please delete the file/directory or specify a new path to save to."
+        )
 
-    if hf_repo is not None:
-        create_model_card(dst_path, hf_repo)
+    print("[INFO] Loading")
+    model_path = get_model_path(hf_path, revision=revision)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
 
+    weights = dict(tree_flatten(model.parameters()))
+    dtype = getattr(mx, dtype)
+    weights = {k: v.astype(dtype) for k, v in weights.items()}
 
-def common_prefix_len(list1, list2):
-    """
-    Calculates the length of the common prefix of two lists.
+    if quantize and dequantize:
+        raise ValueError("Choose either quantize or dequantize, not both.")
 
-    Args:
-        list1: The first list of strings.
-        list2: The second list of strings.
+    if quantize:
+        print("[INFO] Quantizing")
+        model.load_weights(list(weights.items()))
+        weights, config = quantize_model(
+            model, config, q_group_size, q_bits, quant_predicate=quant_predicate
+        )
 
-    Returns:
-        The length of the common prefix. Returns 0 if lists are empty
-        or do not match at the first element.
-    """
-    # Determine the maximum possible length of the common prefix
-    min_len = min(len(list1), len(list2))
+    if dequantize:
+        print("[INFO] Dequantizing")
+        model = dequantize_model(model)
+        weights = dict(tree_flatten(model.parameters()))
 
-    # Iterate up to the length of the shorter list
-    for i in range(min_len):
-        if list1[i] != list2[i]:
-            # Mismatch found, the common prefix length is the current index
-            return i
+    del model
+    save_weights(mlx_path, weights, donate_weights=True)
 
-    # No mismatch found within the bounds of the shorter list,
-    # so the common prefix length is the length of the shorter list.
-    return min_len
+    py_files = glob.glob(str(model_path / "*.py"))
+    for file in py_files:
+        shutil.copy(file, mlx_path)
+
+    tokenizer.save_pretrained(mlx_path)
+
+    save_config(config, config_path=mlx_path / "config.json")
+
+    if upload_repo is not None:
+        upload_to_hub(mlx_path, upload_repo, hf_path)

setup.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from setuptools import setup
 
 package_dir = Path(__file__).parent / "mlx_lm"
-with open("requirements.txt") as fid:
+with open(package_dir / "requirements.txt") as fid:
     requirements = [l.strip() for l in fid.readlines()]
 
 sys.path.append(str(package_dir))
@@ -21,7 +21,7 @@ setup(
     readme="README.md",
     author_email="mlx@group.apple.com",
     author="MLX Contributors",
-    url="https://github.com/ml-explore/mlx-lm",
+    url="https://github.com/ml-explore/mlx-examples",
     license="MIT",
     install_requires=requirements,
     packages=["mlx_lm", "mlx_lm.models", "mlx_lm.tuner"],
@@ -29,12 +29,9 @@ setup(
     extras_require={
         "test": ["datasets"],
         "evaluate": ["lm-eval", "tqdm"],
-        "lwq": ["datasets"],
     },
     entry_points={
         "console_scripts": [
-            "mlx_lm.awq = mlx_lm.awq:main",
-            "mlx_lm.dwq = mlx_lm.dwq:main",
             "mlx_lm.cache_prompt = mlx_lm.cache_prompt:main",
             "mlx_lm.chat = mlx_lm.chat:main",
             "mlx_lm.convert = mlx_lm.convert:main",
@@ -45,7 +42,6 @@ setup(
             "mlx_lm.merge = mlx_lm.merge:main",
             "mlx_lm.server = mlx_lm.server:main",
             "mlx_lm.manage = mlx_lm.manage:main",
-            "mlx_lm.upload = mlx_lm.upload:main",
         ]
     },
 )

tests/test_datsets.py

@@ -6,9 +6,8 @@ import tempfile
 import types
 import unittest
 
-from transformers import AutoTokenizer
-
 from mlx_lm.tuner import datasets
+from transformers import AutoTokenizer
 
 HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
 
@@ -44,7 +43,7 @@ class TestDatasets(unittest.TestCase):
         self.assertEqual(len(test), 0)
         self.assertTrue(len(train[0]) > 0)
         self.assertTrue(len(valid[0]) > 0)
-        self.assertTrue(isinstance(train, datasets.TextDataset))
+        self.assertTrue(isinstance(train, datasets.Dataset))
 
     def test_completions(self):
         data = {"prompt": "What is the capital of France?", "completion": "Paris."}
@@ -79,15 +78,14 @@ class TestDatasets(unittest.TestCase):
         self.assertTrue(isinstance(train, datasets.ChatDataset))
 
     def test_hf(self):
-        hf_args = {
-            "path": "billsum",
-            "prompt_feature": "text",
-            "completion_feature": "summary",
-            "train_split": "train[:2%]",
-            "valid_split": "train[-2%:]",
-        }
         args = types.SimpleNamespace(
-            hf_dataset=hf_args,
+            hf_dataset={
+                "name": "billsum",
+                "prompt_feature": "text",
+                "completion_feature": "summary",
+                "train_split": "train[:2%]",
+                "valid_split": "train[-2%:]",
+            },
             test=False,
             train=True,
         )
@@ -99,16 +97,6 @@ class TestDatasets(unittest.TestCase):
         self.assertTrue(len(valid[0]) > 0)
         self.assertEqual(len(test), 0)
 
-        args = types.SimpleNamespace(
-            hf_dataset=[hf_args, hf_args],
-            test=False,
-            train=True,
-        )
-        train_double, valid_double, test_double = datasets.load_dataset(args, tokenizer)
-        self.assertEqual(2 * len(train), len(train_double))
-        self.assertEqual(2 * len(valid), len(valid_double))
-        self.assertEqual(2 * len(test), len(test_double))
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_finetune.py

@@ -11,7 +11,6 @@ import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as opt
 from mlx.utils import tree_flatten
-
 from mlx_lm import lora, tuner
 from mlx_lm.tuner.dora import DoRAEmbedding, DoRALinear
 from mlx_lm.tuner.lora import LoRAEmbedding, LoRALinear
@@ -22,7 +21,7 @@ from mlx_lm.tuner.utils import build_schedule
 @contextmanager
 def swapped_with_identity(obj, func):
     old_func = getattr(obj, func)
-    setattr(obj, func, lambda x, **kwargs: x)
+    setattr(obj, func, lambda x: x)
     yield
     setattr(obj, func, old_func)
 
@@ -67,7 +66,7 @@ class TestLora(unittest.TestCase):
             )
             self.assertEqual(trainable_params, expected_trainable_parameters)
 
-        params = {"rank": 8, "dropout": 0.0, "scale": 10.0}
+        params = {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}
         check_config(params)
 
         params["rank"] = 1
@@ -108,7 +107,7 @@ class TestLora(unittest.TestCase):
         )
 
         num_lora_layers = 4
-        params = {"rank": 8, "dropout": 0.0, "scale": 10.0}
+        params = {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}
 
         model = gpt_neox.Model(args)
         model.freeze()
@@ -365,6 +364,7 @@ class TestScheduleConfig(unittest.TestCase):
     def test_evaluate_calls(self):
         mock_model = MagicMock()
         mock_dataset = MagicMock()
+        mock_tokenizer = MagicMock()
         mock_default_loss = MagicMock()
         mock_iterate_batches = MagicMock()
 
@@ -387,6 +387,7 @@ class TestScheduleConfig(unittest.TestCase):
             evaluate(
                 model=mock_model,
                 dataset=mock_dataset,
+                tokenizer=mock_tokenizer,
                 batch_size=2,
                 num_batches=2,
                 max_seq_length=2048,
@@ -396,6 +397,7 @@ class TestScheduleConfig(unittest.TestCase):
 
         mock_iterate_batches.assert_called_once_with(
             dataset=mock_dataset,
+            tokenizer=mock_tokenizer,
             batch_size=2,
             max_seq_length=2048,
         )
@@ -404,6 +406,7 @@ class TestScheduleConfig(unittest.TestCase):
     def test_evaluate_infinite_batches(self):
         mock_model = MagicMock()
         mock_dataset = MagicMock()
+        mock_tokenizer = MagicMock()
         mock_default_loss = MagicMock()
         mock_iterate_batches = MagicMock()
 
@@ -423,6 +426,7 @@ class TestScheduleConfig(unittest.TestCase):
             evaluate(
                 model=mock_model,
                 dataset=mock_dataset,
+                tokenizer=mock_tokenizer,
                 batch_size=2,
                 num_batches=-1,
                 max_seq_length=2048,
@@ -432,6 +436,7 @@ class TestScheduleConfig(unittest.TestCase):
 
         mock_iterate_batches.assert_called_once_with(
             dataset=mock_dataset,
+            tokenizer=mock_tokenizer,
             batch_size=2,
             max_seq_length=2048,
         )

tests/test_generate.py

@@ -1,23 +1,17 @@
 # Copyright © 2024 Apple Inc.
 
 import unittest
-from typing import List
 
-from mlx_lm.generate import (
-    GenerationResponse,
-    generate,
-    stream_generate,
-)
-from mlx_lm.sample_utils import make_logits_processors, make_sampler
-from mlx_lm.utils import load
+from mlx_lm.sample_utils import make_logits_processors
+from mlx_lm.utils import generate, load
 
 
 class TestGenerate(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        cls.HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
-        cls.model, cls.tokenizer = load(cls.HF_MODEL_PATH)
+        HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
+        cls.model, cls.tokenizer = load(HF_MODEL_PATH)
 
     def test_generate(self):
         # Simple test that generation runs
@@ -57,39 +51,6 @@ class TestGenerate(unittest.TestCase):
         )
         self.assertEqual(len(all_toks), len(init_toks) + 5)
 
-    def test_stream_generate_speculative(self):
-        # Use same model as draft model, this is not a speed test
-        draft_model, _ = load(self.HF_MODEL_PATH)
-
-        results: List[GenerationResponse] = []
-        drafted: List[bool] = []
-
-        # make a determinate sampler
-        sampler = make_sampler(temp=0.0)
-        messages = [{"role": "user", "content": "hello"}]
-        prompt = self.tokenizer.apply_chat_template(
-            messages, add_generation_prompt=True
-        )
-
-        for generation_result in stream_generate(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            prompt=prompt,
-            max_tokens=5,
-            draft_model=draft_model,
-            num_draft_tokens=2,
-            sampler=sampler,
-        ):
-            drafted.append(generation_result.from_draft)
-            results.append(generation_result)
-
-        self.assertEqual(len(results), 6)
-        drafted.pop()
-        # since num_draft_tokens is 2 and draft model is the same, the
-        # first 2 generations should be drafts, the third should come
-        # from the target model, and last two should be drafts
-        self.assertEqual(drafted, [True, True, False, True, True])
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_gguf.py

@@ -5,7 +5,6 @@ from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import mlx.core as mx
-
 from mlx_lm.gguf import convert_to_gguf
 
 

tests/test_models.py

@@ -4,9 +4,8 @@ import unittest
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.utils import tree_map
-
 from mlx_lm.models import rope_utils
-from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
+from mlx_lm.models.base import create_causal_mask
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
 
 
@@ -166,42 +165,6 @@ class TestModels(unittest.TestCase):
         )
         self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
 
-    def test_quantized_sdpa(self):
-        cache = KVCache()
-
-        k = 1e-1 * mx.random.normal(shape=(1, 1, 256, 32))
-        v = 1e-1 * mx.random.normal(shape=(1, 1, 256, 32))
-
-        cache.update_and_fetch(k, v)
-        quant_cache = cache.to_quantized(group_size=32, bits=8)
-
-        k = 1e-1 * mx.random.normal(shape=(1, 1, 1, 32))
-        v = 1e-1 * mx.random.normal(shape=(1, 1, 1, 32))
-
-        k_up, v_up = cache.update_and_fetch(k, v)
-        qk_up, qv_up = quant_cache.update_and_fetch(k, v)
-
-        q = 1e-1 * mx.random.normal(shape=(1, 4, 257, 32))
-
-        mask = "causal"
-        out = scaled_dot_product_attention(
-            q,
-            k_up,
-            v_up,
-            cache=cache,
-            mask=mask,
-            scale=1.0,
-        )
-        qout = scaled_dot_product_attention(
-            q,
-            qk_up,
-            qv_up,
-            cache=quant_cache,
-            mask=mask,
-            scale=1.0,
-        )
-        self.assertTrue(mx.allclose(out, qout, rtol=1e-2, atol=1e-2))
-
     def model_test_runner(self, model, model_type, vocab_size, num_layers):
 
         self.assertEqual(len(model.layers), num_layers)
@@ -220,7 +183,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 != "mamba":
                 mask = create_causal_mask(inputs.shape[1], 0).astype(t)
                 outputs = model(inputs, mask=mask)
                 self.assertEqual(outputs.shape, (1, 2, vocab_size))
@@ -343,56 +306,6 @@ class TestModels(unittest.TestCase):
             args.n_layers,
         )
 
-    def test_qwen3_moe(self):
-        from mlx_lm.models import qwen3_moe
-
-        args = qwen3_moe.ModelArgs(
-            model_type="qwen3_moe",
-            hidden_size=1024,
-            num_hidden_layers=4,
-            intermediate_size=2048,
-            num_attention_heads=4,
-            num_key_value_heads=4,
-            rms_norm_eps=1e-5,
-            head_dim=128,
-            vocab_size=10_000,
-            decoder_sparse_step=1,
-            mlp_only_layers=[],
-            num_experts_per_tok=4,
-            num_experts=16,
-            moe_intermediate_size=1024,
-            rope_theta=1000,
-            max_position_embeddings=4096,
-            tie_word_embeddings=False,
-            norm_topk_prob=True,
-        )
-        model = qwen3_moe.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
-    def test_qwen3(self):
-        from mlx_lm.models import qwen3
-
-        args = qwen3.ModelArgs(
-            model_type="qwen3",
-            hidden_size=1024,
-            num_hidden_layers=4,
-            intermediate_size=2048,
-            num_attention_heads=4,
-            num_key_value_heads=4,
-            rms_norm_eps=1e-5,
-            vocab_size=10_000,
-            head_dim=128,
-            max_position_embeddings=4096,
-            tie_word_embeddings=False,
-            rope_theta=1000,
-        )
-        model = qwen3.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
     def test_qwen2_moe(self):
         from mlx_lm.models import qwen2_moe
 
@@ -423,7 +336,6 @@ class TestModels(unittest.TestCase):
             num_hidden_layers=4,
             intermediate_size=2048,
             num_attention_heads=4,
-            num_key_value_heads=4,
             rms_norm_eps=1e-5,
             vocab_size=10_000,
         )
@@ -460,23 +372,6 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
-    def test_plamo2(self):
-        from mlx_lm.models import plamo2
-
-        args = plamo2.ModelArgs(
-            model_type="plamo2",
-            hidden_size=1024,
-            num_hidden_layers=4,
-            intermediate_size=2048,
-            num_attention_heads=8,
-            rms_norm_eps=1e-5,
-            vocab_size=10_000,
-        )
-        model = plamo2.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
     def test_stablelm(self):
         from mlx_lm.models import stablelm
 
@@ -787,43 +682,6 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
-    def test_deepseek_v3(self):
-        from mlx_lm.models import deepseek_v3
-
-        args = deepseek_v3.ModelArgs(
-            model_type="deepseek_v3",
-            vocab_size=1024,
-            hidden_size=128,
-            intermediate_size=256,
-            moe_intermediate_size=256,
-            num_hidden_layers=4,
-            num_attention_heads=4,
-            num_key_value_heads=2,
-            n_routed_experts=4,
-            n_group=2,
-            topk_group=1,
-            num_experts_per_tok=2,
-            n_shared_experts=1,
-            kv_lora_rank=4,
-            q_lora_rank=4,
-            qk_rope_head_dim=32,
-            v_head_dim=16,
-            qk_nope_head_dim=32,
-            rope_scaling={
-                "beta_fast": 32,
-                "beta_slow": 1,
-                "factor": 40,
-                "mscale": 1.0,
-                "mscale_all_dim": 1.0,
-                "original_max_position_embeddings": 4096,
-                "type": "yarn",
-            },
-        )
-        model = deepseek_v3.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
     def test_gemma2(self):
         from mlx_lm.models import gemma2
 
@@ -843,26 +701,6 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
-    def test_gemma3_text(self):
-        from mlx_lm.models import gemma3_text
-
-        args = gemma3_text.ModelArgs(
-            model_type="gemma3_text",
-            hidden_size=128,
-            num_hidden_layers=12,
-            intermediate_size=256,
-            num_attention_heads=4,
-            head_dim=32,
-            rms_norm_eps=1e-4,
-            num_key_value_heads=1,
-            sliding_window=1024,
-            sliding_window_pattern=6,
-        )
-        model = gemma3_text.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
     def test_gpt_bigcode(self):
         from mlx_lm.models import gpt_bigcode
 
@@ -1052,23 +890,6 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
-    def test_internlm3(self):
-        from mlx_lm.models import internlm3
-
-        args = internlm3.ModelArgs(
-            model_type="internlm3",
-            hidden_size=1024,
-            num_hidden_layers=4,
-            intermediate_size=2048,
-            num_attention_heads=4,
-            rms_norm_eps=1e-5,
-            vocab_size=10_000,
-        )
-        model = internlm3.Model(args)
-        self.model_test_runner(
-            model, args.model_type, args.vocab_size, args.num_hidden_layers
-        )
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_prompt_cache.py

@@ -6,10 +6,7 @@ import tempfile
 import unittest
 
 import mlx.core as mx
-
-from mlx_lm.generate import generate_step
 from mlx_lm.models.cache import (
-    ChunkedKVCache,
     KVCache,
     MambaCache,
     QuantizedKVCache,
@@ -19,7 +16,7 @@ from mlx_lm.models.cache import (
     save_prompt_cache,
     trim_prompt_cache,
 )
-from mlx_lm.utils import load
+from mlx_lm.utils import generate_step, load
 
 HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
 
@@ -96,13 +93,7 @@ class TestPromptCache(unittest.TestCase):
     def test_save_load_mixed_cache(self):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
 
-        cache = [
-            MambaCache(),
-            KVCache(),
-            RotatingKVCache(8),
-            MambaCache(),
-            ChunkedKVCache(256),
-        ]
+        cache = [MambaCache(), KVCache(), RotatingKVCache(8), MambaCache()]
         for c in cache:
             if isinstance(c, MambaCache):
                 c[0] = mx.random.uniform(shape=(4, 4, 4))
@@ -307,7 +298,7 @@ class TestPromptCache(unittest.TestCase):
         ):
             i += 1
             self.assertEqual(tok, toks[i])
-            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=4e-2))
+            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=2e-2))
 
 
 if __name__ == "__main__":

tests/test_sample_utils.py

@@ -1,120 +1,67 @@
 import unittest
 
 import mlx.core as mx
-
-from mlx_lm.sample_utils import apply_min_p, apply_top_k, apply_top_p, apply_xtc
+from mlx_lm.sample_utils import min_p_sampling, top_k_sampling, top_p_sampling
 
 
 class TestSampleUtils(unittest.TestCase):
-    def test_apply_top_p(self):
+    def test_top_p_sampling(self):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
+        temperature = 1.0
 
-        new_logits = apply_top_p(logits, 0.3)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertEqual(actual_probs.tolist(), [1.0, 0.0, 0.0, 0.0])
+        token = top_p_sampling(logits, 0.3, temperature).item()
+        self.assertEqual(token, 0)
 
-        new_logits = apply_top_p(logits, 0.95)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertTrue(mx.allclose(probs.squeeze(), actual_probs))
+        token = top_p_sampling(logits, 0.95, temperature).item()
+        self.assertTrue(token in (0, 3))
 
         probs = mx.array([0.0, 0.5, 0.4, 0.1])[None]
         logits = mx.log(probs)
-        new_logits = apply_top_p(logits, 0.4)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertEqual(actual_probs.tolist(), [0.0, 1.0, 0.0, 0.0])
 
-        new_logits = apply_top_p(logits, 0.6)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertEqual(
-            [round(p, 4) for p in actual_probs.tolist()], [0.0, 0.5556, 0.4444, 0.0]
-        )
+        token = top_p_sampling(logits, 0.4, temperature).item()
+        self.assertEqual(token, 1)
 
-        new_logits = apply_top_p(logits, 0.95)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        actual_rounded = [round(p, 4) for p in actual_probs.tolist()]
-        expected_rounded = [0.0, 0.5, 0.4, 0.1]
-        self.assertEqual(actual_rounded, expected_rounded)
-        self.assertAlmostEqual(sum(actual_probs.tolist()), 1.0)
+        token = top_p_sampling(logits, 0.6, temperature).item()
+        self.assertTrue(token in (1, 2))
 
-        # Batch mode works
-        probs = mx.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.8, 0.1, 0.1]])
-        logits = mx.log(probs)
-        new_logits = apply_top_p(logits, 0.5)
-        actual_probs = mx.softmax(new_logits, axis=-1)
-        self.assertEqual(
-            actual_probs.tolist(), [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]
-        )
+        token = top_p_sampling(logits, 0.95, temperature).item()
+        self.assertTrue(token in (1, 2, 3))
 
-    def test_apply_min_p(self):
+    def test_min_p_sampling(self):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
-        new_logits = apply_min_p(logits, 0.8)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertEqual(actual_probs.tolist(), [1.0, 0.0, 0.0, 0.0])
+        temperature = 1.0
+        token = min_p_sampling(logits, 0.8)
+        self.assertEqual(token, 0)
 
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
-        new_logits = apply_min_p(logits, 0.05)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertTrue(mx.allclose(actual_probs, mx.squeeze(probs)))
+        temperature = 1.0
+        for _ in range(5):
+            token = min_p_sampling(logits, 0.05)
+            self.assertTrue(token in (0, 3))
 
-        # Batch mode works
-        probs = mx.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.8, 0.0, 0.1]])
-        logits = mx.log(probs)
-        new_logits = apply_min_p(logits, 0.7)
-        actual_probs = mx.softmax(new_logits, axis=-1)
-        self.assertEqual(
-            actual_probs.tolist(), [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]
-        )
-
-    def test_apply_top_k(self):
+    def test_top_k_sampling(self):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
 
-        new_logits = apply_top_k(logits, 1)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertEqual(actual_probs.tolist(), [1.0, 0.0, 0.0, 0.0])
+        token = top_k_sampling(logits, 1).item()
+        self.assertEqual(token, 0)
 
-        probs = mx.array([0.6, 0.0, 0.1, 0.3])[None]
-        logits = mx.log(probs)
-        new_logits = apply_top_k(logits, 2)
-        actual_probs = mx.softmax(new_logits.squeeze())
-        self.assertEqual(
-            [round(p, 4) for p in actual_probs.tolist()], [0.6667, 0.0, 0.0, 0.3333]
-        )
+        probs = mx.array([0.5, 0.0, 0.0, 0.5])[None]
+        tokens = set()
+        for _ in range(100):
+            token = top_k_sampling(logits, 2)
+            tokens.add(token.item())
+        self.assertEqual(tokens, {0, 3})
 
         # Batch mode works
         probs = mx.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.8, 0.0, 0.1]])
         logits = mx.log(probs)
 
-        new_logits = apply_top_k(logits, 1)
-        actual_probs = mx.softmax(new_logits, axis=-1)
-        self.assertEqual(
-            actual_probs.tolist(), [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]
-        )
-
-    def test_apply_xtc(self):
-        # Test the threshold
-        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
-        new_probs = mx.softmax(apply_xtc(mx.log(probs), 1, 0.2, []), -1)
-        expected = mx.array([[0, 0.5, 0.25, 0.25]])
-        self.assertTrue(mx.allclose(new_probs, expected))
-        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
-        new_probs = mx.softmax(apply_xtc(mx.log(probs), 1, 0.1, []), -1)
-        expected = mx.array([[0, 0.0, 0.5, 0.5]])
-        self.assertTrue(mx.allclose(new_probs, expected))
-
-        # Test the special tokens
-        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
-        new_probs = mx.softmax(apply_xtc(mx.log(probs), 1, 0.1, [0]), -1)
-        expected = mx.array([[4 / 7, 0.0, 1.5 / 7, 1.5 / 7]])
-        self.assertTrue(mx.allclose(new_probs, expected))
-
-        # Test that with probability 0 the probs don't change
-        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
-        new_probs = mx.softmax(apply_xtc(mx.log(probs), 0, 0.1, [0]), -1)
-        self.assertTrue(mx.allclose(new_probs, probs))
+        tokens = top_k_sampling(logits, 1)
+        self.assertEqual(tokens.tolist(), [0, 1])
 
 
 if __name__ == "__main__":

tests/test_server.py

@@ -6,37 +6,17 @@ import threading
 import unittest
 
 import requests
-
 from mlx_lm.server import APIHandler
 from mlx_lm.utils import load
 
 
 class DummyModelProvider:
-    def __init__(self, with_draft=False):
+    def __init__(self):
         HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
         self.model, self.tokenizer = load(HF_MODEL_PATH)
         self.model_key = (HF_MODEL_PATH, None)
 
-        # Add draft model support
-        self.draft_model = None
-        self.draft_model_key = None
-        self.cli_args = type(
-            "obj",
-            (object,),
-            {
-                "adapter_path": None,
-                "chat_template": None,
-                "use_default_chat_template": False,
-                "trust_remote_code": False,
-            },
-        )
-
-        if with_draft:
-            # Use the same model as the draft model for testing
-            self.draft_model, _ = load(HF_MODEL_PATH)
-            self.draft_model_key = HF_MODEL_PATH
-
-    def load(self, model, adapter=None, draft_model=None):
+    def load(self, model, adapter=None):
         assert model in ["default_model", "chat_model"]
         return self.model, self.tokenizer
 
@@ -100,29 +80,6 @@ class TestServer(unittest.TestCase):
         self.assertIn("id", response_body)
         self.assertIn("choices", response_body)
 
-    def test_handle_chat_completions_with_content_fragments(self):
-        url = f"http://localhost:{self.port}/v1/chat/completions"
-        chat_post_data = {
-            "model": "chat_model",
-            "max_tokens": 10,
-            "temperature": 0.7,
-            "top_p": 0.85,
-            "repetition_penalty": 1.2,
-            "messages": [
-                {
-                    "role": "system",
-                    "content": [
-                        {"type": "text", "text": "You are a helpful assistant."}
-                    ],
-                },
-                {"role": "user", "content": [{"type": "text", "text": "Hello!"}]},
-            ],
-        }
-        response = requests.post(url, json=chat_post_data)
-        response_body = response.text
-        self.assertIn("id", response_body)
-        self.assertIn("choices", response_body)
-
     def test_handle_models(self):
         url = f"http://localhost:{self.port}/v1/models"
         response = requests.get(url)
@@ -149,324 +106,5 @@ class TestServer(unittest.TestCase):
         self.assertFalse(sequence_overlap([1, 2, 3], [4, 1, 2, 3]))
 
 
-class TestServerWithDraftModel(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model_provider = DummyModelProvider(with_draft=True)
-        cls.server_address = ("localhost", 0)
-        cls.httpd = http.server.HTTPServer(
-            cls.server_address,
-            lambda *args, **kwargs: APIHandler(cls.model_provider, *args, **kwargs),
-        )
-        cls.port = cls.httpd.server_port
-        cls.server_thread = threading.Thread(target=cls.httpd.serve_forever)
-        cls.server_thread.daemon = True
-        cls.server_thread.start()
-
-    @classmethod
-    def tearDownClass(cls):
-        cls.httpd.shutdown()
-        cls.httpd.server_close()
-        cls.server_thread.join()
-
-    def test_handle_completions_with_draft_model(self):
-        url = f"http://localhost:{self.port}/v1/completions"
-
-        post_data = {
-            "model": "default_model",
-            "prompt": "Once upon a time",
-            "max_tokens": 10,
-            "temperature": 0.0,
-            "top_p": 1.0,
-        }
-
-        response = requests.post(url, json=post_data)
-        self.assertEqual(response.status_code, 200)
-
-        response_body = json.loads(response.text)
-        self.assertIn("id", response_body)
-        self.assertIn("choices", response_body)
-        self.assertIn("usage", response_body)
-
-        # Check that tokens were generated
-        self.assertTrue(response_body["usage"]["completion_tokens"] > 0)
-
-    def test_handle_chat_completions_with_draft_model(self):
-        url = f"http://localhost:{self.port}/v1/chat/completions"
-
-        chat_post_data = {
-            "model": "chat_model",
-            "max_tokens": 10,
-            "temperature": 0.0,
-            "messages": [
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": "Hello!"},
-            ],
-        }
-
-        response = requests.post(url, json=chat_post_data)
-        self.assertEqual(response.status_code, 200)
-
-        response_body = json.loads(response.text)
-        self.assertIn("id", response_body)
-        self.assertIn("choices", response_body)
-        self.assertIn("usage", response_body)
-
-        # Check that tokens were generated
-        self.assertTrue(response_body["usage"]["completion_tokens"] > 0)
-
-    def test_streaming_with_draft_model(self):
-        url = f"http://localhost:{self.port}/v1/chat/completions"
-
-        chat_post_data = {
-            "model": "chat_model",
-            "max_tokens": 10,
-            "temperature": 0.0,
-            "stream": True,
-            "messages": [
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": "Hello!"},
-            ],
-        }
-
-        response = requests.post(url, json=chat_post_data, stream=True)
-        self.assertEqual(response.status_code, 200)
-
-        chunk_count = 0
-        for chunk in response.iter_lines():
-            if chunk:
-                data = chunk.decode("utf-8")
-                if data.startswith("data: ") and data != "data: [DONE]":
-                    chunk_data = json.loads(data[6:])  # Skip the "data: " prefix
-                    self.assertIn("choices", chunk_data)
-                    self.assertEqual(len(chunk_data["choices"]), 1)
-                    self.assertIn("delta", chunk_data["choices"][0])
-                    chunk_count += 1
-
-        # Make sure we got some streaming chunks
-        self.assertGreater(chunk_count, 0)
-
-    def test_prompt_cache_with_draft_model(self):
-        url = f"http://localhost:{self.port}/v1/chat/completions"
-
-        # First request to initialize cache
-        chat_post_data = {
-            "model": "chat_model",
-            "max_tokens": 5,
-            "temperature": 0.0,
-            "messages": [
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": "Tell me a story about"},
-            ],
-        }
-
-        first_response = requests.post(url, json=chat_post_data)
-        self.assertEqual(first_response.status_code, 200)
-
-        # Second request with same prefix should use cache
-        chat_post_data = {
-            "model": "chat_model",
-            "max_tokens": 5,
-            "temperature": 0.0,
-            "messages": [
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": "Tell me a story about dragons."},
-            ],
-        }
-
-        second_response = requests.post(url, json=chat_post_data)
-        self.assertEqual(second_response.status_code, 200)
-
-        # Both responses should have content
-        first_response_body = json.loads(first_response.text)
-        second_response_body = json.loads(second_response.text)
-
-        self.assertIn("choices", first_response_body)
-        self.assertIn("choices", second_response_body)
-        self.assertIn("message", first_response_body["choices"][0])
-        self.assertIn("message", second_response_body["choices"][0])
-        self.assertIn("content", first_response_body["choices"][0]["message"])
-        self.assertIn("content", second_response_body["choices"][0]["message"])
-
-        # Ensure both generated content
-        self.assertIsNotNone(first_response_body["choices"][0]["message"]["content"])
-        self.assertIsNotNone(second_response_body["choices"][0]["message"]["content"])
-
-
-# --- Tests for get_prompt_cache ---
-
-from unittest.mock import MagicMock, patch
-
-from mlx_lm.server import PromptCache
-
-
-class TestGetPromptCache(unittest.TestCase):
-
-    def setUp(self):
-        """Set up mocks and a handler instance for each test."""
-        self.mock_model_provider = MagicMock()
-        # Simulate tokenizer needed for decoding in original debug logs (though not strictly needed for cache logic)
-        self.mock_model_provider.tokenizer = MagicMock()
-        self.mock_model_provider.tokenizer.decode = lambda x: f"decoded({x})"
-        self.mock_model_provider.model_key = ("model_v1", None, None)
-        self.mock_model_provider.draft_model = None  # Start without draft model
-
-        # --- Prevent BaseHTTPRequestHandler.__init__ from running ---
-        # It tries to handle a request immediately, which fails with mocks.
-        # We only need the APIHandler instance with its attributes set.
-        with patch(
-            "http.server.BaseHTTPRequestHandler.__init__", lambda *args, **kwargs: None
-        ):
-            # APIHandler init still requires args for BaseHTTPRequestHandler signature,
-            # but they won't be used by the patched __init__.
-            mock_request = MagicMock()
-            mock_client_address = ("127.0.0.1", 8080)
-            mock_server = MagicMock()
-
-            self.prompt_cache_instance = PromptCache()
-            self.handler = APIHandler(
-                self.mock_model_provider,
-                mock_request,
-                mock_client_address,
-                mock_server,
-                prompt_cache=self.prompt_cache_instance,  # Inject our cache instance
-            )
-            # Manually set attributes usually set by APIHandler.__init__ if needed
-            # self.handler.created = MagicMock()
-            # self.handler.system_fingerprint = MagicMock()
-            # (Not strictly necessary for get_prompt_cache testing)
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    def test_initial_request_empty_cache(self, mock_make_cache):
-        """Test first request when the cache is empty."""
-        mock_make_cache.return_value = "new_cache_obj"
-        prompt = [1, 2, 3]
-
-        processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        self.assertEqual(processed_prompt, [1, 2, 3])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj")
-        self.assertEqual(self.handler.prompt_cache.model_key, ("model_v1", None, None))
-        mock_make_cache.assert_called_once()
-
-    def test_identical_request_full_hit(self):
-        """Test when the new prompt is identical to the cached one."""
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3]
-
-        # Mock common_prefix_len to return the full length
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process nothing, cache remains unchanged
-        self.assertEqual(processed_prompt, [])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
-        self.assertEqual(self.handler.prompt_cache.cache, "existing_cache_obj")
-
-    def test_cache_is_prefix(self):
-        """Test when the cached prompt is a prefix of the new prompt."""
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3, 4, 5]
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the suffix, cache tokens updated
-        self.assertEqual(processed_prompt, [4, 5])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 4, 5])
-        self.assertEqual(self.handler.prompt_cache.cache, "existing_cache_obj")
-
-    @patch("mlx_lm.server.trim_prompt_cache")
-    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=True)
-    def test_partial_match_trim_success(self, mock_can_trim, mock_trim_cache):
-        """Test partial match where cache is longer and trimming succeeds."""
-        self.handler.prompt_cache.tokens = [1, 2, 3, 4, 5]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3, 6, 7]  # Diverges after token 3
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the new suffix, cache trimmed and updated
-        self.assertEqual(processed_prompt, [6, 7])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 6, 7])
-        mock_can_trim.assert_called_once_with("existing_cache_obj")
-        # Called with cache object and num_to_trim (5 - 3 = 2)
-        mock_trim_cache.assert_called_once_with("existing_cache_obj", 2)
-        self.assertEqual(
-            self.handler.prompt_cache.cache, "existing_cache_obj"
-        )  # Cache obj itself isn't changed by mock
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    @patch("mlx_lm.server.trim_prompt_cache")
-    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=False)
-    def test_partial_match_trim_fail(
-        self, mock_can_trim, mock_trim_cache, mock_make_cache
-    ):
-        """Test partial match where cache is longer but trimming fails."""
-        mock_make_cache.return_value = "new_cache_obj_on_reset"
-        self.handler.prompt_cache.tokens = [1, 2, 3, 4, 5]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3, 6, 7]  # Diverges after token 3
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the full prompt, cache reset
-        self.assertEqual(processed_prompt, [1, 2, 3, 6, 7])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 6, 7])
-        mock_can_trim.assert_called_once_with("existing_cache_obj")
-        mock_trim_cache.assert_not_called()
-        mock_make_cache.assert_called_once()  # Cache was reset
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj_on_reset")
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    def test_no_common_prefix(self, mock_make_cache):
-        """Test when there is no common prefix between cache and prompt."""
-        mock_make_cache.return_value = "new_cache_obj"
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [4, 5, 6]
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=0):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the full prompt, cache reset
-        self.assertEqual(processed_prompt, [4, 5, 6])
-        self.assertEqual(self.handler.prompt_cache.tokens, [4, 5, 6])
-        mock_make_cache.assert_called_once()
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj")
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    def test_model_changed(self, mock_make_cache):
-        """Test cache reset when the model key changes."""
-        mock_make_cache.return_value = "new_cache_obj_model_change"
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)  # Original key
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-
-        # Simulate model provider having a new key
-        self.mock_model_provider.model_key = ("model_v2", None, None)
-        prompt = [1, 2, 3, 4]
-
-        # No need to mock common_prefix_len, model check happens first
-        processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the full prompt, cache reset
-        self.assertEqual(processed_prompt, [1, 2, 3, 4])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 4])
-        mock_make_cache.assert_called_once()
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj_model_change")
-        self.assertEqual(self.handler.prompt_cache.model_key, ("model_v2", None, None))
-
-
 if __name__ == "__main__":
     unittest.main()

tests/test_tokenizers.py

@@ -4,7 +4,6 @@ import unittest
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
-
 from mlx_lm.tokenizer_utils import (
     BPEStreamingDetokenizer,
     NaiveStreamingDetokenizer,

tests/test_tuner_utils.py

@@ -6,7 +6,6 @@ from io import StringIO
 from unittest.mock import MagicMock
 
 import mlx.nn as nn
-
 from mlx_lm.tuner.lora import LoRALinear
 from mlx_lm.tuner.utils import print_trainable_parameters
 

tests/test_utils.py

@@ -7,8 +7,7 @@ import unittest
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.utils import tree_flatten
-
-from mlx_lm import convert, utils
+from mlx_lm import utils
 
 HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
 
@@ -77,14 +76,14 @@ class TestUtils(unittest.TestCase):
     def test_convert(self):
         mlx_path = os.path.join(self.test_dir, "mlx_model")
 
-        convert(HF_MODEL_PATH, mlx_path=mlx_path, quantize=True)
+        utils.convert(HF_MODEL_PATH, mlx_path=mlx_path, quantize=True)
         model, _ = utils.load(mlx_path)
         self.assertTrue(isinstance(model.layers[0].mlp.up_proj, nn.QuantizedLinear))
         self.assertTrue(isinstance(model.layers[-1].mlp.up_proj, nn.QuantizedLinear))
 
         # Check model weights have right type
         mlx_path = os.path.join(self.test_dir, "mlx_model_bf16")
-        convert(HF_MODEL_PATH, mlx_path=mlx_path, dtype="bfloat16")
+        utils.convert(HF_MODEL_PATH, mlx_path=mlx_path, dtype="bfloat16")
         model, _ = utils.load(mlx_path)
 
         self.assertEqual(model.layers[0].mlp.up_proj.weight.dtype, mx.bfloat16)

tests/test_utils_load_model.py

@@ -2,7 +2,6 @@ import unittest
 from pathlib import Path
 
 import mlx.nn as nn
-
 from mlx_lm.models.qwen2 import Model as Qwen2Model
 from mlx_lm.utils import get_model_path, load_model
 
@@ -18,7 +17,7 @@ class TestLoadModelCustomGetClasses(unittest.TestCase):
                 self.config = args
                 self.custom_attribute = "This is a custom model"
 
-            def load_weights(self, weights, **kwargs):
+            def load_weights(self, weights):
                 self.qwenWeights = weights
 
         class CustomQwenConfig: