fix test

* refactor out of utils

* remove comment

* format

* move generate as well

.circleci/config.yml

@@ -0,0 +1,66 @@
+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

@@ -0,0 +1,139 @@
+# 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

@@ -0,0 +1,11 @@
+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,13 +5,8 @@ with a short description of your contribution(s) below. For example:
 
 - Jane Smith: Added the `foo` example.
 
-MLX Examples was developed with contributions from the following individuals:
+MLX LM 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.
-- Shiyu Li: Added the `Segment Anything Model`.
-- Gökdeniz Gülmez: Added support for `MiniCPM`, `Mamba` and support for `full-fine-tuning`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`.

CONTRIBUTING.md

@@ -1,11 +1,54 @@
 # 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.
 
-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:
+From this directory, do an editable install:
 
 ```shell
 pip install -e .
@@ -17,7 +60,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-examples/tree/main/llms/mlx_lm/models)
+[`mlx_lm/models`](https://github.com/ml-explore/mlx-lm/tree/main/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.
 
@@ -35,12 +78,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-examples/blob/main/llms/mlx_lm/tuner/utils.py#L27-L60)
+[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-lm/blob/main/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-examples/blob/main/llms/tests/test_models.py).
+tests](https://github.com/ml-explore/mlx-lm/blob/main/tests/test_models.py).
 
-From the `llms/` directory, you can run the tests with:
+You can run the tests with:
 
 ```shell
 python -m unittest discover tests/

MANIFEST.in

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

README.md

@@ -1,4 +1,17 @@
-## Generate Text with LLMs and MLX
+## 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`
 
 The easiest way to get started is to install the `mlx-lm` package:
 
@@ -14,18 +27,12 @@ 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 "Hi!"
+mlx_lm.generate --prompt "How tall is Mt Everest?"
 ```
 
 To chat with an LLM use:
@@ -71,7 +78,7 @@ To see a description of all the arguments you can do:
 ```
 
 Check out the [generation
-example](https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/examples/generate_response.py)
+example](https://github.com/ml-explore/mlx-lm/tree/main/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
@@ -123,6 +130,18 @@ 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:
@@ -164,7 +183,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 
@@ -201,17 +220,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 to avoid
+Prompt caching can also be used in the Python API in order 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-examples/blob/main/llms/mlx_lm/examples/chat.py)
+[example](https://github.com/ml-explore/mlx-lm/blob/main/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-examples/issues/new) or better yet,
+[issue](https://github.com/ml-explore/mlx-lm/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/LORA.md

@@ -76,6 +76,14 @@ 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:
@@ -241,14 +249,25 @@ 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
@@ -270,7 +289,7 @@ 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:
   name: "billsum"
   prompt_feature: "text"
@@ -279,11 +298,27 @@ hf_dataset:
 
 - Use `prompt_feature` and `completion_feature` to specify keys for a
   `completions` dataset. Use `text_feature` to specify the key for a `text`
-  dataset. 
+  dataset. Use `chat_feature` to specify the key for a chat 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:
+  - name: "Open-Orca/OpenOrca"
+    train_split: "train[:90%]"
+    valid_split: "train[-10%:]"
+    prompt_feature: "question"
+    completion_feature: "response"
+  - name: "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/__init__.py

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

mlx_lm/__main__.py

@@ -0,0 +1,14 @@
+# Copyright © 2025 Apple Inc.
+
+import importlib
+import sys
+
+if __name__ == "__main__":
+    subcommands = {"convert"}
+    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.20.4"
+__version__ = "0.22.2"

mlx_lm/cache_prompt.py

@@ -7,8 +7,9 @@ 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 generate_step, load
+from .utils import load
 
 DEFAULT_QUANTIZED_KV_START = 5000
 
@@ -152,7 +153,7 @@ def main():
     print("Saving...")
     metadata = {}
     metadata["model"] = args.model
-    metadata["chat_template"] = tokenizer.chat_template
+    metadata["chat_template"] = json.dumps(tokenizer.chat_template)
     metadata["tokenizer_config"] = json.dumps(tokenizer_config)
     save_prompt_cache(args.prompt_cache_file, cache, metadata)
 

mlx_lm/chat.py

@@ -5,13 +5,14 @@ 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, stream_generate
+from .utils import load
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
-DEFAULT_SEED = 0
+DEFAULT_SEED = None
 DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 
@@ -36,7 +37,12 @@ def setup_arg_parser():
     parser.add_argument(
         "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
     )
-    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,
@@ -57,7 +63,8 @@ def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
-    mx.random.seed(args.seed)
+    if args.seed is not None:
+        mx.random.seed(args.seed)
 
     model, tokenizer = load(
         args.model,
@@ -65,12 +72,25 @@ def main():
         tokenizer_config={"trust_remote_code": True},
     )
 
-    print(f"[INFO] Starting chat session with {args.model}. To exit, enter 'q'.")
+    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()
     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(

mlx_lm/convert.py

@@ -1,8 +1,137 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+import glob
+import shutil
+from pathlib import Path
+from typing import Callable, Optional, Union
 
-from .utils import convert
+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_config,
+    save_weights,
+    upload_to_hub,
+)
+
+
+def mixed_quant_predicate_builder(
+    low_bits: int = 4, high_bits: int = 4, group_size: int = 64
+) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+    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(".")[2]) if len(path.split(".")) > 2 else 0
+
+        num_layers = config["num_hidden_layers"]
+        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_quant_predicate_builder(low_bits=3, high_bits=6),
+    "mixed_3_6": mixed_quant_predicate_builder(low_bits=2, high_bits=6),
+}
+
+
+def quant_args(arg):
+    if arg not in QUANT_RECIPES:
+        raise argparse.ArgumentTypeError(
+            f"Invalid q-recipe {arg!r}. Choose from: {list(QUANT_RECIPES.keys())}"
+        )
+    else:
+        return QUANT_RECIPES[arg]
+
+
+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,
+):
+    # 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)
+
+    weights = dict(tree_flatten(model.parameters()))
+    dtype = getattr(mx, dtype)
+    weights = {k: v.astype(dtype) 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_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")
+
+    if upload_repo is not None:
+        upload_to_hub(mlx_path, upload_repo, hf_path)
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -29,6 +158,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: {list(QUANT_RECIPES.keys())}",
+        type=quant_args,
+        required=False,
+    )
     parser.add_argument(
         "--dtype",
         help="Type to save the non-quantized parameters.",
@@ -59,4 +194,8 @@ 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/evaluate.py

@@ -20,8 +20,9 @@ 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.cache import make_prompt_cache
-from .utils import load, stream_generate
+from .utils import load
 
 PAD = 0
 
@@ -111,7 +112,7 @@ class MLXLM(LM):
             )
 
             mx.eval(score, ig)
