fix test

* refactor out of utils

* remove comment

* format

* move generate as well

ACKNOWLEDGMENTS.md

@@ -9,4 +9,4 @@ MLX LM was developed with contributions from the following individuals:
 
 - Shunta Saito: Added support for PLaMo models.
 - Prince Canuma: Helped add support for `Starcoder2` models.
-- - Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, and Allenai's `OLMoE`; added support for the following training algorithms: `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`.

MANIFEST.in

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

README.md

@@ -220,7 +220,7 @@ 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-lm/blob/main/mlx_lm/examples/chat.py)

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.22.1"
+__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
 

mlx_lm/chat.py

@@ -5,9 +5,10 @@ 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

mlx_lm/convert.py

@@ -1,23 +1,137 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+import glob
+import shutil
+from pathlib import Path
+from typing import Callable, Optional, Union
 
-from . import utils
-from .utils import convert
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
 
-QUANT_RECIPES = [
-    "mixed_2_6",
-    "mixed_3_6",
-]
+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: {QUANT_RECIPES}"
+            f"Invalid q-recipe {arg!r}. Choose from: {list(QUANT_RECIPES.keys())}"
         )
     else:
-        return getattr(utils, arg)
+        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:
@@ -46,7 +160,7 @@ def configure_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument(
         "--quant-predicate",
-        help=f"Mixed-bit quantization recipe. Choices: {QUANT_RECIPES}",
+        help=f"Mixed-bit quantization recipe. Choices: {list(QUANT_RECIPES.keys())}",
         type=quant_args,
         required=False,
     )
@@ -80,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)

mlx_lm/examples/lora_config.yaml

@@ -81,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/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
@@ -162,6 +185,526 @@ def setup_arg_parser():
     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()
@@ -284,4 +827,8 @@ def main():
 
 
 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
@@ -167,6 +168,12 @@ 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",
@@ -191,6 +198,7 @@ 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(
@@ -236,10 +244,9 @@ 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()
-
     # Initialize the selected optimizer
     lr = build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
 
@@ -268,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/models/base.py

@@ -42,19 +42,24 @@ def create_causal_mask(
     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)
+                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,7 +83,7 @@ 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:]
@@ -170,10 +170,22 @@ class CohereModel(nn.Module):
 
         if mask is None:
             j = self.args.sliding_window_pattern
-            mask = create_attention_mask(h, cache[j - 1 : j])
+            full_mask = create_attention_mask(h, cache[j - 1 : j])
+            sliding_window_mask = create_attention_mask(h, cache)
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
+        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/gemma2.py

@@ -172,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

@@ -1,15 +1,14 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional
+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, create_attention_mask
-from .cache import KVCache, RotatingKVCache
+from .base import BaseModelArgs
 
 
 @dataclass

mlx_lm/models/gemma3_text.py

@@ -86,7 +86,7 @@ class Attention(nn.Module):
             keys = self.rope(keys)
 
         # Sliding window
-        if mask is not None and mask.shape[-1] != keys.shape[-2]:
+        if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
             mask = mask[..., -keys.shape[-2] :]
 
         output = mx.fast.scaled_dot_product_attention(
@@ -179,17 +179,18 @@ class Gemma3Model(nn.Module):
             sliding_window_mask = create_attention_mask(h, cache)
 
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            is_sliding = (
+            is_global = (
                 i % self.args.sliding_window_pattern
                 == self.args.sliding_window_pattern - 1
             )
 
-            if mask is None and is_sliding:
-                mask = sliding_window_mask
+            local_mask = mask
+            if mask is None and is_global:
+                local_mask = full_mask
             elif mask is None:
-                mask = full_mask
+                local_mask = sliding_window_mask
 
-            h = layer(h, mask, c)
+            h = layer(h, local_mask, c)
 
         return self.norm(h)
 

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/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 .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():
@@ -710,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(
@@ -764,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,
@@ -789,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/tuner/datasets.py

@@ -1,4 +1,3 @@
-import itertools
 import json
 import types
 from pathlib import Path
@@ -7,7 +6,7 @@ from typing import Any, Dict, List, Optional
 from transformers import PreTrainedTokenizer
 
 
-class Dataset:
+class TextDataset:
     """
     Light-weight wrapper to hold a dataset.
     """
@@ -18,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]
@@ -43,17 +47,24 @@ class ChatDataset:
         chat_key: str = "messages",
         mask_prompt: bool = False,
     ):
-        self._data = []
-        for d in data:
-            messages = d[chat_key]
-            tools = d.get("tools", None)
-            tokens = tokenizer.apply_chat_template(messages, tools=tools)
-            if mask_prompt:
-                messages = messages[:-1]
-                offset = len(tokenizer.apply_chat_template(messages, tools=tools))
-                self._data.append((tokens, offset))
-            else:
-                self._data.append(tokens)
+        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]
@@ -77,23 +88,28 @@ class CompletionsDataset:
         completion_key: str,
         mask_prompt: bool,
     ):
