Initial example

* gptq

* speedup

.circleci/config.yml

@@ -20,7 +20,7 @@ jobs:
   mlx_lm_build_and_test:
     macos:
       xcode: "15.2.0"
-    resource_class: macos.m1.large.gen1
+    resource_class: m2pro.medium
     steps:
       - checkout
       - run:
@@ -30,6 +30,7 @@ jobs:
             python3.9 -m venv env
             source env/bin/activate
             pip install --upgrade pip
+            pip install sentencepiece
             pip install unittest-xml-reporting
             pip install -e ".[test]"
       - run:
@@ -40,6 +41,30 @@ jobs:
       - store_test_results:
           path: test-results
 
+  build_release:
+    macos:
+      xcode: "15.2.0"
+    resource_class: m2pro.medium
+    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 build
+            pip install twine
+      - run:
+          name: Build and upload
+          command: |
+            source env/bin/activate
+            python -m build
+            twine upload dist/*
+      - store_artifacts:
+          path: dist/
+
 workflows:
   build_and_test:
     when:
@@ -50,6 +75,15 @@ workflows:
       - mlx_lm_build_and_test
       - linux_build_and_test
 
+  build_pypi_release:
+    jobs:
+      - build_release:
+          filters:
+            tags:
+              only: /^v.*/
+            branches:
+              ignore: /.*/
+
   prb:
     when:
       matches:

ACKNOWLEDGMENTS.md

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

mlx_lm/LEARNED_QUANTS.md

@@ -5,13 +5,15 @@ To reduce the quality loss from quantization MLX LM has several options:
 - Distilled Weight Quantization (DWQ)
 - Activation-aware Weight Quantization (AWQ)[^1]
 - Dynamic quantization
+- GPT Quantization (GPTQ)[^2]
 
 All methods use calibration data to tune parameters or hyper-parameters of the
 model. DWQ fine-tunes non-quantized parameters (including quantization scales
 and biases) using the non-quantized model as a teacher. AWQ scales and clips
 the weights prior to quantization. Dynamic quantization estimates the
 sensitivity of a model's outputs to each layer and uses a higher precision for
-layers which have higher sensitivity.
+layers which have higher sensitivity. GPTQ finds quantized weights which
+minimize the squared error of each layer's output given the provided input.
 
 Dynamic quantization is the fastest to run. DWQ takes longer but typically
 yields better results. You can also cascade methods. For example a dynamically
@@ -28,7 +30,7 @@ pip install mlx-lm[quant]
 Use `mlx_lm.dwq` to run DWQ on a given model. For example:
 
 ```bash
-mlx_lm.dwq --model mistralai/Mistral-7B-Instruct-v0.3
+mlx_lm.dwq --model Qwen/Qwen3-0.6B
 ```
 
 Some important options, along with their default values are:
@@ -77,7 +79,7 @@ Use `mlx_lm.dynamic_quant` to generate a dynamic quantization of given model.
 For example:
 
 ```bash
-mlx_lm.dynamic_quant --model mistralai/Mistral-7B-Instruct-v0.3
+mlx_lm.dynamic_quant --model Qwen/Qwen3-0.6B
 ```
 
 The script will estimate the sensitivity for each quantizable layer in the
@@ -101,7 +103,7 @@ Some important options are:
 Use `mlx_lm.awq` to run AWQ on a given model. For example:
 
 ```bash
-mlx_lm.awq --model mistralai/Mistral-7B-Instruct-v0.3
+mlx_lm.awq --model Qwen/Qwen3-0.6B
 ```
 
 The script can take anywhere form a few minutes to several hours to run
@@ -122,10 +124,27 @@ For a full list of options run:
 mlx_lm.awq --help
 ```
 
+### GPTQ
+
+Use `mlx_lm.gptq` to run GPTQ on a given model. For example:
+
+```bash
+mlx_lm.awq --model Qwen/Qwen3-0.6B
+```
+
+The script can take anywhere from a few minutes to several hours depending on
+the model size.
+
+Some important options, along with their default values, are:
+
+- `--mlx-path mlx_model`: The location to save the AWQ model.
+- `--bits 4`: Precision of the quantization.
+
+
 ### Evaluate
 
-Once the training script finishes, you can evaluate the quality of the model
-on downstream tasks using `mlx_lm.evaluate`. For example:
+Once the quantization training finishes, you can evaluate the quality of the
+model on downstream tasks using `mlx_lm.evaluate`. For example:
 
 ```bash
 mlx_lm.evaluate \