-            mx.metal.clear_cache()
+            mx.clear_cache()
 
             is_greedy.append(ig)
             scores.append(score)
@@ -289,15 +290,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, until in tqdm(zip(contexts, untils), total=len(contexts)):
-            context = self._tokenize(context)
+        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
+            )
             max_tokens = min(
-                self._max_tokens,
+                opt.get("max_gen_tokens", self._max_tokens),
                 self.tokenizer.model_max_length - len(context),
             )
             text = ""
@@ -332,9 +333,9 @@ def main():
     )
     parser.add_argument(
         "--limit",
-        default=1.0,
+        default=100,
         help="Limit the number of examples per task.",
-        type=float,
+        type=int,
     )
     parser.add_argument("--seed", type=int, default=123, help="Random seed.")
     parser.add_argument(

mlx_lm/examples/chat.py

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

mlx_lm/examples/lora_config.yaml

@@ -7,6 +7,15 @@ 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"
 
@@ -72,7 +81,7 @@ lora_parameters:
 #  arguments: [1e-5, 1000, 1e-7] # passed to scheduler
 
 #hf_dataset:
-#  name: "billsum"
+#  path: "billsum"
 #  train_split: "train[:1000]"
 #  valid_split: "train[-100:]"
 #  prompt_feature: "text"

mlx_lm/examples/pipeline_generate.py

@@ -0,0 +1,128 @@
+# Copyright © 2024 Apple Inc.
+
+"""
+Run with:
+
+```
+mlx.launch \
+ --hostfile /path/to/hosts.txt \
+ --backend mpi \
+ /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
+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
+
+
+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(backend="mpi")
+    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(backend="mpi")
+    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

@@ -0,0 +1,73 @@
+# 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/generate.py

@@ -1,14 +1,37 @@
 # 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.cache import QuantizedKVCache, load_prompt_cache
+from .models import cache
+from .models.cache import (
+    QuantizedKVCache,
+    load_prompt_cache,
+    make_prompt_cache,
+    trim_prompt_cache,
+)
 from .sample_utils import make_sampler
-from .utils import generate, load
+from .tokenizer_utils import TokenizerWrapper
+from .utils import load
 
 DEFAULT_PROMPT = "hello"
 DEFAULT_MAX_TOKENS = 100
@@ -16,7 +39,7 @@ DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
 DEFAULT_MIN_P = 0.0
 DEFAULT_MIN_TOKENS_TO_KEEP = 1
-DEFAULT_SEED = 0
+DEFAULT_SEED = None
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 DEFAULT_QUANTIZED_KV_START = 5000
 
@@ -60,6 +83,11 @@ 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",
@@ -82,7 +110,12 @@ def setup_arg_parser():
         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",
@@ -93,6 +126,12 @@ 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,
@@ -131,14 +170,547 @@ 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()
 
-    mx.random.seed(args.seed)
+    if args.seed is not None:
+        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
@@ -183,11 +755,15 @@ 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 = metadata["chat_template"]
+        tokenizer.chat_template = json.loads(metadata["chat_template"])
 
     prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
     prompt = sys.stdin.read() if prompt == "-" else prompt
@@ -197,8 +773,16 @@ 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, add_generation_prompt=True
+            messages,
+            tokenize=False,
+            continue_final_message=has_prefill,
+            add_generation_prompt=not has_prefill,
+            **template_kwargs,
         )
 
         # Treat the prompt as a suffix assuming that the prefix is in the
@@ -208,14 +792,20 @@ def main():
             test_prompt = tokenizer.apply_chat_template(
                 messages,
                 tokenize=False,
-                add_generation_prompt=True,
+                continue_final_message=has_prefill,
+                add_generation_prompt=not has_prefill,
             )
             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)
     response = generate(
         model,
@@ -229,10 +819,16 @@ 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,6 +7,7 @@ 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
@@ -43,6 +44,11 @@ CONFIG_DEFAULTS = {
     "model": "mlx_model",
     "train": False,
     "fine_tune_type": "lora",
+    "optimizer": "adam",
+    "optimizer_config": {
+        "adam": {},
+        "adamw": {},
+    },
     "data": "data/",
     "seed": 0,
     "num_layers": 16,
@@ -58,8 +64,11 @@ CONFIG_DEFAULTS = {
     "test": False,
     "test_batches": 500,
     "max_seq_length": 2048,
+    "config": None,
+    "grad_checkpoint": False,
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
+    "mask_prompt": False,
 }
 
 
@@ -67,6 +76,7 @@ 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.",
     )
 
@@ -89,9 +99,21 @@ 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,
@@ -146,10 +168,16 @@ def build_parser():
         type=int,
         help="Maximum sequence length.",
     )
+    parser.add_argument(
+        "--seq-step-size",
+        type=int,
+        default=None,
+        help="",
+    )
     parser.add_argument(
         "-c",
         "--config",
-        default=None,
+        type=str,
         help="A YAML configuration file with the training options",
     )
     parser.add_argument(
@@ -158,7 +186,7 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
-    parser.add_argument("--seed", type=int, default=None, help="The PRNG seed")
+    parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
 
@@ -170,9 +198,16 @@ def train_model(
     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[-min(args.num_layers, 0) :]:
+        for l in model.layers[-max(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
@@ -209,14 +244,24 @@ def train_model(
         adapter_file=adapter_file,
         max_seq_length=args.max_seq_length,
         grad_checkpoint=args.grad_checkpoint,
+        seq_step_size=args.seq_step_size,
     )
 
-    model.train()
-    opt = optim.Adam(
-        learning_rate=(
-            build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
-        )
-    )
+    # 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)
+
     # Train model
     train(
         model=model,
@@ -230,8 +275,6 @@ def train_model(
 
 
 def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
-    model.eval()
-
     test_loss = evaluate(
         model=model,
         dataset=test_set,

mlx_lm/manage.py

@@ -2,7 +2,22 @@ import argparse
 from typing import List, Union
 
 from huggingface_hub import scan_cache_dir
-from transformers.commands.user import tabulate
+
+
+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)
 
 
 def ask_for_confirmation(message: str) -> bool:

mlx_lm/models/base.py

@@ -33,29 +33,33 @@ 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 * -1e9
+        mask = mask & (rinds < lengths)
+    return mask
 
 
-def create_attention_mask(h: mx.array, cache: Optional[Any] = None):
+def create_attention_mask(
+    h: mx.array, cache: Optional[Any] = None, return_array: bool = False
+):
     T = h.shape[1]
     if T > 1:
-        window_size = None
         offset = 0
+        window_size = None
         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
-            else:
-                offset = c.offset
-        mask = create_causal_mask(T, offset, window_size=window_size)
-        mask = mask.astype(h.dtype)
+                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:
         mask = None
     return mask

mlx_lm/models/cohere2.py

@@ -83,15 +83,22 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        if self.use_sliding_window and mask is not None:
+        if self.use_sliding_window and isinstance(mask, mx.array):
             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, keys, values, cache=cache, scale=self.scale, mask=mask
-        )
-
+            queries.astype(sdpa_type),
+            keys,
+            values,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,
+        ).astype(queries.dtype)
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -126,9 +133,11 @@ 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
 
 
@@ -156,14 +165,27 @@ 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)
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
+        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)
 

mlx_lm/models/deepseek_v2.py

@@ -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)
+            group_scores = scores.max(axis=-1, keepdims=True)
             k = self.n_group - self.topk_group
-            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
+            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
+            )
             scores = scores.reshape(bsz, seq_len, -1)
 
         k = self.top_k
@@ -364,8 +364,32 @@ 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,
@@ -374,14 +398,31 @@ class DeepseekV2Model(nn.Module):
     ) -> mx.array:
         h = self.embed_tokens(x)
 
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+        # Hack to avoid time-outs during prompt-processing
+        dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
         if mask is None:
             mask = create_attention_mask(h, cache)
 
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = [None] * self.num_layers
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1), stream=dist_stream)
+
+        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, stream=dist_stream
+            )
+
+        # Broadcast h while keeping it in the graph
+        h = mx.distributed.all_gather(h, stream=dist_stream)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -418,4 +459,4 @@ class Model(nn.Module):
 
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.layers[self.model.start_idx : self.model.end_idx]

mlx_lm/models/deepseek_v3.py

@@ -0,0 +1,535 @@
+# 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 * base ** (
+            mx.arange(0, dim, 2, dtype=mx.float32) / dim
+        )
+        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)
+            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)
+        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))
+    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(scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
+        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
+        # Hack to avoid time-outs during prompt-processing
+        dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
+        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), stream=dist_stream)
+
+        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, stream=dist_stream
+            )
+
+        # Broadcast h while keeping it in the graph
+        h = mx.distributed.all_gather(h, stream=dist_stream)[: 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):
+            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)
+            )
+            scale_inv = scale_inv.astype(weight.dtype)
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n]
+
+        # 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]

mlx_lm/models/gemma2.py

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

mlx_lm/models/gemma3.py

@@ -0,0 +1,61 @@
+# 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

@@ -0,0 +1,238 @@
+# 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
+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 = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, 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))
+
+
+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 = x + self.post_attention_layernorm(r)
+        r = self.mlp(self.pre_feedforward_layernorm(h))
+        out = 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/gpt_bigcode.py

@@ -145,16 +145,16 @@ class GPTBigCodeModel(nn.Module):
         hidden_states = self.wte(inputs)
 
         mask = None
-        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 mask is not None and hidden_states.shape[1] > 1:
+            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

@@ -0,0 +1,195 @@
+# 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

@@ -0,0 +1,185 @@
+# 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,7 +76,6 @@ 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(
@@ -107,7 +106,6 @@ 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
@@ -198,7 +196,10 @@ class DecoderLayer(nn.Module):
         super().__init__()
         self.hidden_size = args.hidden_size
         self.self_attn = Attention(kv_proj, args)
-        self.mlp = MoeBlock(args)
+        if args.num_experts == 1:
+            self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        else:
+            self.mlp = MoeBlock(args)
 
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
         self.post_attention_layernorm = nn.RMSNorm(
@@ -231,7 +232,10 @@ 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=(i % args.cla_share_factor) == 0)
+            DecoderLayer(
+                args=args,
+                kv_proj=(not args.use_cla) or (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)
@@ -251,7 +255,7 @@ class HunYuanModel(nn.Module):
             cache = [None] * len(self.layers)
 
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            if i % self.args.cla_share_factor == 0:
+            if (not self.args.use_cla) or i % self.args.cla_share_factor == 0:
                 shared_kv_states = None
             h, shared_kv_states = layer(h, mask, c, shared_kv_states)
 
@@ -275,6 +279,29 @@ 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

@@ -0,0 +1,241 @@
+# 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/llama.py

@@ -69,12 +69,14 @@ 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 = 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)
+        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)
@@ -88,7 +90,7 @@ class Attention(nn.Module):
             queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
-        output = output.transpose(0, 2, 1, 3).flatten(-2, -1)
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
 
@@ -194,9 +196,12 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         # Remove unused precomputed rotary freqs
-        return {
+        weights = {
             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/mamba.py

@@ -1,4 +1,4 @@
-# Copyright © 2024 Apple Inc.
+# Copyright © 2024-2025 Apple Inc.
 
 import math
 from dataclasses import dataclass
@@ -123,17 +123,16 @@ class MambaBlock(nn.Module):
             self.intermediate_size, self.hidden_size, bias=args.use_bias
         )
 
-    def ssm_step(self, x, state=None):
-        A = -mx.exp(self.A_log)
+    def ssm_step(self, x, A, state=None):
         D = self.D
         deltaBC = self.x_proj(x)
-        delta, B, C = mx.split(
-            deltaBC,
-            indices_or_sections=[
-                self.time_step_rank,
-                self.time_step_rank + self.ssm_state_size,
-            ],
-            axis=-1,
+        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,
+            ),
         )
         if self.use_bcdt_rms:
             delta, B, C = map(self.mixer_norm, (delta, B, C))
@@ -145,25 +144,40 @@ class MambaBlock(nn.Module):
         y = y + D * x
         return y, new_state
 
-    def __call__(self, x, cache):
+    def _process_sequence(self, x, conv_cache, state_cache):
         B, T, D = x.shape
-        if cache is None:
-            cache = [None, None]
+        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)
 
         outputs = []
+        current_state = state_cache
+        y = []
         for t in range(T):
-            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)
+            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
+
         return output
 
 

mlx_lm/models/minicpm.py

@@ -1,11 +1,10 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2023-2025 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, 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
 
@@ -138,9 +137,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 / np.sqrt(self.num_hidden_layers))
+        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 / np.sqrt(self.num_hidden_layers))
+        out = h + r * (self.scale_depth / self.num_hidden_layers**0.5)
         return out
 
 

mlx_lm/models/minicpm3.py

@@ -0,0 +1,250 @@
+# 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/olmoe.py

@@ -0,0 +1,217 @@
+# 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,8 +23,10 @@ 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:
@@ -59,9 +61,10 @@ 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(
-                head_dim,
+                rope_dim,
                 base=args.rope_theta,
                 max_position_embeddings=args.max_position_embeddings,
                 original_max_position_embeddings=args.original_max_position_embeddings,
@@ -74,7 +77,7 @@ class Attention(nn.Module):
                 assert isinstance(args.rope_scaling["factor"], float)
                 rope_scale = 1 / args.rope_scaling["factor"]
             self.rope = nn.RoPE(
-                head_dim,
+                rope_dim,
                 traditional=args.rope_traditional,
                 base=args.rope_theta,
                 scale=rope_scale,
@@ -190,7 +193,8 @@ class Model(nn.Module):
         super().__init__()
         self.model_type = args.model_type
         self.model = Phi3Model(args)
-        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)
         self.args = args
 
     def __call__(
@@ -200,7 +204,11 @@ class Model(nn.Module):
         cache=None,
     ):
         out = self.model(inputs, mask, cache)
-        return self.lm_head(out)
+        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):

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)
+            mask = create_attention_mask(h, cache, return_array=True)
 
         if cache is None:
             cache = [None] * len(self.layers)

mlx_lm/models/plamo2.py

@@ -0,0 +1,609 @@
+# 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 _rms_norm(hidden_states: mx.array, eps: float) -> mx.array:
+    input_dtype = hidden_states.dtype
+    hidden_states = hidden_states.astype(mx.float32)
+    variance = mx.power(hidden_states, 2).mean(-1, keepdims=True)
+    hidden_states = hidden_states * mx.rsqrt(variance + eps)
+    hidden_states = hidden_states.astype(input_dtype)
+
+    return hidden_states
+
+
+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 = _rms_norm(q, 1e-6) * self.q_weight[:, None]
+        k = _rms_norm(k, 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,6 +7,7 @@ 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
@@ -18,24 +19,13 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     rms_norm_eps: float
     vocab_size: int
-    num_key_value_heads: Optional[int] = None
+    num_key_value_heads: int
+    max_position_embeddings: int = 32768
     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):
@@ -54,16 +44,12 @@ 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)
 
-        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(
+        self.rope = initialize_rope(
             head_dim,
-            traditional=args.rope_traditional,
             base=args.rope_theta,
-            scale=rope_scale,
+            traditional=args.rope_traditional,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
         )
 
     def __call__(

mlx_lm/models/rope_utils.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import math
 from typing import Optional
 
 import mlx.core as mx
@@ -61,6 +62,78 @@ 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,
@@ -87,5 +160,25 @@ 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,
     ):
         """
-        Phi3Su Scaled Rotary Embedding layer for Phi-3 models.
+        Su Scaled Rotary Embedding layer.
 
         Args:
             dims (int): The feature dimensions to be rotated.
@@ -51,11 +51,13 @@ 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(
-            self.scale * x,
-            x.shape[-1],
+            x,
+            self.dim,
             traditional=False,
             base=None,
             scale=1.0,

mlx_lm/sample_utils.py

@@ -35,14 +35,25 @@ def make_sampler(
     """
     if temp == 0:
         return lambda x: mx.argmax(x, axis=-1)
-    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)
+
+    # 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))
+
+    # 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
 
 
 def make_logits_processors(
@@ -85,10 +96,9 @@ def make_logits_processors(
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def top_k_sampling(
+def apply_top_k(
     logprobs: mx.array,
     top_k: int,
-    temperature=1.0,
 ) -> mx.array:
     """
     Sample from only the top K tokens ranked by probability.
@@ -103,20 +113,18 @@ def top_k_sampling(
             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 mx.random.categorical(masked_logprobs, axis=-1)
+    return masked_logprobs
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def min_p_sampling(
+def apply_min_p(
     logprobs: mx.array,
     min_p: float,
     min_tokens_to_keep: int = 1,
-    temperature=1.0,
 ) -> mx.array:
     """
     Apply min-p sampling to the logprobs.
@@ -144,14 +152,12 @@ def min_p_sampling(
         )
     # 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).squeeze(0)
-    sorted_logprobs = logprobs[..., sorted_indices]
+    sorted_indices = mx.argsort(-logprobs, axis=-1)
+    sorted_logprobs = mx.take_along_axis(logprobs, sorted_indices, axis=-1)
 
     # Top probability
-    top_logprobs = logprobs[..., sorted_indices[0]]
+    top_logprobs = sorted_logprobs[:, 0:1]
 
     # Calculate the min_p threshold
     scaled_min_p = top_logprobs + math.log(min_p)
@@ -163,29 +169,35 @@ def min_p_sampling(
     # Create pool of tokens with probability less than scaled min_p
     selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)
 
-    # Return sampled token
-    sorted_token = mx.random.categorical(selected_logprobs)
-    return sorted_indices[sorted_token]
+    # 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
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.array:
+def apply_top_p(logits: mx.array, top_p: 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 * (1 / temperature), axis=-1)
+    probs = mx.softmax(logits, axis=-1)
 
     # sort probs in ascending order
     sorted_indices = mx.argsort(probs, axis=-1)
-    sorted_probs = probs[..., sorted_indices.squeeze(0)]
+    sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
 
     cumulative_probs = mx.cumsum(sorted_probs, axis=-1)
 
@@ -196,10 +208,15 @@ def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.arr
         0,
     )
 
-    sorted_token = mx.random.categorical(mx.log(top_probs))
-    token = sorted_indices.squeeze(0)[sorted_token]
+    # 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)
 
-    return 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)

mlx_lm/server.py

@@ -4,6 +4,7 @@ import argparse
 import json
 import logging
 import platform
+import socket
 import time
 import uuid
 import warnings
@@ -26,9 +27,10 @@ import mlx.core as mx
 from huggingface_hub import scan_cache_dir
 
 from ._version import __version__
-from .models.cache import make_prompt_cache
+from .generate import stream_generate
+from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
 from .sample_utils import make_logits_processors, make_sampler
-from .utils import load, stream_generate
+from .utils import load
 
 
 def get_system_fingerprint():
@@ -114,6 +116,33 @@ 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)
@@ -425,14 +454,30 @@ class APIHandler(BaseHTTPRequestHandler):
 
     def get_prompt_cache(self, prompt):
         cache_len = len(self.prompt_cache.tokens)
+        prompt_len = len(prompt)
+        prefix_len = min(cache_len, prompt_len)
         if (
             self.prompt_cache.model_key != self.model_provider.model_key
-            or cache_len >= len(prompt)
-            or self.prompt_cache.tokens != prompt[:cache_len]
+            or prompt[:prefix_len] != self.prompt_cache.tokens[:prefix_len]
         ):
             self.prompt_cache.model_key = self.model_provider.model_key
             self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
+            self.prompt_cache.tokens = []
+        elif cache_len >= prompt_len:
+            # Trim the cache if it contains the prompt as a prefix. This case
+            # is useful for reusing the cache for multiple queries with a long
+            # prompt
+            if can_trim_prompt_cache(self.prompt_cache.cache):
+                num_to_trim = cache_len - prompt_len + 1
+                trim_prompt_cache(self.prompt_cache.cache, num_to_trim)
+                self.prompt_cache.tokens = self.prompt_cache.tokens[:-num_to_trim]
+                prompt = prompt[-1:]
+            else:
+                self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
+                self.prompt_cache.tokens = []
         else:
+            # Trim the prompt if it contains the cache as a prefix. This case
+            # is to avoid recomputing the cache in multi-turn chats.
             prompt = prompt[cache_len:]
         self.prompt_cache.tokens.extend(prompt)
         return prompt
@@ -591,8 +636,10 @@ 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(
-                body["messages"],
+                messages,
                 body.get("tools", None),
                 add_generation_prompt=True,
             )
@@ -665,6 +712,10 @@ 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(
@@ -719,13 +770,6 @@ 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,
@@ -744,11 +788,6 @@ 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))
 
 

mlx_lm/tokenizer_utils.py

@@ -1,5 +1,6 @@
 import json
 from functools import partial
+from typing import List
 
 from transformers import AutoTokenizer
 
@@ -351,7 +352,7 @@ def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
 
     tokenizer_file = model_path / "tokenizer.json"
     if tokenizer_file.exists():
-        with open(tokenizer_file, "r") as fid:
+        with open(tokenizer_file, "r", encoding="utf-8") as fid:
             tokenizer_content = json.load(fid)
         if "decoder" in tokenizer_content:
             if _is_spm_decoder(tokenizer_content["decoder"]):
@@ -368,3 +369,8 @@ def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
         detokenizer_class,
         eos_token_ids=eos_token_ids,
     )
+
+
+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

mlx_lm/tuner/datasets.py

@@ -1,11 +1,12 @@
 import json
+import types
 from pathlib import Path
-from typing import Dict, List
+from typing import Any, Dict, List, Optional
 
 from transformers import PreTrainedTokenizer
 
 
-class Dataset:
+class TextDataset:
     """
     Light-weight wrapper to hold a dataset.
     """
@@ -16,10 +17,15 @@ class Dataset:
         tokenizer: PreTrainedTokenizer,
         text_key: str = "text",
     ):
-        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)
+        self._data = [d for d in 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
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -34,14 +40,31 @@ class ChatDataset:
     https://platform.openai.com/docs/guides/fine-tuning/example-format
     """
 
-    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 __init__(
+        self,
+        data: List[Dict[str, str]],
+        tokenizer: PreTrainedTokenizer,
+        chat_key: str = "messages",
+        mask_prompt: bool = False,
+    ):
+        self._data = [d for d in 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(tokenizer.apply_chat_template(messages, tools=tools))
+            return (tokens, offset)
+        else:
+            return tokens
+
+    def itemlen(idx: int):
+        return len(self._data[idx])
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -61,18 +84,32 @@ class CompletionsDataset:
         self,
         data: List[Dict[str, str]],
         tokenizer: PreTrainedTokenizer,
-        prompt_key: str = "prompt",
-        completion_key: str = "completion",
+        prompt_key: str,
+        completion_key: str,
+        mask_prompt: bool,
     ):
-        self._data = [
-            tokenizer.apply_chat_template(
-                [
-                    {"role": "user", "content": d[prompt_key]},
-                    {"role": "assistant", "content": d[completion_key]},
-                ],
+        self._data = [d for d in 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]}]
+                )
             )
-            for d in data
-        ]
+            return (tokens, offset)
+
+        return tokens
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -81,15 +118,63 @@ class CompletionsDataset:
         return len(self._data)
 
 
-def create_dataset(data, tokenizer: PreTrainedTokenizer):
-    sample = data[0]
+class ConcatenatedDataset:
+    def __init__(self, data: List[Any]):
+        self._data = data
+        self._len = sum(len(d) for d in self._data)
 
-    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)
+    def __getitem__(self, idx: int):
+        for data in self._data:
+            j = idx - len(data)
+            if j < 0:
+                break
+            idx = j
+        return data[idx]
+
+    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")
+    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)
     else:
         raise ValueError(
             "Unsupported data format, check the supported formats here:\n"
@@ -97,20 +182,28 @@ def create_dataset(data, tokenizer: PreTrainedTokenizer):
         )
 
 