-        self._data = []
-        for d in data:
-            tokens = tokenizer.apply_chat_template(
-                [
-                    {"role": "user", "content": d[prompt_key]},
-                    {"role": "assistant", "content": d[completion_key]},
-                ],
-            )
-            if mask_prompt:
-                offset = len(
-                    tokenizer.apply_chat_template(
-                        [{"role": "user", "content": d[prompt_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]}]
                 )
-                self._data.append((tokens, offset))
-            else:
-                self._data.append(tokens)
+            )
+            return (tokens, offset)
+
+        return tokens
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -104,10 +120,33 @@ class CompletionsDataset:
 
 class ConcatenatedDataset:
     def __init__(self, data: List[Any]):
-        self._data = list(itertools.chain(*data))
+        self._data = data
+        self._len = sum(len(d) for d in self._data)
 
     def __getitem__(self, idx: int):
-        return self._data[idx]
+        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)
@@ -135,7 +174,7 @@ def create_dataset(
     elif text_feature in sample:
         if mask_prompt:
             raise ValueError("Prompt masking not supported for text dataset.")
-        return Dataset(data, tokenizer, text_key=text_feature)
+        return TextDataset(data, tokenizer, text_key=text_feature)
     else:
         raise ValueError(
             "Unsupported data format, check the supported formats here:\n"
@@ -204,8 +243,8 @@ def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
 
     collection = []
     for ds in dataset_collection:
-        ds_name = ds["name"]
-        print(f"Loading Hugging Face dataset {ds_name}.")
+        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", {})
@@ -213,13 +252,13 @@ def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
             train_split = ds.get("train_split", "train[:80%]")
             valid_split = ds.get("valid_split", "train[-10%:]")
             train = create_hf_dataset(
-                ds_name,
+                ds_path,
                 config,
                 train_split,
                 hf_config,
             )
             valid = create_hf_dataset(
-                ds_name,
+                ds_path,
                 config,
                 valid_split,
                 hf_config,
@@ -230,7 +269,7 @@ def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
         if args.test:
             test_split = ds.get("test_split")
             test = create_hf_dataset(
-                ds_name,
+                ds_path,
                 config,
                 test_split,
                 hf_config,

mlx_lm/tuner/trainer.py

@@ -4,6 +4,7 @@ import glob
 import shutil
 import time
 from dataclasses import dataclass, field
+from functools import partial
 from pathlib import Path
 from typing import List, Optional, Tuple
 
@@ -11,10 +12,19 @@ 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 .datasets import CompletionsDataset
+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):
@@ -64,16 +74,21 @@ 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, batch, lengths):
+def default_loss(model, batch, lengths, cache=None):
     inputs = batch[:, :-1]
     targets = batch[:, 1:]
 
-    logits = model(inputs)
+    offset = cache[0].offset if cache is not None else 0
+    logits = model(inputs, cache=cache)
     logits = logits.astype(mx.float32)
 
-    steps = mx.arange(1, targets.shape[1] + 1)
+    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) * mask
@@ -91,7 +106,11 @@ def iterate_batches(
     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}"
@@ -126,9 +145,9 @@ def iterate_batches(
                     "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)
@@ -155,12 +174,17 @@ def evaluate(
     max_seq_length=2048,
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
+    seq_step_size: Optional[int] = None,
 ):
+    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(
@@ -170,10 +194,15 @@ 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, stream=mx.cpu)
     ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
@@ -203,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()
@@ -213,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)
@@ -227,6 +260,44 @@ 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
@@ -259,7 +330,9 @@ def train(
                 num_batches=args.val_batches,
                 max_seq_length=args.max_seq_length,
                 iterate_batches=iterate_batches,
+                seq_step_size=seq_step_size,
             )
+            model.train()
             val_time = time.perf_counter() - tic
             if rank == 0:
                 print(
@@ -279,23 +352,28 @@ def train(
 
             tic = time.perf_counter()
 
-        lvalue, toks = step(batch)
+        if batch[0].shape[1] > seq_step_size:
+            lvalue, toks = seq_split_step(batch)
+        else:
+            lvalue, toks = step(batch)
+
         losses += lvalue
         n_tokens += toks
         steps += 1
         mx.eval(state, losses, n_tokens)
+
         train_time += time.perf_counter() - tic
 
         # Report training loss if needed
         if it % args.steps_per_report == 0 or it == args.iters:
             train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
-            train_loss /= steps * mx.distributed.init().size()
+            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 / 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}, "
@@ -325,7 +403,7 @@ def train(
             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 = (
@@ -338,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/utils.py

@@ -1,25 +1,17 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import contextlib
 import copy
-import functools
 import glob
 import importlib
 import json
 import logging
 import os
-import shutil
-import time
-from dataclasses import dataclass
 from pathlib import Path
 from textwrap import dedent
 from typing import (
     Any,
     Callable,
     Dict,
-    Generator,
-    List,
-    NamedTuple,
     Optional,
     Tuple,
     Type,
@@ -43,8 +35,6 @@ 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
@@ -58,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):
@@ -68,70 +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.
-        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
-
-
-@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.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.
@@ -208,458 +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,
-) -> 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.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 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.metal.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.metal.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.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:
@@ -1020,105 +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 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
-
-
-mixed_3_6 = mixed_quant_predicate_builder(low_bits=3, high_bits=6)
-mixed_2_6 = mixed_quant_predicate_builder(low_bits=2, high_bits=6)
-
-
-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.22.0
+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))

tests/test_datsets.py

@@ -44,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."}
@@ -80,7 +80,7 @@ class TestDatasets(unittest.TestCase):
 
     def test_hf(self):
         hf_args = {
-            "name": "billsum",
+            "path": "billsum",
             "prompt_feature": "text",
             "completion_feature": "summary",
             "train_split": "train[:2%]",

tests/test_finetune.py

@@ -370,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 = [
@@ -412,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

@@ -3,14 +3,13 @@
 import unittest
 from typing import List
 
-from mlx_lm.sample_utils import make_logits_processors
-from mlx_lm.utils import (
+from mlx_lm.generate import (
     GenerationResponse,
     generate,
-    load,
-    make_sampler,
     stream_generate,
 )
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
+from mlx_lm.utils import load
 
 
 class TestGenerate(unittest.TestCase):

tests/test_prompt_cache.py

@@ -7,6 +7,7 @@ import unittest
 
 import mlx.core as mx
 
+from mlx_lm.generate import generate_step
 from mlx_lm.models.cache import (
     KVCache,
     MambaCache,
@@ -17,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"
 
@@ -299,7 +300,7 @@ class TestPromptCache(unittest.TestCase):
         ):
             i += 1
             self.assertEqual(tok, toks[i])
-            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=3e-2))
+            self.assertTrue(mx.allclose(logits, all_logits[i], rtol=4e-2))
 
 
 if __name__ == "__main__":

tests/test_utils.py

@@ -8,7 +8,7 @@ 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"
 
@@ -77,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)