@@ -146,4 +165,6 @@ mlx_lm.upload \
 
 [^1]: Refer to the [paper](https://arxiv.org/abs/2306.00978)
 and [github repository](https://github.com/mit-han-lab/llm-awq) for more
-details.
+details on AWQ.
+[^2]: Refer to the [paper](https://arxiv.org/abs/2210.17323) for more details
+on GPTQ.

mlx_lm/UPLOAD.md

@@ -1,37 +0,0 @@
-### Packaging for PyPI
-
-Install `build` and `twine`:
-
-```
-pip install --user --upgrade build
-pip install --user --upgrade twine
-```
-
-Generate the source distribution and wheel:
-
-```
-python -m build
-```
-
-> [!warning]
-> Use a test server first
-
-#### Test Upload
-
-Upload to test server:
-
-```
-python -m twine upload --repository testpypi dist/*
-```
-
-Install from test server and check that it works:
-
-```
-python -m pip install --index-url https://test.pypi.org/simple/ --no-deps mlx-lm
-```
-
-#### Upload
-
-```
-python -m twine upload dist/*
-```

mlx_lm/__main__.py

@@ -8,6 +8,7 @@ if __name__ == "__main__":
         "quant.awq",
         "quant.dwq",
         "quant.dynamic_quant",
+        "quant.gptq",
         "cache_prompt",
         "chat",
         "convert",

mlx_lm/_version.py

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

mlx_lm/convert.py

@@ -21,7 +21,6 @@ from .utils import (
 def mixed_quant_predicate_builder(
     recipe: str, model: nn.Module
 ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
-
     high_bits = 6
     group_size = 64
 
@@ -56,12 +55,6 @@ def mixed_quant_predicate_builder(
         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
-        if module.weight.shape[1] % group_size != 0:
-            return False
-
         index = (
             int(path.split(".")[layer_location])
             if len(path.split(".")) > layer_location
@@ -102,6 +95,7 @@ def convert(
     quant_predicate: Optional[
         Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
     ] = None,
+    trust_remote_code: bool = False,
 ):
     # Check the save path is empty
     if isinstance(mlx_path, str):
@@ -115,18 +109,12 @@ def convert(
 
     print("[INFO] Loading")
     model_path, hf_path = get_model_path(hf_path, revision=revision)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
-
-    def base_quant_predicate(path, module, config):
-        if not hasattr(module, "to_quantized"):
-            return False
-        if module.weight.shape[1] % q_group_size != 0:
-            return False
-        return True
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=trust_remote_code
+    )
 
     if isinstance(quant_predicate, str):
         quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
-    quant_predicate = quant_predicate or base_quant_predicate
 
     if dtype is None:
         dtype = config.get("torch_dtype", None)

mlx_lm/evaluate.py

@@ -367,8 +367,8 @@ def main():
     parser.add_argument(
         "--chat-template-args",
         type=json.loads,
-        help="""A JSON formatted string of arguments for the tokenizer's "
-        "apply_chat_template, e.g. '{"enable_thinking":false}'""",
+        help="""A JSON formatted string of arguments for the tokenizer's
+        apply_chat_template, e.g. '{"enable_thinking":false}'""",
         default="{}",
     )
 

mlx_lm/examples/afm/README.md

@@ -0,0 +1,10 @@
+Apple Foundation Model in MLX
+=============================
+
+This example provides information about porting the AFM model to MLX-LM and
+training adapters with it or using it as any other open-weights model. It is
+paired with https://developer.apple.com/apple-intelligence/foundation-models-adapter/ that
+was published during WWDC 25 and to get the weights one needs to follow these
+instructions to download the toolkit.
+
+

mlx_lm/examples/afm/extract_quantized.py

@@ -0,0 +1,38 @@
+import argparse
+
+import mlx.core as mx
+
+from mlx_lm.convert import convert
+
+
+def mixed_quant(layer_path, layer, cfg):
+    if "embedding" in layer_path:
+        return {"group_size": 32, "bits": 8}
+    return hasattr(layer, "to_quantized")
+
+
+def main(argv):
+    parser = argparse.ArgumentParser(
+        description="Quantize the AFM according to its original quantization"
+    )
+    parser.add_argument("source", help="The mlx model containing the fp32 weights")
+    parser.add_argument("destination", help="The folder to save the quantized model to")
+    parser.add_argument("--copy-adapters", action="store_true")
+    parser.add_argument(
+        "--dtype", choices=["bfloat16", "float16", "float32"], default="float32"
+    )
+    args = parser.parse_args(argv)
+
+    convert(
+        args.source,
+        args.destination,
+        quantize=True,
+        q_group_size=128,
+        q_bits=2,
+        dtype=getattr(mx, args.dtype),
+        quant_predicate=mixed_quant,
+    )
+
+
+if __name__ == "__main__":
+    main(None)

mlx_lm/examples/afm/extract_weights.py

@@ -0,0 +1,249 @@
+import argparse
+import json
+import textwrap
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+from transformers import LlamaTokenizerFast
+
+
+def share_data(a, b):
+    return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr()
+
+
+def get_model_config():
+    return {
+        "model_type": "afm7",
+        "vocab_size": 153600,
+        "hidden_dim": 2048,
+        "num_layers": 56,
+        "num_kv_reuse_layers": 21,
+        "num_heads": 16,
+        "num_kv_heads": 2,
+        "hidden_dim_scale_factor": 3.25,
+        "rope_theta": 500000.0,
+    }
+
+
+def get_adapter_config():
+    return {
+        "num_layers": 56,
+        "lora_parameters": {
+            "rank": 32,
+            "scale": 0.5,
+            "dropout": 0.0,
+            "keys": [
+                "mlp.gate_proj",
+                "mlp.down_proj",
+                "mlp.up_proj",
+                "self_attn.qkv_proj",
+                "self_attn.q_proj",
+                "self_attn.out_proj",
+            ],
+        },
+    }
+
+
+def get_chat_template():
+    return textwrap.dedent(
+        """
+        {%- set default_system_message = "A conversation between a user and a helpful assistant." %}
+        {%- if messages[0]['role'] == 'system' %}
+            {%- set system_message = messages[0]['content'] %}
+            {%- set loop_messages = messages[1:] %}
+        {%- else %}
+            {%- set system_message = default_system_message %}
+            {%- set loop_messages = messages %}
+        {%- endif %}
+        {{- '<turn_start> system<n>' + system_message -}}
+        {% if tools %}
+            {{- ('<n>system tools: ' + (tools | map('tojson') | join('<n>'))) -}}
+        {% endif %}
+        {{- '<turn_end>' -}}
+        {% for message in loop_messages %}
+            {{- '<turn_start> ' + message['role'] + '<n>' + message['content'] + '<turn_end>' -}}
+        {% endfor %}
+        {% if add_generation_prompt is defined and add_generation_prompt %}
+            {% if messages[-1]['role'] != 'assistant' %}
+                {{- '<turn_start> assistant<n>' -}}
+            {% endif %}
+        {% endif %}"""
+    ).strip()
+
+
+def map_model_keys(state):
+    model_keys = {}
+    for old in state:
+        if "adapter" in old:
+            continue
+        if "kv_quantizer" in old:
+            continue
+
+        new = old
+        if new.startswith("layers."):
+            new = new[7:]
+            new = new.replace("layer_", "")
+            new = new.replace("attention.norm", "input_layernorm")
+            new = new.replace(".attention.", ".self_attn.")
+            new = new.replace("self_attn.output_transform", "self_attn.out_proj")
+            new = new.replace("feed_forward.norm", "post_attention_layernorm")
+            new = new.replace(".feed_forward.", ".mlp.")
+            new = new.replace("hidden_transform.linear_0", "gate_proj")
+            new = new.replace("hidden_transform.linear_1", "up_proj")
+            new = new.replace("mlp.output_transform", "mlp.down_proj")
+            if new.startswith("segment_0"):
+                new = new.replace("segment_0", "layers")
+                new = new.replace(".qkv_transform.", ".qkv_proj.")
+                new = new.replace(".fused_linear.", ".")
+                new = new.replace(".qk_norm.query_norm.", ".q_norm.")
+                new = new.replace(".qk_norm.key_norm.", ".k_norm.")
+            elif new.startswith("segment_1"):
+                new = new.replace("segment_1", "kv_reuse_layers")
+                new = new.replace(".q_transform.", ".q_proj.")
+                new = new.replace(".q_norm.query_norm.", ".q_norm.")
+        new = new.replace(".wrapped.", ".")
+        new = "model." + new
+        model_keys[old] = new
+    return model_keys
+
+
+def map_adapter_keys(state):
+    adapter_keys = {}
+    for old in state:
+        if "adapter" not in old:
+            continue
+
+        new = old
+        new = new[7:]
+        new = new.replace("layer_", "")
+        new = new.replace(".attention.", ".self_attn.")
+        new = new.replace("self_attn.output_transform", "self_attn.out_proj")
+        new = new.replace(".feed_forward.", ".mlp.")
+        new = new.replace("hidden_transform.linear_0", "gate_proj")
+        new = new.replace("hidden_transform.linear_1", "up_proj")
+        new = new.replace("mlp.output_transform", "mlp.down_proj")
+        if new.startswith("segment_0"):
+            new = new.replace("segment_0", "layers")
+            new = new.replace(".qkv_transform.", ".qkv_proj.")
+            new = new.replace(".fused_linear.", ".")
+        elif new.startswith("segment_1"):
+            new = new.replace("segment_1", "kv_reuse_layers")
+            new = new.replace(".q_transform.", ".q_proj.")
+
+        new = new.replace(".lora_0.b_transpose", ".b_transpose.0")
+        new = new.replace(".lora_1.b_transpose", ".b_transpose.1")
+        new = new.replace(".lora_2.b_transpose", ".b_transpose.2")
+        new = new.replace(".lora_0.a_transpose", ".a_transpose.0")
+        new = new.replace(".lora_1.a_transpose", ".a_transpose.1")
+        new = new.replace(".lora_2.a_transpose", ".a_transpose.2")
+        new = new.replace("adapters.base_adapter.b_transpose", "lora_b")
+        new = new.replace("adapters.base_adapter.a_transpose", "lora_a")
+        new = "model." + new
+        adapter_keys[old] = new
+    return adapter_keys
+
+
+def add_kv_quant_weights(new_state, old_state, dt):
+    for k, v in old_state.items():
+        if "range" not in k:
+            continue
+
+        v = v.tolist()
+        weight = "quant_key_scale" if "key_quantizer" in k else "quant_value_scale"
+        new_k = k[: k.find("kv_quantizer")]
+        new_k = new_k.replace("segment_0.layer_", "")
+        new_k = new_k.replace("attention", "self_attn")
+        new_k = "model." + new_k + weight
+        quant_scale = torch.tensor(max(v[0] / (-128), v[1] / 127), dtype=dt)
+        new_state[new_k] = quant_scale
+
+
+def cast(x, dt):
+    info = torch.finfo(dt)
+    a, b = info.min, info.max
+    return x.clip(a, b).to(dt)
+
+
+def main(argv):
+    parser = argparse.ArgumentParser(
+        description="Map the PT weights to MLX-LM safetensors"
+    )
+    parser.add_argument("source", help="The source weights in PT format")
+    parser.add_argument("tokenizer", help="The source tokenizer file")
+    parser.add_argument("destination", help="The folder to write the model weights in")
+    parser.add_argument(
+        "--dtype", choices=["bfloat16", "float16", "float32"], default="float32"
+    )
+    parser.add_argument(
+        "--adapter-dtype", choices=["bfloat16", "float16", "float32"], default="float32"
+    )
+    parser.add_argument(
+        "--force",
+        "-f",
+        action="store_true",
+        help="If set overwrite the weight files in the destination folder",
+    )
+    args = parser.parse_args(argv)
+
+    destination = Path(args.destination)
+    if not destination.exists():
+        destination.mkdir()
+    model_file = destination / "model.safetensors"
+    adapter_file = destination / "adapters.safetensors"
+    if (model_file.exists() or adapter_file.exists()) and not args.force:
+        print("Model files already exist. Delete them or use --force to overwrite them")
+        return
+
+    # Write the configuration files
+    with (destination / "config.json").open("w") as f:
+        json.dump(get_model_config(), f, indent=4)
+    with (destination / "adapter_config.json").open("w") as f:
+        json.dump(get_adapter_config(), f, indent=4)
+
+    # Pop the tied output transform
+    state = torch.load(args.source)
+    if share_data(state["embedding.weight"], state["output_transform.weight"]):
+        state.pop("output_transform.weight")
+
+    # Map the weights
+    model_keys = map_model_keys(state)
+    adapter_keys = map_adapter_keys(state)
+
+    # Make the new weight dictionaries
+    dt = getattr(torch, args.dtype)
+    adapter_dt = getattr(torch, args.adapter_dtype)
+    adapters = {
+        k_new: cast(state[k_old], adapter_dt) for k_old, k_new in adapter_keys.items()
+    }
+    model = {k_new: cast(state[k_old], dt) for k_old, k_new in model_keys.items()}
+    add_kv_quant_weights(model, state, dt)
+
+    # Save them to disk
+    save_file(model, model_file)
+    save_file(adapters, adapter_file)
+
+    # Save the tokenizer
+    tok = LlamaTokenizerFast(vocab_file=args.tokenizer)
+    tok.chat_template = get_chat_template()
+    tok.eos_token_ids = tok.convert_tokens_to_ids("<turn_end>")
+    tok.save_pretrained(str(destination))
+    with (destination / "tokenizer_config.json").open("r+") as f:
+        config = json.load(f)
+        config["tokenizer_class"] = "NewlineTokenizer"
+        f.seek(0)
+        json.dump(config, f, indent=4)
+        f.truncate()
+    with (destination / "tokenizer.json").open("r+") as f:
+        tok = json.load(f)
+        tok["decoder"]["decoders"].insert(
+            1,
+            {"type": "Replace", "pattern": {"String": "<n>"}, "content": "\n"},
+        )
+        f.seek(0)
+        json.dump(tok, f, indent=4)
+        f.truncate()
+
+
+if __name__ == "__main__":
+    main(None)

mlx_lm/examples/afm/requirements.txt

@@ -0,0 +1,3 @@
+tamm==0.1.0
+transformers
+torch

mlx_lm/examples/openai_tool_use.py

@@ -52,7 +52,7 @@ response = client.chat.completions.create(
 
 # Call the function
 function = response.choices[0].message.tool_calls[0].function
-tool_result = functions[function.name](**function.arguments)
+tool_result = functions[function.name](**json.loads(function.arguments))
 
 # Put the result of the function in the messages and generate the final
 # response:

mlx_lm/generate.py

@@ -329,18 +329,20 @@ def generate_step(
            when ``kv_bits`` is non-None. Default: ``0``.
         prompt_progress_callback (Callable[int, int]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
-        input_embeddings (mx.array, optional): Input embeddings to use in place of
-           prompt tokens. Default: ``None``.
+        input_embeddings (mx.array, optional): Input embeddings to use in conjunction
+           with prompt tokens. Default: ``None``.
 
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
+    if len(prompt) == 0:
+        raise ValueError("Prompt must be non-empty.")
     if input_embeddings is not None:
         if not does_model_support_input_embeddings(model):
             raise ValueError("Model does not support input embeddings.")
-        if len(prompt) != 0:
+        elif prompt.shape[0] != input_embeddings.shape[0]:
             raise ValueError(
-                "If using input embeddings, prompt tokens must be an empty array."
+                "If using input embeddings, the sequence length must match that of the prompt."
             )
 
     tokens = None
@@ -363,47 +365,67 @@ def generate_step(
 
     sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
 
-    def _model_call(y):
-        if y.ndim == 3:
-            return model(None, cache=prompt_cache, input_embeddings=y)
+    def _model_call(input_tokens: mx.array, input_embeddings: Optional[mx.array]):
+        if input_embeddings is not None:
+            return model(
+                input_tokens, cache=prompt_cache, input_embeddings=input_embeddings
+            )
         else:
-            return model(y, cache=prompt_cache)
+            return model(input_tokens, cache=prompt_cache)
 
-    def _step(y):
+    def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None):
         nonlocal tokens
 
         with mx.stream(generation_stream):
-            logits = _model_call(y[None])
+            logits = _model_call(
+                input_tokens=input_tokens[None],
+                input_embeddings=(
+                    input_embeddings[None] if input_embeddings is not None else None
+                ),
+            )
 
             logits = logits[:, -1, :]
 
-            if logits_processors and input_embeddings is None:
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
+            if logits_processors:
+                tokens = (
+                    mx.concat([tokens, input_tokens])
+                    if tokens is not None
+                    else input_tokens
+                )
                 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)
+            sampled = sampler(logprobs)
+            return sampled, logprobs.squeeze(0)
 
-    using_embeddings = input_embeddings is not None
-
-    y = input_embeddings if using_embeddings else prompt
     with mx.stream(generation_stream):
-        total_prompt_tokens = y.shape[0]
+        total_prompt_tokens = prompt.shape[0]
         prompt_processed_tokens = 0
-        while y.shape[0] > prefill_step_size:
-            _model_call(y[:prefill_step_size][None])
+        while total_prompt_tokens - prompt_processed_tokens > prefill_step_size:
+            _model_call(
+                input_tokens=prompt[:prefill_step_size][None],
+                input_embeddings=(
+                    input_embeddings[:prefill_step_size][None]
+                    if input_embeddings is not None
+                    else None
+                ),
+            )
             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:]
+            prompt = prompt[prefill_step_size:]
+            input_embeddings = (
+                input_embeddings[prefill_step_size:]
+                if input_embeddings is not None
+                else input_embeddings
+            )
             mx.clear_cache()
 
-        y, logprobs = _step(y)
+        y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
 
     mx.async_eval(y, logprobs)
     n = 0

mlx_lm/models/bitlinear_layers.py

@@ -3,6 +3,33 @@
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.nn.layers.quantized import QuantizedLinear
+from mlx.utils import tree_flatten, tree_unflatten
+
+
+def bitnet_quantize(model, quantization_config: dict):
+    quantize_layers = []
+    modules_to_not_convert = quantization_config.get("modules_to_not_convert", [])
+    invert_weight_scales = (
+        quantization_config.get("linear_class", "") != "autobitlinear"
+    )
+
+    for name, module in tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module):
+
+        # Replace nn.Linear layers, but skip any layer from the `modules_to_not_convert` list
+        if name not in modules_to_not_convert and isinstance(module, nn.Linear):
+            old_weight = module.weight
+            out_features, in_features = old_weight.shape
+            bias = "bias" in module
+            new_layer = BitLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                invert_weight_scales=invert_weight_scales,
+            )
+            quantize_layers.append((name, new_layer))
+    if len(quantize_layers) > 0:
+        model.update_modules(tree_unflatten(quantize_layers))
+    return model
 
 
 def make_bitlinear_kernel():

mlx_lm/models/bitnet.py

@@ -112,6 +112,7 @@ class MLP(nn.Module):
         self.gate_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
         self.down_proj = BitLinear(hidden_dim, dim, bias=mlp_bias)
         self.up_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
+
         self.ffn_sub_norm = nn.RMSNorm(args.intermediate_size, eps=args.rms_norm_eps)
 
     def __call__(self, x) -> mx.array:

mlx_lm/models/ernie4_5_moe.py

@@ -0,0 +1,291 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    hidden_size: int
+    intermediate_size: int
+    model_type: str
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_hidden_layers: int
+    rms_norm_eps: float
+    vocab_size: int
+    rope_theta: float
+    use_bias: bool
+    tie_word_embeddings: bool
+    moe_num_experts: int
+    moe_layer_start_index: int = 0
+    moe_intermediate_size: int = 0
+    moe_capacity: list[int] = field(default_factory=list)
+    moe_k: int = 1
+    moe_layer_interval: int = 1
+    moe_use_aux_free: bool = False
+    moe_num_shared_experts: int = 0
+    moe_layer_end_index: Optional[int] = None
+    head_dim: Optional[int] = None
+    moe_gate_act: str = "softmax"
+
+
+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 dim // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.use_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.use_bias)
+
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=True,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class Ernie4_5_MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, use_bias=False):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=use_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Ernie4_5_MoeMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.k = args.moe_k
+        self.moe_intermediate_size = (
+            args.moe_intermediate_size
+            if args.moe_intermediate_size
+            else args.intermediate_size
+        )
+
+        self.gate = nn.Linear(args.hidden_size, args.moe_num_experts, bias=False)
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            self.moe_intermediate_size,
+            args.moe_num_experts,
+            bias=args.use_bias,
+        )
+
+        if getattr(args, "moe_num_shared_experts", 0) > 0:
+            shared_intermediate_size = (
+                args.moe_intermediate_size * args.moe_num_shared_experts
+                if getattr(args, "moe_intermediate_size", None)
+                else args.intermediate_size * args.moe_num_shared_experts
+            )
+            self.shared_experts = Ernie4_5_MLP(
+                args.hidden_size, shared_intermediate_size, args.use_bias
+            )
+        else:
+            self.shared_experts = None
+
+        if args.moe_gate_act == "softmax":
+            self.gate_act = nn.Softmax()
+        elif args.moe_gate_act == "sigmoid":
+            self.gate_act = nn.Sigmoid()
+        else:
+            raise ValueError(f"{args.moe_gate_act} is not supported.")
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.gate(x)
+        gates = self.gate_act(gates.astype(mx.float32))
+
+        k = self.k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+
+        scores = scores / mx.maximum(scores.sum(axis=-1, keepdims=True), 1e-12)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class Ernie4_5_DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(args)
+
+        moe_layer_start_index = (
+            min(args.moe_layer_start_index)
+            if isinstance(args.moe_layer_start_index, (tuple, list))
+            else args.moe_layer_start_index
+        )
+
+        if args.moe_layer_end_index is None:
+            moe_layer_end_index = args.num_hidden_layers - 1
+        else:
+            moe_layer_end_index = (
+                max(args.moe_layer_end_index)
+                if isinstance(args.moe_layer_end_index, (tuple, list))
+                else args.moe_layer_end_index
+            )
+
+        if (
+            ((layer_idx + 1) % args.moe_layer_interval == 0)
+            and layer_idx >= moe_layer_start_index
+            and layer_idx <= moe_layer_end_index
+        ):
+            self.mlp = Ernie4_5_MoeMLP(args)
+        else:
+            self.mlp = Ernie4_5_MLP(
+                args.hidden_size, args.intermediate_size, args.use_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))
+        return h + r
+
+
+class Ernie45Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Ernie4_5_DecoderLayer(args, i) for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Ernie45Model(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
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def sanitize(self, weights):
+        remove_patterns = [
+            "mtp_block.",
+            "mtp_linear_proj.",
+            "mtp_hidden_norm.",
+            "mtp_emb_norm.",
+            "e_score_correction_bias",
+        ]
+
+        weights = {
+            key: value
+            for key, value in weights.items()
+            if not any(pattern in key for pattern in remove_patterns)
+        }
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for m in ["gate_proj", "down_proj", "up_proj"]:
+                if f"{prefix}.mlp.experts.0.{m}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.mlp.experts.{e}.{m}.weight")
+                        for e in range(self.args.moe_num_experts)
+                    ]
+                    weights[f"{prefix}.mlp.switch_mlp.{m}.weight"] = mx.stack(to_join)
+
+        return weights

mlx_lm/models/hunyuan.py

@@ -29,6 +29,7 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float
     use_cla: bool
     cla_share_factor: 2
+    moe_intermediate_size: Optional[Union[int, list]] = None
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = False
 
@@ -40,6 +41,12 @@ class ModelArgs(BaseModelArgs):
                 raise ValueError(f"rope_scaling must contain keys {required_keys}")
 
 
+def _int_or_list(arg, idx):
+    if isinstance(arg, list):
+        return arg[idx]
+    return arg
+
+
 class DynamicNTKAlphaRoPE(nn.Module):
     def __init__(
         self,
@@ -154,20 +161,29 @@ class Gate(nn.Module):
 
 
 class MoeBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, layer_idx: int = 0):
         super().__init__()
         dim = args.hidden_size
         intermediate_size = args.intermediate_size
         self.use_shared_mlp = args.use_mixed_mlp_moe
 
         if args.use_mixed_mlp_moe:
-            self.shared_mlp = MLP(dim, intermediate_size * args.num_shared_expert)
+            num_shared = _int_or_list(args.num_shared_expert, layer_idx)
+            self.shared_mlp = MLP(dim, int(intermediate_size * num_shared))
 
         self.num_experts = num_experts = args.num_experts
-        self.top_k = args.moe_topk
+        self.top_k = _int_or_list(args.moe_topk, layer_idx)
 
         self.gate = Gate(dim, num_experts)
-        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
+
+        # Use moe_intermediate_size if available, otherwise use intermediate_size
+        expert_intermediate_size = intermediate_size
+        if args.moe_intermediate_size is not None:
+            expert_intermediate_size = _int_or_list(
+                args.moe_intermediate_size, layer_idx
+            )
+
+        self.switch_mlp = SwitchGLU(dim, expert_intermediate_size, num_experts)
 
     def __call__(
         self,
@@ -181,7 +197,7 @@ class MoeBlock(nn.Module):
         scores = mx.take_along_axis(gates, inds, axis=-1)
 
         y = self.switch_mlp(x, inds)
-        y = (y * scores[..., None]).sum(axis=-2)
+        y = (y * scores[..., None].astype(mx.float32)).sum(axis=-2).astype(y.dtype)
 
         if self.use_shared_mlp:
             shared_expert_output = self.shared_mlp(x)
@@ -191,14 +207,14 @@ class MoeBlock(nn.Module):
 
 
 class DecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs, kv_proj: bool):
+    def __init__(self, args: ModelArgs, kv_proj: bool, layer_idx: int = 0):
         super().__init__()
         self.hidden_size = args.hidden_size
         self.self_attn = Attention(kv_proj, args)
         if args.num_experts == 1:
             self.mlp = MLP(args.hidden_size, args.intermediate_size)
         else:
-            self.mlp = MoeBlock(args)
+            self.mlp = MoeBlock(args, layer_idx)
 
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
         self.post_attention_layernorm = nn.RMSNorm(
@@ -234,6 +250,7 @@ class HunYuanModel(nn.Module):
             DecoderLayer(
                 args=args,
                 kv_proj=(not args.use_cla) or (i % args.cla_share_factor) == 0,
+                layer_idx=i,
             )
             for i in range(args.num_hidden_layers)
         ]

mlx_lm/models/smollm3.py

@@ -0,0 +1,76 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from . import llama
+
+
+@dataclass
+class ModelArgs(llama.ModelArgs):
+    model_type: str
+    no_rope_layer_interval: int = 4
+    no_rope_layers: Optional[list[int]] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.no_rope_layers is None:
+            self.no_rope_layers = [
+                int((i + 1) % self.no_rope_layer_interval != 0)
+                for i in range(self.num_hidden_layers)
+            ]
+        elif len(self.no_rope_layers) != self.num_hidden_layers:
+            raise ValueError("`no_rope_layers` length mismatch")
+
+
+class NoPE(nn.Module):
+    """No-op used to disable rotary embeddings in selected layers."""
+
+    def __call__(self, x, offset: int = 0):
+        return x
+
+
+class Model(nn.Module):
+    """Wrapper around Llama that respects NoPE layers in SmolLM-3."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type: str = args.model_type
+
+        self.model = llama.LlamaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+        for idx, use_rope in enumerate(args.no_rope_layers):
+            if not use_rope:
+                self.model.layers[idx].self_attn.rope = NoPE()
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, mask, cache, input_embeddings)
+        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
+
+    def sanitize(self, weights: dict):
+        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

mlx_lm/models/switch_layers.py

@@ -53,12 +53,6 @@ class QuantizedSwitchLinear(nn.Module):
         # Freeze this model's parameters
         self.freeze()
 
-    def unfreeze(self, *args, **kwargs):
-        """Wrap unfreeze so that we unfreeze any layers we might contain but
-        our parameters will remain frozen."""
-        super().unfreeze(*args, **kwargs)
-        self.freeze(recurse=False)
-
     @property
     def input_dims(self):
         return self.scales.shape[2] * self.group_size
@@ -163,7 +157,8 @@ class SwitchGLU(nn.Module):
         inv_order = None
         if do_sort:
             x, idx, inv_order = _gather_sort(x, indices)
-
+        if self.training:
+            idx = mx.stop_gradient(idx)
         x_up = self.up_proj(x, idx, sorted_indices=do_sort)
         x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
         x = self.down_proj(
@@ -203,7 +198,8 @@ class SwitchMLP(nn.Module):
         inv_order = None
         if do_sort:
             x, idx, inv_order = _gather_sort(x, indices)
-
+        if self.training:
+            idx = mx.stop_gradient(idx)
         x = self.fc1(x, idx, sorted_indices=do_sort)
         x = self.activation(x)
         x = self.fc2(x, idx, sorted_indices=do_sort)

mlx_lm/quant/dwq.py

@@ -13,7 +13,8 @@ from mlx.utils import tree_flatten, tree_map
 from tqdm import tqdm
 
 from mlx_lm.tuner.datasets import load_dataset
-from mlx_lm.tuner.trainer import iterate_batches
+from mlx_lm.tuner.losses import kl_div_loss
+from mlx_lm.tuner.trainer import grad_checkpoint, iterate_batches
 from mlx_lm.tuner.utils import print_trainable_parameters
 from mlx_lm.utils import (
     fetch_from_hub,
@@ -43,6 +44,7 @@ def dwq_quantize(
     activation_layer_step: float = 0.25,
     activation_loss_weight: float = 1.0,
     dtype: mx.Dtype = mx.bfloat16,
+    gradient_checkpoint: bool = False,
 ):
     group = mx.distributed.init()
     world_size = group.size()
@@ -62,22 +64,22 @@ def dwq_quantize(
         model.layers[lid] = Catcher(model.layers[lid])
         q_model.layers[lid] = Catcher(q_model.layers[lid])
 
-    def log_norm(x):
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
+    if gradient_checkpoint:
+        grad_checkpoint(q_model.layers[0])
 
     def forward(model, inputs):
-        logprobs = log_norm(model(inputs).astype(mx.float32))
+        logits = model(inputs)
         extra_targets = [
             model.layers[lid].outputs.astype(mx.float32) for lid in layer_ids
         ]
         for lid in layer_ids:
             model.layers[lid].outputs = None
-        return logprobs, extra_targets
+        return logits, extra_targets
 
     def loss_fn(params, x, targets, extra_targets, lengths):
         q_model.update(tree_map(lambda x: x.astype(dtype), params))
-        logprobs, q_extra_targets = forward(q_model, x)
-        losses = nn.losses.kl_div_loss(logprobs, targets, reduction="none")
+        logits, q_extra_targets = forward(q_model, x)
+        losses = kl_div_loss(logits, targets)
         mask = mx.arange(1, 1 + targets.shape[1]) < lengths[:, 1:]
         ntoks = mask.sum()
         kl_loss = (mask * losses).sum() / ntoks
@@ -194,6 +196,11 @@ def main():
         default="allenai/tulu-3-sft-mixture",
         help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
     )
+    parser.add_argument(
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
+    )
     args = parser.parse_args()
 
     group = mx.distributed.init()
@@ -232,6 +239,7 @@ def main():
         calibration_data,
         batch_size=args.batch_size,
         max_seq_length=args.max_seq_length,
+        gradient_checkpoint=args.grad_checkpoint,
     )
     save(
         args.mlx_path,

mlx_lm/quant/dynamic_quant.py

@@ -12,6 +12,8 @@ from mlx.utils import tree_flatten, tree_map, tree_unflatten
 from tqdm import tqdm
 
 from mlx_lm.quant.utils import load_data
+from mlx_lm.tuner.losses import kl_div_loss
+from mlx_lm.tuner.trainer import grad_checkpoint
 from mlx_lm.utils import (
     compute_bits_per_weight,
     fetch_from_hub,
@@ -42,17 +44,17 @@ def estimate_sensitivities(
     low_group_size,
     high_bits,
     high_group_size,
+    batch_size: int = 4,
+    gradient_accum_dtype: mx.Dtype = mx.float32,
+    gradient_checkpoint: bool = False,
 ):
-    batch_size = 4
-    layers = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module)
-    layers = {k: l for k, l in layers if hasattr(l, "to_quantized")}
-
-    q_model = copy.deepcopy(model)
-
     def qdq(w, bits, group_size):
         w, s, b = mx.quantize(w, bits=bits, group_size=group_size)
         return mx.dequantize(w, scales=s, biases=b, bits=bits, group_size=group_size)
 
+    layers = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module)
+    layers = {k: l for k, l in layers if hasattr(l, "to_quantized")}
+    q_model = copy.deepcopy(model)
     q_layers = copy.deepcopy(layers)
     for l in q_layers.values():
         l.weight = qdq(l.weight, low_bits, low_group_size)
@@ -62,25 +64,27 @@ def estimate_sensitivities(
     q_model.freeze()
     q_model.update_modules(tree_unflatten(list(q_layers.items())))
 
-    def log_norm(x):
-        x = x.astype(mx.float32)
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
-
     def loss_fn(batch, targets):
-        logprobs = log_norm(q_model(batch))
-        return nn.losses.kl_div_loss(logprobs, targets, reduction="mean")
+        return kl_div_loss(q_model(batch), targets).mean()
 
-    grad_accum = tree_map(lambda x: mx.zeros(x.shape), q_model.trainable_parameters())
+    if gradient_checkpoint:
+        grad_checkpoint(q_model.layers[0])
+
+    grad_accum = tree_map(
+        lambda x: mx.zeros(x.shape, dtype=gradient_accum_dtype),
+        q_model.trainable_parameters(),
+    )
     for e, s in tqdm(
         enumerate(range(0, len(data), batch_size)),
         total=len(data) // batch_size,
         desc="Estimating sensitivities",
     ):
         batch = data[s : s + batch_size]
-        targets = log_norm(model(batch))
+        targets = model(batch)
         mx.eval(targets)
         _, grads = nn.value_and_grad(q_model, loss_fn)(batch, targets)
         grad_accum = tree_map(lambda x, y: x + y, grad_accum, grads)
+        del grads
         mx.eval(grad_accum)
 
     def compute_sensitivity(gradient, low_q_weight, original_weight):
@@ -169,7 +173,17 @@ def main():
         action="store_true",
         help="Compute the perplexity of the base and quantized models.",
     )
-
+    parser.add_argument(
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
+    )
+    parser.add_argument(
+        "--accumulation-dtype",
+        default="float32",
+        choices=["float32", "bfloat16"],
+        help="What type to use to accumulate the gradients for the sensitivities",
+    )
     args = parser.parse_args()
 
     group = mx.distributed.init()
@@ -186,6 +200,8 @@ def main():
             args.low_group_size,
             args.high_bits,
             args.high_group_size,
+            gradient_accum_dtype=getattr(mx, args.accumulation_dtype),
+            gradient_checkpoint=args.grad_checkpoint,
         )
         model_name = args.model.replace("/", "_")
         with open(f"{model_name}_sensitivities.json", "w") as fid:
@@ -241,6 +257,7 @@ def main():
         config,
         hf_repo=hf_repo,
     )
+    print(f"Peak memory used: {mx.get_peak_memory() / 1000**3:.3f}GB")
 
 
 if __name__ == "__main__":

mlx_lm/quant/gptq.py

@@ -0,0 +1,232 @@
+# Copyright © 2025 Apple Inc.
+
+"""
+Implements GPTQ
+
+- https://arxiv.org/abs/2210.17323
+- https://github.com/AutoGPTQ
+"""
+
+import argparse
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+from tqdm import tqdm
+
+from mlx_lm.models.switch_layers import QuantizedSwitchLinear, SwitchLinear
+from mlx_lm.quant.utils import load_data
+from mlx_lm.utils import (
+    compute_bits_per_weight,
+    fetch_from_hub,
+    get_model_path,
+    save,
+)
+
+
+def quantize(w, bits, scales, biases):
+    assert bits in {2, 4, 8}, f"Unsupported bits {bits}"
+    el_per_int = 32 // bits
+    n_bins = 2**bits - 1
+    w = mx.unflatten(w, -1, (scales.shape[-1], -1))
+    w = mx.clip(
+        mx.round((w - biases[..., None]) / scales[..., None]), 0.0, n_bins
+    ).astype(mx.uint32)
+    shifts = mx.power(2, mx.arange(0, 32, bits, mx.uint32))
+    w = mx.unflatten(w, -1, (-1, el_per_int))
+    w = mx.sum(w * shifts, axis=-1)
+    return w.flatten(-2, -1)
+
+
+class Catcher(nn.Module):
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+        self.H = mx.array(0.0)
+
+    def __call__(self, x, *args, **kwargs):
+        xf = x.flatten(0, -2)
+        self.H = self.H + xf.T @ xf
+        return self.module(x, *args, **kwargs)
+
+
+def gptq_quantize(
+    model,
+    data,
+    bits,
+    group_size,
+    fallback_bits,
+    fallback_group_size,
+    batch_size=8,
+):
+    layers = []
+    gptq_types = {nn.Linear, SwitchLinear}
+    for k, l in tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module):
+        if type(l) in gptq_types:
+            layers.append((k, Catcher(l)))
+    model.update_modules(tree_unflatten(layers))
+
+    # Evaluate the Hessians for all quantizable layers
+    for e, s in tqdm(
+        enumerate(range(0, len(data), batch_size)),
+        total=len(data) // batch_size,
+        desc="Computing Hessians",
+    ):
+        batch = data[s : s + batch_size]
+        model(batch)
+        mx.eval(layers)
+
+    def compute_inverse_hessian(H):
+        with mx.stream(mx.cpu):
+            damp = 1e-2 * mx.mean(mx.diag(H))
+            diag = mx.arange(H.shape[0])
+            H[diag, diag] += damp
+            H = mx.linalg.cholesky(H)
+            H = mx.linalg.cholesky_inv(H)
+            Hinv = mx.linalg.cholesky(H, upper=True)
+            return Hinv
+
+    @mx.compile
+    def gptq_error(w, d, scales, biases):
+        n_bins = 2**bits - 1
+        q = mx.clip(mx.round((w - biases) / scales), 0.0, n_bins)
+        q = scales * q + biases
+        return (w - q) / d
+
+    for lid, (key, l) in tqdm(
+        enumerate(layers),
+        total=len(layers),
+        desc="Quantizing",
+    ):
+        Hinv = compute_inverse_hessian(l.H)
+        del l.H
+        mx.eval(Hinv)
+
+        orig_type = l.module.weight.dtype
+        W = l.module.weight.astype(mx.float32)
+
+        all_scales = []
+        all_biases = []
+        for i in range(0, W.shape[-1], group_size):
+            j = i + group_size
+            Wl = W[..., i:j]
+            err = mx.zeros_like(Wl)
+
+            # Find scales and biases
+            _, scales, biases = mx.quantize(Wl, bits=bits, group_size=group_size)
+
+            all_scales.append(scales)
+            all_biases.append(biases)
+            for k in range(group_size):
+                k += i
+                w = W[..., k : k + 1]
+                d = Hinv[k, k]
+
+                e = gptq_error(w, d, scales, biases)
+
+                W[..., k : k + j] -= e @ Hinv[k : k + 1, k : k + j]
+                err[..., k : k + 1] = e
+                mx.eval(err, W)
+
+            W[..., j:] -= err @ Hinv[i:j, j:]
+
+        # Quantize with the given scales and biases
+        scales = mx.concatenate(all_scales, axis=-1)
+        biases = mx.concatenate(all_biases, axis=-1)
+        Wq = quantize(W, bits, scales, biases)
+        layer = l.module.to_quantized(bits=bits, group_size=group_size)
+        layer.weight = Wq
+        layer.scales = scales
+        layer.biases = biases
+        layer.set_dtype(orig_type)
+        mx.eval(layer)
+        layers[lid] = (key, layer)
+
+    model.update_modules(tree_unflatten(layers))
+
+    layers = tree_flatten(
+        model.leaf_modules(),
+        is_leaf=nn.Module.is_module,
+    )
+    config = {"bits": bits, "group_size": group_size}
+    fallback_config = {"bits": fallback_bits, "group_size": fallback_group_size}
+    q_layers = []
+    for e, (k, l) in enumerate(layers):
+        if hasattr(l, "to_quantized"):
+            config[k] = fallback_config
+            q_layers.append((k, l.to_quantized(**fallback_config)))
+    if len(q_layers) > 0:
+        model.update_modules(tree_unflatten(q_layers))
+    return model, config
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", "-m", default="Qwen/Qwen3-0.6B-base")
+    parser.add_argument("--mlx-path", default="mlx_model")
+    parser.add_argument(
+        "--bits", type=int, default=4, help="Quantization bits for GPTQ layers"
+    )
+    parser.add_argument(
+        "--group-size",
+        type=int,
+        default=64,
+        help="Quantization group size for GPTQ layers",
+    )
+    parser.add_argument(
+        "--fallback-bits",
+        type=int,
+        default=6,
+        help="Quantization bits for non-GPTQ layers",
+    )
+    parser.add_argument(
+        "--fallback-group-size",
+        type=int,
+        default=64,
+        help="Quantization group size for non-GPTQ layers",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=-1,
+        help="Number of samples from the calibration dataset, use -1 for all.",
+    )
+    parser.add_argument(
+        "--sequence-length",
+        type=int,
+        default=512,
+        help="Sequence length for the calibration data.",
+    )
+    parser.add_argument("--seed", type=int, default=123)
+    args = parser.parse_args()
+
+    mx.random.seed(args.seed)
+
+    model_path, hf_repo = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    calibration_data = load_data(tokenizer, args.num_samples, args.sequence_length)
+
+    model, config["quantization"] = gptq_quantize(
+        model,
+        calibration_data,
+        args.bits,
+        args.group_size,
+        args.fallback_bits,
+        args.fallback_group_size,
+    )
+
+    bpw = compute_bits_per_weight(model)
+    print(f"Quantized model with {bpw:.3f} bits per weight.")
+
+    save(
+        args.mlx_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_repo,
+    )
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/tuner/losses.py

@@ -0,0 +1,378 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def _make_kl_forward_kernel():
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx;
+
+    float lse_q_minus_p;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q_minus_p = max_q + metal::fast::log(sum_exp_q) - lse_p;
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                kl += metal::fast::exp(vals_p[j] - lse_p) * (vals_p[j] - vals_q[j] + lse_q_minus_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                kl += metal::fast::exp(vals_p[j] - lse_p) * (vals_p[j] - vals_q[j] + lse_q_minus_p);
+            }
+        }
+
+        // Add the kl across the threadgroup
+        kl = simd_sum(kl);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id] = kl;
+        }
+        threadgroup_barrier(mem_flags::mem_none);
+        kl = shared[simd_lane_id];
+        kl = simd_sum(kl);
+
+        if (thread_index_in_threadgroup == 0) {
+            out[0] = static_cast<T>(kl);
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="kl_forward",
+        input_names=["logits_q", "logits_p"],
+        output_names=["out"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+def _make_kl_backward_kernel():
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx * V;
+    cotan += out_idx;
+
+    float lse_q;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q = max_q + metal::fast::log(sum_exp_q);
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl = 0;
+        float c = cotan[0];
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                out[offset + j] = static_cast<T>(
+                    c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p)));
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                if (offset + j < V) {
+                    out[offset + j] = static_cast<T>(
+                        c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p)));
+                }
+            }
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="kl_backward",
+        input_names=["logits_q", "logits_p", "cotan"],
+        output_names=["out"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+_kl_forward_kernel = _make_kl_forward_kernel()
+_kl_backward_kernel = _make_kl_backward_kernel()
+
+
+@mx.custom_function
+def _kl_div_loss(logits_q, logits_p):
+    n_outs = logits_q.size // logits_q.shape[-1]
+    dt = logits_q.dtype
+
+    return _kl_forward_kernel(
+        inputs=[logits_q, logits_p],
+        output_shapes=[logits_q.shape[:-1]],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, n_outs, 1),
+        threadgroup=(1024, 1, 1),
+    )[0]
+
+
+@_kl_div_loss.vjp
+def _kl_div_loss(primals, cotangent, output):
+    logits_q, logits_p = primals
+    dt = logits_q.dtype
+
+    dp = mx.zeros_like(logits_p)
+    dq = _kl_backward_kernel(
+        inputs=[logits_q, logits_p, cotangent],
+        output_shapes=[logits_q.shape],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, cotangent.size, 1),
+        threadgroup=(1024, 1, 1),
+    )[0]
+
+    return dq, dp
+
+
+def kl_div_loss(logits_q, logits_p):
+    if mx.metal.is_available():
+        return _kl_div_loss(logits_q, logits_p)
+    else:
+        return nn.losses.kl_div_loss(
+            logits_q - mx.logsumexp(logits_q, axis=-1, keepdims=True),
+            logits_p - mx.logsumexp(logits_p, axis=-1, keepdims=True),
+            axis=-1,
+            reduction="none",
+        )

mlx_lm/tuner/utils.py

@@ -84,7 +84,7 @@ def linear_to_lora_layers(
     keys = config.get("keys", None)
     if keys is not None:
         keys = set(keys)
-    elif model.model_type in [
+    elif model.model_type in {
         "mistral",
         "mistral3",
         "llama",
@@ -118,8 +118,9 @@ def linear_to_lora_layers(
         "mimo",
         "ernie4_5",
         "dots1",
-    ]:
-        keys = set(["self_attn.q_proj", "self_attn.v_proj"])
+        "smollm3",
+    }:
+        keys = {"self_attn.q_proj", "self_attn.v_proj"}
         if model.model_type in ["mixtral", "phimoe"]:
             keys.add("block_sparse_moe.gate")
         if model.model_type == "qwen2_moe":
@@ -129,44 +130,40 @@ def linear_to_lora_layers(
             keys.add("mlp.gate")
 
     elif model.model_type == "gpt_bigcode":
-        keys = set(["attn.c_attn"])
+        keys = {"attn.c_attn"}
     elif model.model_type == "gpt2":
-        keys = set(["attn.c_attn"])
+        keys = {"attn.c_attn"}
     elif model.model_type == "gpt_neox":
-        keys = set(["attention.query_key_value"])
+        keys = {"attention.query_key_value"}
     elif model.model_type == "olmo":
-        keys = set(["att_proj"])
+        keys = {"att_proj"}
     elif model.model_type == "openelm":
-        keys = set(["attn.qkv_proj"])
+        keys = {"attn.qkv_proj"}
     elif model.model_type == "phi3":
-        keys = set(["self_attn.qkv_proj"])
+        keys = {"self_attn.qkv_proj"}
     elif model.model_type == "phi-msft":
-        keys = set(["mixer.Wqkv", "moe.gate"])
+        keys = {"mixer.Wqkv", "moe.gate"}
     elif model.model_type == "dbrx":
-        keys = set(["norm_attn_norm.attn.Wqkv", "ffn.router.layer"])
+        keys = {"norm_attn_norm.attn.Wqkv", "ffn.router.layer"}
     elif model.model_type == "internlm2":
-        keys = set(["attention.wqkv", "attention.wo"])
-    elif model.model_type == "deepseek_v2" or model.model_type == "minicpm3":
-        keys = set(
-            [
-                "self_attn.q_proj",
-                "self_attn.q_a_proj",
-                "self_attn.q_b_proj",
-                "self_attn.kv_a_proj_with_mqa",
-                "self_attn.kv_b_proj",
-            ]
-        )
+        keys = {"attention.wqkv", "attention.wo"}
+    elif model.model_type in {"deepseek_v2", "deepseek_v3", "minicpm3"}:
+        keys = {
+            "self_attn.q_proj",
+            "self_attn.q_a_proj",
+            "self_attn.q_b_proj",
+            "self_attn.kv_a_proj_with_mqa",
+            "self_attn.kv_b_proj",
+        }
     elif model.model_type == "mamba":
-        keys = set(
-            [
-                "mixer.in_proj",
-                "mixer.x_proj",
-                "mixer.dt_proj",
-                "mixer.out_proj",
-            ]
-        )
+        keys = {
+            "mixer.in_proj",
+            "mixer.x_proj",
+            "mixer.dt_proj",
+            "mixer.out_proj",
+        }
     elif model.model_type == "exaone":
-        keys = set(["attn.attention.q_proj", "attn.attention.v_proj"])
+        keys = {"attn.attention.q_proj", "attn.attention.v_proj"}
     else:
         raise ValueError(f"Lora does not support {model.model_type}")
 

mlx_lm/utils.py

@@ -113,11 +113,14 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
             )
         )
     else:
-
         from huggingface_hub import ModelCard
 
-        card = ModelCard.load(model_path / "README.md")
-        hf_path = card.data.base_model
+        card_path = model_path / "README.md"
+        if card_path.is_file():
+            card = ModelCard.load(card_path)
+            hf_path = card.data.base_model
+        else:
+            hf_path = None
     return model_path, hf_path
 
 
@@ -194,7 +197,6 @@ def load_model(
                 return config["quantization"][p]
             if not hasattr(m, "to_quantized"):
                 return False
-            # Handle legacy models which may not have everything quantized
             return f"{p}.scales" in weights
 
         nn.quantize(
@@ -203,6 +205,15 @@ def load_model(
             bits=quantization["bits"],
             class_predicate=class_predicate,
         )
+    elif quantization_config := config.get("quantization_config", False):
+        # Handle legacy quantization config
+        quant_method = quantization_config["quant_method"]
+        if quant_method == "bitnet":
+            from .models.bitlinear_layers import bitnet_quantize
+
+            model = bitnet_quantize(model, quantization_config)
+        else:
+            raise ValueError(f"Unsupported quantization method {quant_method}")
 
     model.load_weights(list(weights.items()), strict=strict)
 
@@ -255,11 +266,13 @@ def load(
 
 
 def fetch_from_hub(
-    model_path: Path, lazy: bool = False
+    model_path: Path, lazy: bool = False, trust_remote_code: bool = False
 ) -> Tuple[nn.Module, dict, PreTrainedTokenizer]:
     model, config = load_model(model_path, lazy)
     tokenizer = load_tokenizer(
-        model_path, eos_token_ids=config.get("eos_token_id", None)
+        model_path,
+        eos_token_ids=config.get("eos_token_id", None),
+        tokenizer_config_extra={"trust_remote_code": trust_remote_code},
     )
     return model, config, tokenizer
 
@@ -460,9 +473,18 @@ def quantize_model(
     quantized_config = copy.deepcopy(config)
     quantized_config["quantization"] = {"group_size": q_group_size, "bits": q_bits}
 
+    def base_predicate(path, module):
+        if not hasattr(module, "to_quantized"):
+            return False
+        if module.weight.shape[-1] % q_group_size != 0:
+            return False
+        return True
+
     # Add any custom quantization parameters to the config as we go
-    def _class_predicate(p, m):
-        bool_or_params = quant_predicate(p, m, config)
+    def wrapped_predicate(p, m):
+        bool_or_params = base_predicate(p, m)
+        if bool_or_params:
+            bool_or_params = quant_predicate(p, m, config)
         quantized_config["quantization"][p] = bool_or_params
         return bool_or_params
 
@@ -470,7 +492,7 @@ def quantize_model(
         model,
         q_group_size,
         q_bits,
-        class_predicate=_class_predicate if quant_predicate else None,
+        class_predicate=wrapped_predicate if quant_predicate else base_predicate,
     )
     # support hf model tree #957
     quantized_config["quantization_config"] = quantized_config["quantization"]
@@ -496,6 +518,8 @@ def save_config(
     # Clean unused keys
     config.pop("_name_or_path", None)
     config.pop("vision_config", None)
+    if "quantization" in config:
+        config["quantization_config"] = config["quantization"]
 
     # sort the config for better readability
     config = dict(sorted(config.items()))

requirements.txt

@@ -1,6 +1,6 @@
 mlx>=0.25.0
 numpy
-transformers[sentencepiece]>=4.39.3
+transformers>=4.39.3
 protobuf
 pyyaml
 jinja2

setup.py

@@ -36,6 +36,7 @@ setup(
             "mlx_lm.awq = mlx_lm.quant.awq:main",
             "mlx_lm.dwq = mlx_lm.quant.dwq:main",
             "mlx_lm.dynamic_quant = mlx_lm.quant.dynamic_quant:main",
+            "mlx_lm.gptq = mlx_lm.quant.gptq:main",
             "mlx_lm.cache_prompt = mlx_lm.cache_prompt:main",
             "mlx_lm.chat = mlx_lm.chat:main",
             "mlx_lm.convert = mlx_lm.convert:main",

tests/test_generate.py

@@ -104,7 +104,7 @@ class TestGenerate(unittest.TestCase):
         for generation_result in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
-            prompt=[],  # no prompt tokens passed
+            prompt=prompt,
             max_tokens=5,
             sampler=sampler,
             input_embeddings=prompt_embeddings,
@@ -136,7 +136,7 @@ class TestGenerate(unittest.TestCase):
         for generation_result in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
-            prompt=[],  # no prompt tokens passed
+            prompt=prompt,
             max_tokens=5,
             sampler=sampler,
             input_embeddings=prompt_embeddings,

tests/test_models.py

@@ -1082,6 +1082,23 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_smollm3(self):
+        from mlx_lm.models import smollm3
+
+        args = smollm3.ModelArgs(
+            model_type="smollm3",
+            hidden_size=1024,
+            num_hidden_layers=4,
+            intermediate_size=2048,
+            num_attention_heads=4,
+            rms_norm_eps=1e-5,
+            vocab_size=10_000,
+        )
+        model = smollm3.Model(args)
+        self.model_test_runner(
+            model, "smollm3", args.vocab_size, args.num_hidden_layers
+        )
+
 
 if __name__ == "__main__":
     unittest.main()