-def load_local_dataset(data_path: Path, tokenizer: PreTrainedTokenizer):
+def load_local_dataset(
+    data_path: Path,
+    tokenizer: PreTrainedTokenizer,
+    config,
+):
     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)
+        return create_dataset(data, tokenizer, config)
 
     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):
+def load_hf_dataset(
+    data_id: str,
+    tokenizer: PreTrainedTokenizer,
+    config,
+):
     from datasets import exceptions, load_dataset
 
     try:
@@ -119,7 +212,11 @@ def load_hf_dataset(data_id: str, tokenizer: PreTrainedTokenizer):
         names = ("train", "valid", "test")
 
         train, valid, test = [
-            create_dataset(dataset[n], tokenizer) if n in dataset.keys() else []
+            (
+                create_dataset(dataset[n], tokenizer, config)
+                if n in dataset.keys()
+                else []
+            )
             for n in names
         ]
 
@@ -132,42 +229,61 @@ def load_hf_dataset(data_id: str, tokenizer: PreTrainedTokenizer):
 def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
     import datasets
 
-    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):
+    def create_hf_dataset(dataset_name, config, split, hf_config):
         ds = datasets.load_dataset(
             dataset_name,
             split=split,
-            **hf_args.get("config", {}),
+            **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:
-            raise ValueError(
-                "Specify either a prompt and completion feature or a text "
-                "feature for the Hugging Face dataset."
+        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,
+            )
+        else:
+            train, valid = [], []
 
-    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 args.test:
+            test_split = ds.get("test_split")
+            test = create_hf_dataset(
+                ds_path,
+                config,
+                test_split,
+                hf_config,
+            )
+        else:
+            test = []
 
-    return train, valid, test
+        collection.append((train, valid, test))
+
+    if len(collection) == 1:
+        return collection[0]
+
+    # Otherwise concatenate them
+    return tuple(map(ConcatenatedDataset, zip(*collection)))
 
 
 def load_dataset(args, tokenizer: PreTrainedTokenizer):
@@ -176,10 +292,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)
+            train, valid, test = load_local_dataset(data_path, tokenizer, args)
         else:
             print(f"Loading Hugging Face dataset {args.data}.")
-            train, valid, test = load_hf_dataset(args.data, tokenizer)
+            train, valid, test = load_hf_dataset(args.data, tokenizer, args)
 
     if args.train and len(train) == 0:
         raise ValueError(

mlx_lm/tuner/trainer.py

@@ -4,14 +4,27 @@ import glob
 import shutil
 import time
 from dataclasses import dataclass, field
+from functools import partial
 from pathlib import Path
-from typing import Union
+from typing import List, Optional, Tuple
 
 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 mlx.utils import tree_flatten, tree_map
+from transformers import PreTrainedTokenizer
+
+from ..models.cache import KVCache, make_prompt_cache
+from .datasets import CacheDataset
+
+
+def reset_prompt_cache(cache):
+    for e, c in enumerate(cache):
+        if isinstance(c, KVCache):
+            cache[e] = KVCache()
+        else:
+            raise ValueError("Unsupported cache")
 
 
 def grad_checkpoint(layer):
@@ -61,24 +74,43 @@ class TrainingArgs:
         default=False,
         metadata={"help": "Use gradient checkpointing to reduce memory use."},
     )
+    seq_step_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "The examples are processsed in seq_step_size chunks."},
+    )
 
 
-def default_loss(model, inputs, targets, lengths):
-    logits = model(inputs)
+def default_loss(model, batch, lengths, cache=None):
+    inputs = batch[:, :-1]
+    targets = batch[:, 1:]
+
+    offset = cache[0].offset if cache is not None else 0
+    logits = model(inputs, cache=cache)
     logits = logits.astype(mx.float32)
 
-    length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]
+    steps = mx.arange(1, targets.shape[1] + 1) + offset
+    mask = mx.logical_and(steps >= lengths[:, 0:1], steps <= lengths[:, 1:])
 
-    ce = nn.losses.cross_entropy(logits, targets) * length_mask
-    ntoks = length_mask.sum()
+    ce = nn.losses.cross_entropy(logits, targets) * mask
+    ntoks = mask.sum()
     ce = ce.sum() / ntoks
 
     return ce, ntoks
 
 
-def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False):
+def iterate_batches(
+    dataset,
+    tokenizer,
+    batch_size,
+    max_seq_length,
+    train=False,
+):
     # Sort by length:
-    idx = sorted(range(len(dataset)), key=lambda idx: len(dataset[idx]))
+    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)
     if len(dataset) < batch_size:
         raise ValueError(
             f"Dataset must have at least batch_size={batch_size}"
@@ -101,6 +133,10 @@ def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False)
         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(
@@ -109,9 +145,9 @@ def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False)
                     "Consider pre-splitting your data to save memory."
                 )
 
-            # 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)
+            # 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)
             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)
@@ -123,8 +159,7 @@ def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False)
                     truncated_length  # Update lengths to match truncated lengths
                 )
             batch = mx.array(batch_arr)
-
-            yield batch[:, :-1], batch[:, 1:], mx.array(lengths)
+            yield batch, mx.array(list(zip(offsets, lengths)))
 
         if not train:
             break
@@ -139,12 +174,17 @@ def evaluate(
     max_seq_length=2048,
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
+    seq_step_size: Optional[int] = None,
 ):
-    all_losses = 0
-    ntokens = 0
+    model.eval()
+    all_losses = mx.array(0.0)
+    ntokens = mx.array(0)
 
     index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
 
+    seq_step_size = seq_step_size or max_seq_length
+
+    cache = make_prompt_cache(model)
     for _, batch in zip(
         index_iterator,
         iterate_batches(
@@ -154,13 +194,18 @@ def evaluate(
             max_seq_length=max_seq_length,
         ),
     ):
-        losses, toks = loss(model, *batch)
-        all_losses += losses * toks
-        ntokens += toks
-        mx.eval(all_losses, ntokens)
+        seq_length = batch[0].shape[1]
+        for s in range(0, seq_length, seq_step_size):
+            local_batch = (batch[0][:, s : s + seq_step_size], batch[1])
+            losses, toks = loss(model, *local_batch, cache)
+            all_losses += losses * toks
+            ntokens += toks
+            if s + seq_step_size >= seq_length:
+                reset_prompt_cache(cache)
+            mx.eval(all_losses, ntokens)
 
-    all_losses = mx.distributed.all_sum(all_losses)
-    ntokens = mx.distributed.all_sum(ntokens)
+    all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
+    ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
 
     return (all_losses / ntokens).item()
 
@@ -187,6 +232,7 @@ 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()
@@ -197,8 +243,11 @@ def train(
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
 
-    state = [model.state, optimizer.state]
+    seq_step_size = args.seq_step_size or args.max_seq_length
+    cache = make_prompt_cache(model)
+    state = [model.state, optimizer.state, mx.random.state]
 
+    @partial(mx.compile, inputs=state, outputs=state)
     def step(batch):
         # Forward and backward pass
         (lvalue, toks), grad = loss_value_and_grad(model, *batch)
@@ -211,14 +260,52 @@ def train(
 
         return lvalue, toks
 
+    train_dataset = CacheDataset(train_dataset)
+    val_dataset = CacheDataset(val_dataset)
+
+    loss_value_and_grad = nn.value_and_grad(model, loss)
+
+    model.train()
+    seq_step_size = args.seq_step_size or args.max_seq_length
+
+    def seq_split_step(batch):
+        losses = mx.array(0.0)
+        n_tokens = mx.array(0.0)
+        seq_length = batch[0].shape[1]
+        grad_accum = None
+        for s in range(0, seq_length, seq_step_size):
+            local_batch = (batch[0][:, s : s + seq_step_size], batch[1])
+            (lvalue, toks), grad = loss_value_and_grad(model, *local_batch, cache)
+            prev_n_tokens = n_tokens
+            losses += toks * lvalue
+            n_tokens += toks
+
+            if grad_accum is None:
+                grad_accum = grad
+            else:
+                scale_g = toks / n_tokens
+                scale_acc = prev_n_tokens / n_tokens
+                grad_accum = tree_map(
+                    lambda g, acc: scale_g * g + scale_acc * acc, grad, grad_accum
+                )
+
+            # Let go of the prompt cache before the last eval
+            if s + seq_step_size >= seq_length:
+                reset_prompt_cache(cache)
+            mx.eval(grad_accum, losses, n_tokens)
+
+        grad_accum = average_gradients(grad_accum)
+        optimizer.update(model, grad_accum)
+        return losses / n_tokens, n_tokens
+
     loss_value_and_grad = nn.value_and_grad(model, loss)
 
     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(
@@ -229,10 +316,11 @@ def train(
             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:
-            stop = time.perf_counter()
+            tic = time.perf_counter()
             val_loss = evaluate(
                 model=model,
                 dataset=val_dataset,
@@ -242,8 +330,10 @@ def train(
                 num_batches=args.val_batches,
                 max_seq_length=args.max_seq_length,
                 iterate_batches=iterate_batches,
+                seq_step_size=seq_step_size,
             )
-            val_time = time.perf_counter() - stop
+            model.train()
+            val_time = time.perf_counter() - tic
             if rank == 0:
                 print(
                     f"Iter {it}: "
@@ -260,26 +350,30 @@ def train(
                 }
                 training_callback.on_val_loss_report(val_info)
 
-            start = time.perf_counter()
+            tic = time.perf_counter()
+
+        if batch[0].shape[1] > seq_step_size:
+            lvalue, toks = seq_split_step(batch)
+        else:
+            lvalue, toks = step(batch)
 
-        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:
-            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()
+            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()
             learning_rate = optimizer.learning_rate.item()
-            it_sec = args.steps_per_report / (stop - start)
-            tokens_sec = float(n_tokens) / (stop - start)
+            it_sec = args.steps_per_report / train_time
+            tokens_sec = float(n_tokens) / train_time
             trained_tokens += n_tokens
-            peak_mem = mx.metal.get_peak_memory() / 1e9
+            peak_mem = mx.get_peak_memory() / 1e9
             if rank == 0:
                 print(
                     f"Iter {it}: Train loss {train_loss:.3f}, "
@@ -306,10 +400,10 @@ def train(
             losses = 0
             n_tokens = 0
             steps = 0
-            start = time.perf_counter()
+            train_time = 0
 
         # Save adapter weights
-        if it % args.steps_per_save == 0:
+        if it % args.steps_per_save == 0 and rank == 0:
             adapter_weights = dict(tree_flatten(model.trainable_parameters()))
             mx.save_safetensors(str(args.adapter_file), adapter_weights)
             checkpoint = (
@@ -322,6 +416,7 @@ def train(
             )
 
     # Save final weights
-    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}.")
+    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}.")

mlx_lm/tuner/utils.py

@@ -52,11 +52,6 @@ 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)):
@@ -89,17 +84,25 @@ def linear_to_lora_layers(
         "mixtral",
         "nemotron",
         "stablelm",
+        "hunyuan",
         "qwen2",
         "qwen2_moe",
         "phimoe",
         "gemma",
         "gemma2",
+        "gemma3",
+        "gemma3_text",
+        "granite",
+        "helium",
         "starcoder2",
         "cohere",
         "cohere2",
         "minicpm",
+        "minicpm3",
         "deepseek",
         "olmo2",
+        "olmoe",
+        "internlm3",
     ]:
         keys = set(["self_attn.q_proj", "self_attn.v_proj"])
         if model.model_type in ["mixtral", "phimoe"]:
@@ -107,6 +110,8 @@ 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 == "olmoe":
+            keys.add("mlp.gate")
 
     elif model.model_type == "gpt_bigcode":
         keys = set(["attn.c_attn"])
@@ -126,7 +131,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":
+    elif model.model_type == "deepseek_v2" or model.model_type == "minicpm3":
         keys = set(
             [
                 "self_attn.q_proj",
@@ -150,7 +155,7 @@ def linear_to_lora_layers(
     else:
         raise ValueError(f"Lora does not support {model.model_type}")
 
-    for l in model.layers[-min(num_layers, 0) :]:
+    for l in model.layers[-max(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/utils.py

@@ -1,27 +1,40 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import contextlib
 import copy
 import glob
 import importlib
 import json
 import logging
-import shutil
-import time
-from dataclasses import dataclass
+import os
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 import mlx.core as mx
 import mlx.nn as nn
-from huggingface_hub import snapshot_download
+
+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 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
@@ -35,9 +48,6 @@ 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):
@@ -45,68 +55,6 @@ 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.
@@ -153,11 +101,12 @@ 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(
-                    repo_id=path_or_hf_repo,
+                    path_or_hf_repo,
                     revision=revision,
                     allow_patterns=[
                         "*.json",
@@ -165,7 +114,9 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
                         "*.py",
                         "tokenizer.model",
                         "*.tiktoken",
+                        "tiktoken.model",
                         "*.txt",
+                        "*.jsonl",
                     ],
                 )
             )
@@ -180,280 +131,6 @@ 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:
@@ -467,6 +144,7 @@ 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:
@@ -478,6 +156,8 @@ 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):
@@ -500,7 +180,7 @@ def load_model(
         # Try weight for back-compat
         weight_files = glob.glob(str(model_path / "weight*.safetensors"))
 
-    if not weight_files:
+    if not weight_files and strict:
         logging.error(f"No safetensors found in {model_path}")
         raise FileNotFoundError(f"No safetensors found in {model_path}")
 
@@ -534,7 +214,7 @@ def load_model(
             class_predicate=class_predicate,
         )
 
-    model.load_weights(list(weights.items()))
+    model.load_weights(list(weights.items()), strict=strict)
 
     if not lazy:
         mx.eval(model.parameters())
@@ -561,7 +241,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:
@@ -634,13 +314,18 @@ def upload_to_hub(path: str, upload_repo: str, hf_path: str):
     from . import __version__
 
     card = ModelCard.load(hf_path)
-    card.data.tags = ["mlx"] if card.data.tags is None else card.data.tags + ["mlx"]
+    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 = hf_path
     card.text = dedent(
         f"""
         # {upload_repo}
 
-        The Model [{upload_repo}](https://huggingface.co/{upload_repo}) was
+        This 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__}**.
 
@@ -655,7 +340,7 @@ def upload_to_hub(path: str, upload_repo: str, hf_path: str):
 
         model, tokenizer = load("{upload_repo}")
 
-        prompt="hello"
+        prompt = "hello"
 
         if tokenizer.chat_template is not None:
             messages = [{{"role": "user", "content": prompt}}]
@@ -798,6 +483,7 @@ 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()))
@@ -805,65 +491,3 @@ def save_config(
     # write the updated config to the config_path (if provided)
     with open(config_path, "w") as fid:
         json.dump(config, fid, indent=4)
-
-
-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,
-):
-    # 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)
-
-    weights = dict(tree_flatten(model.parameters()))
-    dtype = getattr(mx, dtype)
-    weights = {k: v.astype(dtype) 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_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")
-
-    if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo, hf_path)

requirements.txt

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

setup.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from setuptools import setup
 
 package_dir = Path(__file__).parent / "mlx_lm"
-with open(package_dir / "requirements.txt") as fid:
+with open("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-examples",
+    url="https://github.com/ml-explore/mlx-lm",
     license="MIT",
     install_requires=requirements,
     packages=["mlx_lm", "mlx_lm.models", "mlx_lm.tuner"],

tests/test_datsets.py

@@ -6,9 +6,10 @@ import tempfile
 import types
 import unittest
 
-from mlx_lm.tuner import datasets
 from transformers import AutoTokenizer
 
+from mlx_lm.tuner import datasets
+
 HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
 
 
@@ -43,7 +44,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.Dataset))
+        self.assertTrue(isinstance(train, datasets.TextDataset))
 
     def test_completions(self):
         data = {"prompt": "What is the capital of France?", "completion": "Paris."}
@@ -78,14 +79,15 @@ 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={
-                "name": "billsum",
-                "prompt_feature": "text",
-                "completion_feature": "summary",
-                "train_split": "train[:2%]",
-                "valid_split": "train[-2%:]",
-            },
+            hf_dataset=hf_args,
             test=False,
             train=True,
         )
@@ -97,6 +99,16 @@ 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,6 +11,7 @@ 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
@@ -21,7 +22,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: x)
+    setattr(obj, func, lambda x, **kwargs: x)
     yield
     setattr(obj, func, old_func)
 
@@ -369,11 +370,10 @@ class TestScheduleConfig(unittest.TestCase):
         mock_iterate_batches = MagicMock()
 
         mock_iterate_batches.return_value = [
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
         ]
 
         mock_default_loss.side_effect = [
@@ -411,9 +411,9 @@ class TestScheduleConfig(unittest.TestCase):
         mock_iterate_batches = MagicMock()
 
         mock_iterate_batches.return_value = [
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
         ]
 
         mock_default_loss.side_effect = [

tests/test_generate.py

@@ -1,17 +1,23 @@
 # Copyright © 2024 Apple Inc.
 
 import unittest
+from typing import List
 
-from mlx_lm.sample_utils import make_logits_processors
-from mlx_lm.utils import generate, load
+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
 
 
 class TestGenerate(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
-        cls.model, cls.tokenizer = load(HF_MODEL_PATH)
+        cls.HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
+        cls.model, cls.tokenizer = load(cls.HF_MODEL_PATH)
 
     def test_generate(self):
         # Simple test that generation runs
@@ -51,6 +57,34 @@ 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)
+
+        for generation_result in stream_generate(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            prompt="hello",
+            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), 5)
+        # 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,6 +5,7 @@ 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,6 +4,7 @@ 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
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
@@ -183,7 +184,7 @@ class TestModels(unittest.TestCase):
             self.assertEqual(outputs.shape, (1, 2, vocab_size))
             self.assertEqual(outputs.dtype, t)
 
-            if model_type != "mamba":
+            if model_type not in ("mamba", "plamo2"):
                 mask = create_causal_mask(inputs.shape[1], 0).astype(t)
                 outputs = model(inputs, mask=mask)
                 self.assertEqual(outputs.shape, (1, 2, vocab_size))
@@ -336,6 +337,7 @@ 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,
         )
@@ -372,6 +374,23 @@ 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
 
@@ -682,6 +701,43 @@ 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
 
@@ -701,6 +757,26 @@ 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
 
@@ -890,6 +966,23 @@ 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,6 +6,8 @@ import tempfile
 import unittest
 
 import mlx.core as mx
+
+from mlx_lm.generate import generate_step
 from mlx_lm.models.cache import (
     KVCache,
     MambaCache,
@@ -16,7 +18,7 @@ from mlx_lm.models.cache import (
     save_prompt_cache,
     trim_prompt_cache,
 )
-from mlx_lm.utils import generate_step, load
+from mlx_lm.utils import load
 
 HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
 
@@ -298,7 +300,7 @@ class TestPromptCache(unittest.TestCase):
         ):
             i += 1
             self.assertEqual(tok, toks[i])
-            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=2e-2))
+            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=4e-2))
 
 
 if __name__ == "__main__":

tests/test_sample_utils.py

@@ -1,67 +1,98 @@
 import unittest
 
 import mlx.core as mx
-from mlx_lm.sample_utils import min_p_sampling, top_k_sampling, top_p_sampling
+
+from mlx_lm.sample_utils import apply_min_p, apply_top_k, apply_top_p
 
 
 class TestSampleUtils(unittest.TestCase):
-    def test_top_p_sampling(self):
+    def test_apply_top_p(self):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
-        temperature = 1.0
 
-        token = top_p_sampling(logits, 0.3, temperature).item()
-        self.assertEqual(token, 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.95, temperature).item()
-        self.assertTrue(token in (0, 3))
+        new_logits = apply_top_p(logits, 0.95)
+        actual_probs = mx.softmax(new_logits.squeeze())
+        self.assertTrue(mx.allclose(probs.squeeze(), actual_probs))
 
         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])
 
-        token = top_p_sampling(logits, 0.4, temperature).item()
-        self.assertEqual(token, 1)
+        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.6, temperature).item()
-        self.assertTrue(token in (1, 2))
+        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.95, temperature).item()
-        self.assertTrue(token in (1, 2, 3))
+        # 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]]
+        )
 
-    def test_min_p_sampling(self):
+    def test_apply_min_p(self):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
-        temperature = 1.0
-        token = min_p_sampling(logits, 0.8)
-        self.assertEqual(token, 0)
+        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])
 
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
-        temperature = 1.0
-        for _ in range(5):
-            token = min_p_sampling(logits, 0.05)
-            self.assertTrue(token in (0, 3))
+        new_logits = apply_min_p(logits, 0.05)
+        actual_probs = mx.softmax(new_logits.squeeze())
+        self.assertTrue(mx.allclose(actual_probs, mx.squeeze(probs)))
 
-    def test_top_k_sampling(self):
+        # 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):
         probs = mx.array([0.9, 0.0, 0.0, 0.1])[None]
         logits = mx.log(probs)
 
-        token = top_k_sampling(logits, 1).item()
-        self.assertEqual(token, 0)
+        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])
 
-        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})
+        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]
+        )
 
         # 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)
 
-        tokens = top_k_sampling(logits, 1)
-        self.assertEqual(tokens.tolist(), [0, 1])
+        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]]
+        )
 
 
 if __name__ == "__main__":

tests/test_server.py

@@ -6,6 +6,7 @@ import threading
 import unittest
 
 import requests
+
 from mlx_lm.server import APIHandler
 from mlx_lm.utils import load
 
@@ -80,6 +81,29 @@ 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)

tests/test_tokenizers.py

@@ -4,6 +4,7 @@ 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,6 +6,7 @@ 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,7 +7,8 @@ import unittest
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.utils import tree_flatten
-from mlx_lm import utils
+
+from mlx_lm import convert, utils
 
 HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
 
@@ -76,14 +77,14 @@ class TestUtils(unittest.TestCase):
     def test_convert(self):
         mlx_path = os.path.join(self.test_dir, "mlx_model")
 
-        utils.convert(HF_MODEL_PATH, mlx_path=mlx_path, quantize=True)
+        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")
-        utils.convert(HF_MODEL_PATH, mlx_path=mlx_path, dtype="bfloat16")
+        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,6 +2,7 @@ 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
 
@@ -17,7 +18,7 @@ class TestLoadModelCustomGetClasses(unittest.TestCase):
                 self.config = args
                 self.custom_attribute = "This is a custom model"
 
-            def load_weights(self, weights):
+            def load_weights(self, weights, **kwargs):
                 self.qwenWeights = weights
 
         class CustomQwenConfig: