log only on rank 0

Co-authored-by: Tristan Stahnke <tristan@melchior.lan>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.31.2"
+__version__ = "0.31.3"

mlx_lm/chat.py

@@ -1,9 +1,16 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+import sys
 
 import mlx.core as mx
 
+from .cli_ui import (
+    make_console,
+    make_corridor_prompt,
+    print_chat_help,
+    print_header_panel,
+)
 from .generate import stream_generate
 from .models.cache import make_prompt_cache
 from .sample_utils import make_sampler
@@ -18,6 +25,19 @@ DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 
 
+def _print_chat_header(args, console):
+    rows = [("model", str(args.model))]
+    if args.adapter_path:
+        rows.append(("adapter", str(args.adapter_path)))
+    rows.append(("max tokens", f"{args.max_tokens:,}"))
+    if args.system_prompt:
+        sp = args.system_prompt
+        if len(sp) > 60:
+            sp = sp[:57] + "..."
+        rows.append(("system", sp))
+    print_header_panel(console, "mlx_lm.chat", rows)
+
+
 def setup_arg_parser():
     """Set up and return the argument parser."""
     parser = argparse.ArgumentParser(description="Chat with an LLM")
@@ -96,6 +116,8 @@ def main():
     pipeline_group = group if args.pipeline else None
     tensor_group = group if not args.pipeline else None
 
+    console = make_console()
+
     def rprint(*args, **kwargs):
         if rank == 0:
             print(*args, **kwargs)
@@ -115,24 +137,38 @@ def main():
             },
         )
 
-    def print_help():
-        rprint("The command list:")
-        rprint("- 'q' to exit")
-        rprint("- 'r' to reset the chat")
-        rprint("- 'h' to display these commands")
+    if rank == 0:
+        _print_chat_header(args, console)
+        print_chat_help(console)
+
+    if rank == 0:
+        prompt_console = make_console(force_terminal=True, color_system="truecolor")
+        _draw_corridor_prompt = make_corridor_prompt(prompt_console)
+    else:
+        _draw_corridor_prompt = lambda: ""
 
-    rprint(f"[INFO] Starting chat session with {args.model}.")
-    print_help()
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
-        query = input(">> " if rank == 0 else "")
+        prompt = _draw_corridor_prompt()
+        query = input(prompt)
+        if rank == 0:
+            # Cursor is now on the bottom-rule row; advance past it.
+            sys.stdout.write("\n")
+            sys.stdout.flush()
         if query == "q":
+            if rank == 0:
+                console.print("[ui.muted]bye[/ui.muted]")
             break
         if query == "r":
             prompt_cache = make_prompt_cache(model, args.max_kv_size)
+            if rank == 0:
+                console.print(
+                    "  [ui.good]reset[/ui.good] [ui.muted]context cleared[/ui.muted]"
+                )
             continue
         if query == "h":
-            print_help()
+            if rank == 0:
+                print_chat_help(console)
             continue
         messages = []
         if args.system_prompt is not None:
@@ -142,6 +178,7 @@ def main():
             messages,
             add_generation_prompt=True,
         )
+        last_response = None
         for response in stream_generate(
             model,
             tokenizer,
@@ -159,7 +196,15 @@ def main():
             prompt_cache=prompt_cache,
         ):
             rprint(response.text, flush=True, end="")
+            last_response = response
         rprint()
+        if rank == 0 and last_response is not None:
+            console.print(
+                f"  [ui.muted]{last_response.generation_tokens} tokens · "
+                f"{last_response.generation_tps:.1f} tok/s · "
+                f"prompt {last_response.prompt_tps:.1f} tok/s · "
+                f"peak {last_response.peak_memory:.2f} GB[/ui.muted]"
+            )
 
 
 if __name__ == "__main__":

mlx_lm/cli_ui.py

@@ -0,0 +1,172 @@
+# Copyright © 2024 Apple Inc.
+
+"""Shared UI helpers for the mlx_lm command-line tools.
+
+Centralizes the rich-based panel/progress/prompt rendering used by the chat
+and training entry points. The theme is hardcoded for a light terminal
+background.
+"""
+
+import os
+import re
+import shutil
+import sys
+
+from rich.box import ROUNDED
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import (
+    Progress,
+    ProgressColumn,
+    TextColumn,
+)
+from rich.text import Text
+from rich.theme import Theme
+
+
+def _terminal_width(default: int = 120) -> int:
+    """Best-effort terminal width.
+
+    Under launchers like ``mlx.launch`` the worker's stdout is a pipe, so
+    Rich's auto-detection falls back to 80 columns. Honor an explicit
+    ``MLX_LM_WIDTH`` override, then ``COLUMNS``, then a real TTY query, and
+    finally a generous default that's nicer than 80 on modern terminals.
+    """
+    for var in ("MLX_LM_WIDTH", "COLUMNS"):
+        value = os.environ.get(var)
+        if value and value.isdigit():
+            return int(value)
+    width = shutil.get_terminal_size(fallback=(0, 0)).columns
+    return width if width > 0 else default
+
+
+def _make_theme() -> Theme:
+    return Theme(
+        {
+            "ui.strong": "bold #000000",
+            "ui.label": "#2a2a2a",
+            "ui.muted": "grey42",
+            "ui.heading": "bold #1a1a1a",
+            "ui.dim": "grey62",
+            "ui.accent": "bold purple",
+            "ui.border": "blue",
+            "ui.good": "bold green",
+            "ui.warn": "yellow",
+            "progress.percentage": "bold blue",
+        }
+    )
+
+
+def make_console(**kwargs) -> Console:
+    """Return a rich Console pre-loaded with the mlx_lm theme."""
+    kwargs.setdefault("highlight", False)
+    kwargs.setdefault("color_system", "truecolor")
+    kwargs.setdefault("width", _terminal_width())
+    return Console(theme=_make_theme(), **kwargs)
+
+
+def print_header_panel(
+    console: Console, title: str, rows: list[tuple[str, str]]
+) -> None:
+    """Render the boxed header used by the chat and training entry points."""
+    label_w = max(len(k) for k, _ in rows)
+    body = "\n".join(
+        f"  [ui.label]{k.ljust(label_w)}[/ui.label]   [ui.strong]{v}[/ui.strong]"
+        for k, v in rows
+    )
+    console.print(
+        Panel(
+            body,
+            title=f"[ui.accent]{title}[/ui.accent]",
+            title_align="left",
+            border_style="ui.border",
+            box=ROUNDED,
+            padding=(0, 2),
+        )
+    )
+
+
+def print_chat_help(console: Console) -> None:
+    console.print(
+        "  [ui.label]commands[/ui.label]    "
+        "[ui.strong]q[/ui.strong] [ui.muted]exit[/ui.muted]   "
+        "[ui.strong]r[/ui.strong] [ui.muted]reset[/ui.muted]   "
+        "[ui.strong]h[/ui.strong] [ui.muted]help[/ui.muted]"
+    )
+
+
+def make_corridor_prompt(console: Console):
+    """Return a callable that draws the chat input corridor.
+
+    The returned callable draws the top/bottom rules around the input line,
+    repositions the cursor onto the middle line, and returns the styled
+    "›" prompt string. Pass that string to ``input()`` so readline treats
+    the marker as part of the prompt — otherwise backspace will erase it.
+    """
+
+    _ANSI_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]")
+
+    def _readline_safe(text: str) -> str:
+        # Wrap escape sequences in \x01..\x02 so readline doesn't count
+        # them when computing the prompt's visible width.
+        return _ANSI_RE.sub(lambda m: f"\x01{m.group(0)}\x02", text)
+
+    def _draw() -> str:
+        width = console.width
+        dashes = "─" * max(width - 1, 10)
+        with console.capture() as cap:
+            console.print(f"[ui.muted]{dashes}[/ui.muted]")
+            console.print()
+            console.print(f"[ui.muted]{dashes}[/ui.muted]")
+        sys.stdout.write(cap.get())
+        # Move the cursor up two rows back onto the blank middle line.
+        sys.stdout.write("\033[2A\r")
+        sys.stdout.flush()
+        with console.capture() as cap2:
+            console.print("[ui.accent]›[/ui.accent] ", end="")
+        return _readline_safe(cap2.get())
+
+    return _draw
+
+
+class SquareBar(ProgressColumn):
+    """Progress bar rendered with █/░ blocks plus eighth-block sub-precision."""
+
+    _EIGHTHS = "▏▎▍▌▋▊▉"  # 1/8 .. 7/8
+
+    def __init__(self, bar_width: int = 40, complete_style: str = "blue"):
+        super().__init__()
+        self.bar_width = bar_width
+        self.complete_style = complete_style
+
+    def render(self, task):
+        if not task.total:
+            return Text("░" * self.bar_width, style="ui.dim")
+        pct = min(max(task.completed / task.total, 0.0), 1.0)
+        total_eighths = int(pct * self.bar_width * 8)
+        full = total_eighths // 8
+        rem = total_eighths % 8
+        text = Text()
+        text.append("█" * full, style=self.complete_style)
+        used = full
+        if rem > 0 and full < self.bar_width:
+            text.append(self._EIGHTHS[rem - 1], style=self.complete_style)
+            used += 1
+        text.append("░" * (self.bar_width - used), style="ui.dim")
+        return text
+
+
+def make_train_progress(console: Console, *, disable: bool = False) -> Progress:
+    return Progress(
+        TextColumn("[bold blue]train[/bold blue]"),
+        SquareBar(bar_width=30, complete_style="blue"),
+        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+        TextColumn("[ui.muted]·[/ui.muted]"),
+        TextColumn(
+            "[bold blue]{task.completed:>5,}[/bold blue]"
+            "[ui.muted]/{task.total:,}[/ui.muted]"
+        ),
+        console=console,
+        transient=False,
+        disable=disable,
+    )

mlx_lm/generate.py

@@ -223,7 +223,7 @@ def setup_arg_parser():
 
 
 # A stream on the default device just for generation
-generation_stream = mx.new_stream(mx.default_device())
+generation_stream = mx.new_thread_local_stream(mx.default_device())
 
 
 @contextlib.contextmanager
@@ -1497,6 +1497,7 @@ class BatchGenerator:
     def __init__(
         self,
         model: nn.Module,
+        *,
         max_tokens: int = 128,
         stop_tokens: Optional[Sequence[Sequence[int]]] = None,
         sampler: Optional[Callable[[mx.array], mx.array]] = None,
@@ -1507,6 +1508,7 @@ class BatchGenerator:
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
         max_kv_size: Optional[int] = None,
+        stream=None,
     ):
         self.model = model
         self.max_tokens = max_tokens
@@ -1518,6 +1520,8 @@ class BatchGenerator:
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
         self.max_kv_size = max_kv_size
 
+        self._stream = stream or generation_stream
+
         self._default_state_machine = SequenceStateMachine(
             {"normal": [(seq, None) for seq in stop_tokens]} if stop_tokens else {},
             initial="normal",
@@ -1544,9 +1548,13 @@ class BatchGenerator:
         else:
             self._old_wired_limit = None
 
+    @property
+    def stream(self):
+        return self._stream
+
     def close(self):
         if self._old_wired_limit is not None:
-            mx.synchronize(generation_stream)
+            mx.synchronize(self._stream)
             mx.set_wired_limit(self._old_wired_limit)
             self._old_wired_limit = None
 
@@ -1843,7 +1851,7 @@ class BatchGenerator:
         Returns:
             Tuple of prompt processing responses and generation responses.
         """
-        with mx.stream(generation_stream):
+        with mx.stream(self._stream):
             return self._next()
 
     def next_generated(self):
@@ -1853,7 +1861,7 @@ class BatchGenerator:
         Returns:
             List of GenerationBatch.Response objects
         """
-        with mx.stream(generation_stream):
+        with mx.stream(self._stream):
             while True:
                 prompt_responses, generation_responses = self._next()
                 if not generation_responses and prompt_responses:

mlx_lm/lora.py

@@ -12,6 +12,7 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml
 
+from .cli_ui import make_console, print_header_panel
 from .tuner.callbacks import get_reporting_callbacks
 from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
@@ -23,6 +24,12 @@ from .tuner.utils import (
 )
 from .utils import _parse_size, load, save_config
 
+
+def printf(*args, **kwargs):
+    if mx.distributed.init().rank() == 0:
+        print(*args, **kwargs)
+
+
 yaml_loader = yaml.SafeLoader
 yaml_loader.add_implicit_resolver(
     "tag:yaml.org,2002:float",
@@ -246,7 +253,7 @@ def train_model(
 
     # Resume from weights if provided
     if args.resume_adapter_file is not None:
-        print(f"Loading fine-tuned weights from {args.resume_adapter_file}")
+        printf(f"Loading fine-tuned weights from {args.resume_adapter_file}")
         model.load_weights(args.resume_adapter_file, strict=False)
 
     print_trainable_parameters(model)
@@ -291,6 +298,8 @@ def train_model(
 
     opt = opt_class(learning_rate=lr, **optimizer_config)
 
+    _print_run_header(args)
+
     # Train model
     train(
         model=model,
@@ -302,18 +311,50 @@ def train_model(
     )
 
 
+def _print_run_header(args):
+    rank = mx.distributed.init().rank()
+    if rank != 0:
+        return
+
+    type_label = args.fine_tune_type
+    if args.fine_tune_type in ("lora", "dora"):
+        rank = args.lora_parameters.get("rank", "?")
+        type_label = f"{args.fine_tune_type} · {args.num_layers} layers · rank {rank}"
+    elif args.fine_tune_type == "full":
+        type_label = f"full · {args.num_layers} layers"
+
+    lr = (
+        args.learning_rate
+        if isinstance(args.learning_rate, (int, float))
+        else "schedule"
+    )
+    lr_str = f"{lr:.1e}" if isinstance(lr, (int, float)) else lr
+
+    rows = [
+        ("model", str(args.model)),
+        ("type", type_label),
+        ("dataset", str(args.data)),
+        ("optimizer", f"{args.optimizer} · lr {lr_str}"),
+        ("batch · iters", f"{args.batch_size} · {args.iters:,}"),
+        ("max seq", f"{args.max_seq_length:,}"),
+    ]
+    print_header_panel(make_console(), "mlx_lm.lora", rows)
+
+
 def evaluate_model(args, model: nn.Module, test_set):
+    rank = mx.distributed.init().rank()
     test_loss = evaluate(
         model=model,
         dataset=CacheDataset(test_set),
         batch_size=args.batch_size,
         num_batches=args.test_batches,
         max_seq_length=args.max_seq_length,
+        progress=(rank == 0),
     )
 
     test_ppl = math.exp(test_loss)
 
-    print(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
+    printf(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
 
 
 def run(args, training_callback: TrainingCallback = None):
@@ -325,10 +366,10 @@ def run(args, training_callback: TrainingCallback = None):
         config=vars(args),
     )
 
-    print("Loading pretrained model")
+    printf("Loading pretrained model")
     model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})
 
-    print("Loading datasets")
+    printf("Loading datasets")
     train_set, valid_set, test_set = load_dataset(args, tokenizer)
 
     if args.test and not args.train:
@@ -337,13 +378,13 @@ def run(args, training_callback: TrainingCallback = None):
             load_adapters(model, args.adapter_path)
 
     elif args.train:
-        print("Training")
+        printf("Training")
         train_model(args, model, train_set, valid_set, training_callback)
     else:
         raise ValueError("Must provide at least one of --train or --test")
 
     if args.test:
-        print("Testing")
+        printf("Testing")
         evaluate_model(args, model, test_set)
 
 
@@ -351,10 +392,11 @@ def main():
     os.environ["TOKENIZERS_PARALLELISM"] = "true"
     parser = build_parser()
     args = parser.parse_args()
+
     config = args.config
     args = vars(args)
     if config:
-        print("Loading configuration file", config)
+        printf("Loading configuration file", config)
         with open(config, "r") as file:
             config = yaml.load(file, yaml_loader)
         # Prefer parameters from command-line arguments

mlx_lm/models/apertus.py

@@ -167,7 +167,8 @@ class Model(nn.Module):
         self.args = args
         self.model_type = args.model_type
         self.model = ApertusModel(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
         self,
@@ -175,12 +176,18 @@ class Model(nn.Module):
         cache: Optional[Any] = None,
     ) -> mx.array:
         out = self.model(inputs, cache)
-        return self.lm_head(out)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
 
     def sanitize(self, weights):
         for k, v in weights.items():
             if k.endswith("alpha_p") or k.endswith("alpha_n"):
                 weights[k] = v.squeeze()
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
         return weights
 
     @property

mlx_lm/models/cache.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map, tree_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
 
 from .base import create_causal_mask
 
@@ -603,6 +603,18 @@ class ArraysCache(_BaseCache):
         if left_padding:
             self.left_padding = mx.array(left_padding)
 
+    @property
+    def batch_size(self):
+        for c in self.cache:
+            if c is not None:
+                return c.shape[0]
+        if self.left_padding is not None:
+            return self.left_padding.size
+        elif self.lengths is not None:
+            return self.lengths.size
+        else:
+            return 1
+
     def __setitem__(self, idx, value):
         self.cache[idx] = value
 
@@ -622,6 +634,8 @@ class ArraysCache(_BaseCache):
         In-place filter to keep just the given indices in the cache.
         """
         self.cache = [c[batch_indices] if c is not None else None for c in self.cache]
+        if self.left_padding is not None:
+            self.left_padding = self.left_padding[batch_indices]
         if self.lengths is not None:
             self.lengths = self.lengths[batch_indices]
 
@@ -630,14 +644,31 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
+        a_batch = self.batch_size
+        b_batch = other.batch_size
+
         def cat(a, b):
+            shape = dtype = None
+            if a is not None:
+                shape = a.shape
+                dtype = a.dtype
+            if b is not None:
+                shape = b.shape
+                dtype = b.dtype
+
+            if shape is None:
+                return None
+
             if a is None:
-                return b
+                a = mx.zeros((a_batch,) + shape[1:], dtype=dtype)
             if b is None:
-                return a
+                b = mx.zeros((b_batch,) + shape[1:], dtype=dtype)
+
             return mx.concatenate([a, b])
 
         self.cache = [cat(c, o) for c, o in zip(self.cache, other.cache)]
+        self.left_padding = cat(self.left_padding, other.left_padding)
+        self.lengths = cat(self.lengths, other.lengths)
 
     def extract(self, idx):
         cache = ArraysCache(len(self.cache))
@@ -675,6 +706,7 @@ class ArraysCache(_BaseCache):
 
         # All caches are empty so return early
         if all(c.empty() for c in caches):
+            cache.left_padding = mx.array([0] * B)
             return cache
 
         for e in range(n_state):
@@ -1024,8 +1056,9 @@ class BatchKVCache(_BaseCache):
         def pad(c):
             k, v = c.keys, c.values
             if k is None:
-                k = mx.array([]).reshape(B, H, 0, D)
-                v = mx.array([]).reshape(B, H, 0, M)
+                Bc = c.offset.shape[0]
+                k = mx.array([]).reshape(Bc, H, 0, D)
+                v = mx.array([]).reshape(Bc, H, 0, M)
             left = max_idx - c._idx
             right = max_size - k.shape[2] - left
             if right < 0:
@@ -1360,8 +1393,9 @@ class BatchRotatingKVCache(_BaseCache):
             left = max_idx - c._idx
             k, v = c.keys, c.values
             if k is None:
-                k = mx.array([]).reshape(B, H, 0, D)
-                v = mx.array([]).reshape(B, H, 0, M)
+                Bc = c.offset.shape[0]
+                k = mx.array([]).reshape(Bc, H, 0, D)
+                v = mx.array([]).reshape(Bc, H, 0, M)
             right = max_size - k.shape[2] - left
             if right < 0:
                 k = k[..., :right, :]

mlx_lm/models/gemma4_text.py

@@ -180,6 +180,7 @@ class Attention(nn.Module):
         self.layer_idx = layer_idx
         self.layer_type = config.layer_types[layer_idx]
         self.is_sliding = self.layer_type == "sliding_attention"
+        self.has_kv = layer_idx < config.num_hidden_layers - config.num_kv_shared_layers
 
         self.head_dim = (
             config.global_head_dim
@@ -202,14 +203,18 @@ class Attention(nn.Module):
         self.scale = 1.0
 
         self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-        if not self.use_k_eq_v:
-            self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+        if self.has_kv:
+            self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+            if not self.use_k_eq_v:
+                self.v_proj = nn.Linear(
+                    dim, self.n_kv_heads * self.head_dim, bias=False
+                )
         self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
 
         self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.v_norm = RMSNormNoScale(self.head_dim, eps=config.rms_norm_eps)
+        if self.has_kv:
+            self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.v_norm = RMSNormNoScale(self.head_dim, eps=config.rms_norm_eps)
 
         # RoPE (with partial rotation support)
         layer_key = "sliding_attention" if self.is_sliding else "full_attention"
@@ -238,6 +243,10 @@ class Attention(nn.Module):
 
         if shared_kv is not None:
             keys, values = shared_kv
+        elif not self.has_kv:
+            raise ValueError(
+                f"Layer {self.layer_idx} is a KV-shared layer but received no shared_kv"
+            )
         else:
             keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
             values = keys
@@ -600,6 +609,7 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         sanitized = {}
+        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
         for k, v in weights.items():
             if any(
                 s in k
@@ -613,6 +623,18 @@ class Model(nn.Module):
             ):
                 continue
 
+            # KV-shared layers reuse K/V from earlier layers — drop their projections
+            if any(
+                s in k
+                for s in (".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.k_norm")
+            ):
+                try:
+                    layer_idx = int(k.split("layers.")[1].split(".")[0])
+                    if layer_idx >= first_kv_shared:
+                        continue
+                except (IndexError, ValueError):
+                    pass
+
             if k.endswith(".experts.gate_up_proj"):
                 base = k.removesuffix(".gate_up_proj")
                 gate, up = map(mx.contiguous, mx.split(v, 2, axis=-2))

mlx_lm/quant/dwq.py

@@ -314,7 +314,11 @@ def main():
 
     if args.target_dir is not None:
         target_dir = Path(args.target_dir)
-        has_targets = target_dir.exists()
+        has_targets = (
+            target_dir.is_dir()
+            and any((target_dir / "train").glob("*.safetensors"))
+            and any((target_dir / "valid").glob("*.safetensors"))
+        )
     else:
         has_targets = False
         target_dir = None

mlx_lm/server.py

@@ -10,7 +10,7 @@ import time
 import uuid
 import warnings
 from collections import deque
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
 from queue import Empty as QueueEmpty
@@ -36,7 +36,6 @@ from ._version import __version__
 from .generate import (
     BatchGenerator,
     SequenceStateMachine,
-    generation_stream,
     stream_generate,
 )
 from .models.cache import (
@@ -78,7 +77,14 @@ class ToolCallFormatter:
 
         result = []
         for tool_text in tool_calls:
-            parsed = self._tool_parser(tool_text, self._tools)
+            try:
+                parsed = self._tool_parser(tool_text, self._tools)
+            except (ValueError, json.JSONDecodeError) as e:
+                logging.warning(
+                    f"Failed to parse tool call ({type(e).__name__}: {e}) — "
+                    f"tool text was likely truncated mid-generation."
+                )
+                continue
             if not isinstance(parsed, list):
                 parsed = [parsed]
             result.extend(self._format(tc) for tc in parsed)
@@ -227,6 +233,22 @@ class Response:
     top_tokens: Tuple[Dict[str, Any]]
 
 
+def _process_control_tokens(ctx, token_stream):
+    buffer_size = max(len(s) for s in ctx.sequences)
+    buffered_stream = deque()
+
+    for tok in token_stream:
+        buffered_stream.append(tok)
+        if tok.match is not None:
+            popped = [buffered_stream.pop() for _ in tok.match]
+            for t in reversed(popped):
+                buffered_stream.append(replace(t, text=""))
+        if len(buffered_stream) >= buffer_size:
+            yield buffered_stream.popleft()
+    while len(buffered_stream) > 0:
+        yield buffered_stream.popleft()
+
+
 class TimeBudget:
     def __init__(self, budget=0.5, iterations=25, sync_frequency=10):
         self._is_distributed = mx.distributed.init().size() > 1
@@ -256,8 +278,7 @@ class TimeBudget:
         self._loops += 1
         self._time_spent += time.time() - self._start
         if self._loops % self._sync_frequency == 0:
-            with mx.stream(generation_stream):
-                loop_time = mx.distributed.all_sum(self._time_spent).item()
+            loop_time = mx.distributed.all_sum(self._time_spent).item()
             avg_loop_time = loop_time / (
                 mx.distributed.init().size() * self._sync_frequency
             )
@@ -285,94 +306,92 @@ class ModelProvider:
         )
         self.is_distributed = group.size() > 1
 
-        # Preload the default model if it is provided
-        self.default_model_map = {}
-        if self.cli_args.model is not None:
-            self.default_model_map[self.cli_args.model] = "default_model"
-            self.load(self.cli_args.model, draft_model_path="default_model")
+        # Maps model and adapter paths the actual paths to be used. Used to
+        # map 'default_model' to the provided model by cli argument but could
+        # be used for more in the future.
+        self._model_map = {}
+        self._adapter_map = {}
+        self._draft_model_map = {}
+        self._model_map["default_model"] = self.cli_args.model
+        self._adapter_map["default_model"] = self.cli_args.adapter_path
+        self._draft_model_map["default_model"] = self.cli_args.draft_model
 
-    # Added in adapter_path to load dynamically
-    def load(self, model_path, adapter_path=None, draft_model_path=None):
-        model_path = self.default_model_map.get(model_path, model_path)
-        if self.model_key == (model_path, adapter_path, draft_model_path):
-            return self.model, self.tokenizer
+        # Build the tokenizer config for later use in load
+        self._tokenizer_config = {
+            "trust_remote_code": True if cli_args.trust_remote_code else None
+        }
+        if cli_args.chat_template:
+            self._tokenizer_config["chat_template"] = cli_args.chat_template
+
+    def _load(self, model_path, adapter_path=None, draft_model_path=None):
+        if self.is_distributed and (
+            adapter_path is not None or draft_model_path is not None
+        ):
+            raise ValueError(
+                "Loading with adapters or draft models not supported in distributed mode"
+            )
 
         # Remove the old model if it exists.
+        self.model_key = None
         self.model = None
         self.tokenizer = None
-        self.model_key = None
         self.draft_model = None
 
-        # Building tokenizer_config
-        tokenizer_config = {
-            "trust_remote_code": True if self.cli_args.trust_remote_code else None
-        }
-        if self.cli_args.chat_template:
-            tokenizer_config["chat_template"] = self.cli_args.chat_template
-
-        if model_path == "default_model":
-            if self.cli_args.model is None:
-                raise ValueError(
-                    "A model path has to be given as a CLI "
-                    "argument or in the HTTP request"
-                )
-            adapter_path = adapter_path or self.cli_args.adapter_path
-            # TODO: Generalize distributed load
-            if self.is_distributed:
-                model, tokenizer = sharded_load(
-                    self.cli_args.model, self.pipeline_group, self.tensor_group
-                )
-            else:
-                model, tokenizer = load(
-                    self.cli_args.model,
-                    adapter_path=adapter_path,
-                    tokenizer_config=tokenizer_config,
-                )
+        # Load the model and tokenizer
+        if self.is_distributed:
+            model, tokenizer = sharded_load(
+                model_path,
+                pipeline_group=self.pipeline_group,
+                tensor_group=self.tensor_group,
+                tokenizer_config=self._tokenizer_config,
+            )
         else:
-            # TODO: Generalize distributed load
-            if self.is_distributed:
-                model, tokenizer = sharded_load(
-                    model_path, self.pipeline_group, self.tensor_group
-                )
-            else:
-                model, tokenizer = load(
-                    model_path,
-                    adapter_path=adapter_path,
-                    tokenizer_config=tokenizer_config,
-                )
+            model, tokenizer = load(
+                model_path,
+                adapter_path=adapter_path,
+                tokenizer_config=self._tokenizer_config,
+            )
 
+        # Use the default chat template if needed
         if self.cli_args.use_default_chat_template:
             if tokenizer.chat_template is None:
                 tokenizer.chat_template = tokenizer.default_chat_template
 
-        self.model_key = (model_path, adapter_path, draft_model_path)
-        self.model = model
-        self.tokenizer = tokenizer
-
-        def validate_draft_tokenizer(draft_tokenizer):
-            # Check if tokenizers are compatible
+        # Load the draft model for speculative decoding
+        draft_model = None
+        if draft_model_path is not None:
+            draft_model, draft_tokenizer = load(draft_model_path)
             if draft_tokenizer.vocab_size != tokenizer.vocab_size:
                 logging.warning(
                     "Draft model tokenizer does not match model tokenizer. "
                     "Speculative decoding may not work as expected."
                 )
 
-        # Load draft model if specified
-        if (
-            draft_model_path == "default_model"
-            and self.cli_args.draft_model is not None
-        ):
-            self.draft_model, draft_tokenizer = load(self.cli_args.draft_model)
-            validate_draft_tokenizer(draft_tokenizer)
+        # Compute batchability
+        is_batchable = draft_model is None
+        is_batchable = is_batchable and all(
+            hasattr(c, "merge") for c in make_prompt_cache(model)
+        )
 
-        elif draft_model_path is not None and draft_model_path != "default_model":
-            self.draft_model, draft_tokenizer = load(draft_model_path)
-            validate_draft_tokenizer(draft_tokenizer)
+        # Update the member variables
+        self.model_key = (model_path, adapter_path, draft_model_path)
+        self.model = model
+        self.tokenizer = tokenizer
+        self.draft_model = draft_model
+        self.is_batchable = is_batchable
 
-        if self.draft_model is None:
-            self.is_batchable = all(
-                hasattr(c, "merge") for c in make_prompt_cache(self.model)
-            )
+    def load_default(self):
+        if self._model_map["default_model"] is not None:
+            self.load("default_model", None, "default_model")
+
+    def load(self, model_path, adapter_path=None, draft_model_path=None):
+        model_path = self._model_map.get(model_path, model_path)
+        adapter_path = self._adapter_map.get(model_path, adapter_path)
+        draft_model_path = self._draft_model_map.get(draft_model_path, draft_model_path)
+
+        model_key = (model_path, adapter_path, draft_model_path)
+        if self.model_key != model_key:
+            self._load(*model_key)
 
         return self.model, self.tokenizer
 
@@ -466,22 +485,21 @@ class ResponseGenerator:
         if not self._is_distributed:
             return obj
 
-        with mx.stream(generation_stream):
-            if self._rank == 0:
-                if obj is None:
-                    mx.eval(mx.distributed.all_sum(0))
-                    return None
-                data = mx.array(pickle.dumps(obj))
-                mx.eval(mx.distributed.all_sum(data.size))
-                mx.eval(mx.distributed.all_sum(data))
-                return obj
-            else:
-                size = mx.distributed.all_sum(0).item()
-                if size == 0:
-                    return None
-                data = mx.zeros(size, dtype=mx.uint8)
-                data = mx.distributed.all_sum(data)
-                return pickle.loads(data)
+        if self._rank == 0:
+            if obj is None:
+                mx.eval(mx.distributed.all_sum(0))
+                return None
+            data = mx.array(pickle.dumps(obj))
+            mx.eval(mx.distributed.all_sum(data.size))
+            mx.eval(mx.distributed.all_sum(data))
+            return obj
+        else:
+            size = mx.distributed.all_sum(0).item()
+            if size == 0:
+                return None
+            data = mx.zeros(size, dtype=mx.uint8)
+            data = mx.distributed.all_sum(data)
+            return pickle.loads(data)
 
     def _share_request(self, request):
         if not self._is_distributed:
@@ -651,10 +669,11 @@ class ResponseGenerator:
             ts = tokenizer.tool_call_start_tokens
             te = tokenizer.tool_call_end_tokens
             transitions["normal"].append((ts, "tool"))
-            transitions["tool"] = [(te, "normal")]
+            transitions["tool"] = [(te, "normal")] if te else []
             transitions["tool"].extend(common_stops)
             sequences[ts] = tokenizer.tool_call_start
-            sequences[te] = tokenizer.tool_call_end
+            if te:
+                sequences[te] = tokenizer.tool_call_end
 
         sm = SequenceStateMachine(transitions, initial=initial_state)
         if len(self._state_machine_cache) > 100:
@@ -667,6 +686,14 @@ class ResponseGenerator:
         return self.model_provider.is_batchable and args.seed is None
 
     def _generate(self):
+        # Local thread stream that we 'll pass to the BatchGenerator to make
+        # sure that all generation runs in the same stream as the
+        # synchronization messages.
+        generation_stream = mx.default_stream(mx.default_device())
+
+        # Load the default model if it is given
+        self.model_provider.load_default()
+
         current_model = None
         current_sampling = None
         current_tokenizer = None
@@ -796,6 +823,7 @@ class ResponseGenerator:
                         completion_batch_size=self.cli_args.decode_concurrency,
                         prefill_batch_size=self.cli_args.prompt_concurrency,
                         prefill_step_size=self.cli_args.prefill_step_size,
+                        stream=generation_stream,
                     )
                     unprocessed_requests.append((rqueue, request, args))
                     continue
@@ -885,12 +913,11 @@ class ResponseGenerator:
 
                 uids_to_remove = self._share_object(uids_to_remove)
                 if uids_to_remove:
-                    with mx.stream(generation_stream):
-                        batch_generator.remove(uids_to_remove)
-                        for uid in uids_to_remove:
-                            # It may have already been removed during
-                            # generation
-                            batch_results.pop(uid, None)
+                    batch_generator.remove(uids_to_remove)
+                    for uid in uids_to_remove:
+                        # It may have already been removed during
+                        # generation
+                        batch_results.pop(uid, None)
 
     def _serve_single(self, request):
         rqueue, request, args = request
@@ -1018,20 +1045,6 @@ class ResponseGenerator:
                     continue
                 yield response
 
-        def _process_control_tokens(ctx, token_stream):
-            buffer_size = max(len(s) for s in ctx.sequences)
-            buffered_stream = deque()
-
-            for tok in token_stream:
-                buffered_stream.append(tok)
-                if tok.match is not None:
-                    for _ in tok.match:
-                        buffered_stream.pop()
-                if len(buffered_stream) >= buffer_size:
-                    yield buffered_stream.popleft()
-            while len(buffered_stream) > 0:
-                yield buffered_stream.popleft()
-
         ctx = response_queue.get()
         if isinstance(ctx, Exception):
             raise ctx

mlx_lm/tokenizer_utils.py

@@ -397,6 +397,8 @@ class TokenizerWrapper:
 
     @property
     def think_start_id(self):
+        if self._think_start_tokens is None:
+            return None
         if len(self._think_start_tokens) > 1:
             raise ValueError("The start thinking sequence is more than 1 token")
         return self._think_start_tokens[0]
@@ -411,6 +413,8 @@ class TokenizerWrapper:
 
     @property
     def think_end_id(self):
+        if self._think_end_tokens is None:
+            return None
         if len(self._think_end_tokens) > 1:
             raise ValueError("The end thinking sequence is more than 1 token")
         return self._think_end_tokens[0]

mlx_lm/tool_parsers/gemma4.py

@@ -5,10 +5,19 @@ from typing import Any, Optional
 
 import regex as re
 
+# Matches <|"|>...<|"|> string literals (Gemma 4's string delimiter).
+_GEMMA4_STR = r'<\|"\|>(?:(?!<\|"\|>)[\s\S])*?<\|"\|>'
+
 # Matches call:name{...} with balanced braces via the regex module's
-# recursive (?R)-style support.  (\{(?:[^{}]|(?2))*\}) recurses on the
-# second capture group so nested objects like {a:{b:1}} are captured whole.
-_tool_call_regex = re.compile(r"call:(\w+)(\{(?:[^{}]|(?2))*\})", re.DOTALL)
+# recursive (?R)-style support.  The inner alternatives handle:
+#   [^{}<]          – any char that is not a brace or start of <|"|>
+#   <(?!\|"\|>)     – a lone '<' that is NOT the start of <|"|>
+#   <|"|>...<|"|>   – a complete string literal (braces inside are ignored)
+#   (?2)            – recursively balanced nested brace group
+_tool_call_regex = re.compile(
+    r"call:([\w-]+)(\{(?:[^{}<]|<(?!\|\"\|>)|" + _GEMMA4_STR + r"|(?2))*\})",
+    re.DOTALL,
+)
 
 
 def _gemma4_args_to_json(text: str) -> str:

mlx_lm/tool_parsers/minimax_m2.py

@@ -157,43 +157,44 @@ def _get_param_types_from_config(param_name: str, param_config: dict) -> list[st
 
 
 def parse_tool_call(text: str, tools: list | None = None):
-    invoke_match = _invoke_complete_regex.findall(text)
-    if not invoke_match:
+    invoke_matches = _invoke_complete_regex.findall(text)
+    if not invoke_matches:
         raise ValueError("No tool call found")
-    invoke_text = invoke_match[0]
 
-    name_match = re.search(r"^([^>]+)", invoke_text)
-    if not name_match:
-        return None
-
-    function_name = _extract_name(name_match.group(1))
-
-    # Get parameter configuration
-    param_config = {}
+    param_config_for = {}
     if tools:
         for tool in tools:
             if func := tool.get("function", False):
-                if func["name"] != function_name:
-                    continue
                 if params := func.get("parameters", False):
-                    param_config = params.get("properties", {})
+                    param_config_for[func["name"]] = params.get("properties", {})
 
-    # Extract parameters
-    param_dict = {}
-    for match in _parameter_complete_regex.findall(invoke_text):
-        param_match = re.search(r"^([^>]+)>(.*)", match, re.DOTALL)
-        if param_match:
-            param_name = _extract_name(param_match.group(1))
-            param_value = param_match.group(2).strip()
-            if param_value.startswith("\n"):
-                param_value = param_value[1:]
-            if param_value.endswith("\n"):
-                param_value = param_value[:-1]
+    calls = []
+    for invoke_text in invoke_matches:
+        name_match = re.search(r"^([^>]+)", invoke_text)
+        if not name_match:
+            continue
+        function_name = _extract_name(name_match.group(1))
+        param_config = param_config_for.get(function_name, {})
 
-            param_type = _get_param_types_from_config(param_name, param_config)
+        param_dict = {}
+        for match in _parameter_complete_regex.findall(invoke_text):
+            param_match = re.search(r"^([^>]+)>(.*)", match, re.DOTALL)
+            if param_match:
+                param_name = _extract_name(param_match.group(1))
+                param_value = param_match.group(2).strip()
+                if param_value.startswith("\n"):
+                    param_value = param_value[1:]
+                if param_value.endswith("\n"):
+                    param_value = param_value[:-1]
 
-            param_dict[param_name] = _convert_param_value_with_types(
-                param_value, param_type
-            )
+                param_type = _get_param_types_from_config(param_name, param_config)
 
-    return dict(name=function_name, arguments=param_dict)
+                param_dict[param_name] = _convert_param_value_with_types(
+                    param_value, param_type
+                )
+
+        calls.append(dict(name=function_name, arguments=param_dict))
+
+    if len(calls) == 1:
+        return calls[0]
+    return calls

mlx_lm/tuner/datasets.py

@@ -5,9 +5,15 @@ import types
 from pathlib import Path
 from typing import Any, Dict, List
 
+import mlx.core as mx
 from transformers import PreTrainedTokenizer
 
 
+def printf(*args, **kwargs):
+    if mx.distributed.init().rank() == 0:
+        print(*args, **kwargs)
+
+
 class TextDataset:
     """
     Light-weight wrapper to hold a dataset.
@@ -264,7 +270,7 @@ def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
     collection = []
     for ds in dataset_collection:
         ds_path = ds["path"]
-        print(f"Loading Hugging Face dataset {ds_path}.")
+        printf(f"Loading Hugging Face dataset {ds_path}.")
         ds["mask_prompt"] = getattr(args, "mask_prompt", False)
         config = types.SimpleNamespace(**ds)
         hf_config = ds.get("config", {})
@@ -314,7 +320,7 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
         if data_path.exists():
             train, valid, test = load_local_dataset(data_path, tokenizer, args)
         else:
-            print(f"Loading Hugging Face dataset {args.data}.")
+            printf(f"Loading Hugging Face dataset {args.data}.")
             train, valid, test = load_hf_dataset(args.data, tokenizer, args)
 
     if args.train and len(train) == 0:
@@ -322,7 +328,7 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
             "Training set not found or empty. Must provide training set for fine-tuning."
         )
     if args.train and len(valid) == 0:
-        print(
+        printf(
             "Warning: Validation set not found or empty. Training will proceed without validation."
         )
     if args.test and len(test) == 0:

mlx_lm/tuner/trainer.py

@@ -13,10 +13,16 @@ from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten, tree_map
 from tqdm import tqdm
 
+from ..cli_ui import make_console, make_train_progress
 from .callbacks import TrainingCallback
 from .datasets import CacheDataset
 
 
+def printf(*args, **kwargs):
+    if mx.distributed.init().rank() == 0:
+        print(*args, **kwargs)
+
+
 def _clear_cache(threshold: int):
     if mx.get_cache_memory() > threshold:
         mx.clear_cache()
@@ -147,7 +153,7 @@ def iterate_batches(
                 offsets = [0] * len(batch)
             lengths = [len(x) for x in batch]
             if max(lengths) > max_seq_length:
-                print(
+                printf(
                     f"[WARNING] Some sequences are longer than {max_seq_length} tokens. "
                     f"The longest sentence {max(lengths)} will be truncated to {max_seq_length}. "
                     "Consider pre-splitting your data to save memory."
@@ -182,6 +188,8 @@ def evaluate(
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
     clear_cache_threshold: int = 0,
+    progress: bool = True,
+    batch_callback: callable = None,
 ):
     model.eval()
     all_losses = mx.array(0.0)
@@ -189,24 +197,30 @@ def evaluate(
 
     index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
 
-    for _, batch in tqdm(
-        zip(
-            index_iterator,
-            iterate_batches(
-                dataset=dataset,
-                batch_size=batch_size,
-                max_seq_length=max_seq_length,
-                comm_group=mx.distributed.init(),
-            ),
+    batch_iter = zip(
+        index_iterator,
+        iterate_batches(
+            dataset=dataset,
+            batch_size=batch_size,
+            max_seq_length=max_seq_length,
+            comm_group=mx.distributed.init(),
         ),
-        desc="Calculating loss...",
-        total=min(len(dataset) // batch_size, num_batches),
-    ):
+    )
+    if progress:
+        batch_iter = tqdm(
+            batch_iter,
+            desc="Calculating loss...",
+            total=min(len(dataset) // batch_size, num_batches),
+        )
+
+    for _, batch in batch_iter:
         losses, toks = loss(model, *batch)
         all_losses += losses * toks
         ntokens += toks
         mx.eval(all_losses, ntokens)
         _clear_cache(clear_cache_threshold)
+        if batch_callback is not None:
+            batch_callback()
 
     all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
     ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
@@ -227,12 +241,13 @@ def train(
 ):
     if mx.metal.is_available():
         mx.set_wired_limit(mx.device_info()["max_recommended_working_set_size"])
-    print(f"Starting training..., iters: {args.iters}")
     world = mx.distributed.init()
     world_size = world.size()
     rank = world.rank()
-    if world_size > 1:
-        print(f"Node {rank} of {world_size}")
+
+    console = make_console()
+    if rank == 0 and world_size > 1:
+        console.print(f"[ui.muted]node {rank} of {world_size}[/ui.muted]")
 
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
@@ -269,119 +284,153 @@ def train(
     train_time = 0
     grad_accum = None
 
-    # Main training loop
-    for it, batch in zip(
-        range(1, args.iters + 1),
-        iterate_batches(
-            dataset=train_dataset,
-            batch_size=args.batch_size,
-            max_seq_length=args.max_seq_length,
-            loop=True,
-            comm_group=world,
-        ),
-    ):
-        tic = time.perf_counter()
-        # Report validation loss if needed, the first validation loss
-        # is always measured before any training.
-        if val_dataset and (
-            it == 1 or it % args.steps_per_eval == 0 or it == args.iters
+    progress = make_train_progress(console, disable=(rank != 0))
+    task = progress.add_task("train", total=args.iters)
+
+    if rank == 0:
+        console.print(
+            "  [ui.heading]iter   train_loss     tok/s     tokens[/ui.heading]"
+        )
+    progress.start()
+
+    prev_train_loss = None
+    try:
+        # Main training loop
+        for it, batch in zip(
+            range(1, args.iters + 1),
+            iterate_batches(
+                dataset=train_dataset,
+                batch_size=args.batch_size,
+                max_seq_length=args.max_seq_length,
+                loop=True,
+                comm_group=world,
+            ),
         ):
             tic = time.perf_counter()
-            val_loss = evaluate(
-                model=model,
-                dataset=val_dataset,
-                loss=loss,
-                batch_size=args.batch_size,
-                num_batches=args.val_batches,
-                max_seq_length=args.max_seq_length,
-                iterate_batches=iterate_batches,
-            )
-            model.train()
-            val_time = time.perf_counter() - tic
-            if rank == 0:
-                print(
-                    f"Iter {it}: "
-                    f"Val loss {val_loss:.3f}, "
-                    f"Val took {val_time:.3f}s",
-                    flush=True,
+            # Run validation periodically (skip iter 1).
+            if val_dataset and (it % args.steps_per_eval == 0 or it == args.iters):
+                if args.val_batches == -1:
+                    val_total = len(val_dataset) // args.batch_size
+                else:
+                    val_total = min(
+                        len(val_dataset) // args.batch_size, args.val_batches
+                    )
+                val_task = (
+                    progress.add_task("[bold blue]val  [/bold blue]", total=val_total)
+                    if rank == 0
+                    else None
                 )
 
-            if training_callback is not None:
-                val_info = {
-                    "iteration": it - 1,
-                    "val_loss": val_loss,
-                    "val_time": val_time,
-                }
-                training_callback.on_val_loss_report(val_info)
+                def _advance_val():
+                    if val_task is not None:
+                        progress.advance(val_task)
 
-            tic = time.perf_counter()
-
-        lvalue, toks, grad_accum = step(
-            batch,
-            grad_accum,
-            it % grad_accum_steps == 0,
-        )
-
-        losses += lvalue
-        n_tokens += toks
-        steps += 1
-        mx.eval(state, losses, n_tokens, grad_accum)
-        _clear_cache(args.clear_cache_threshold)
-        train_time += time.perf_counter() - tic
-
-        # Report training loss if needed
-        if it % args.steps_per_report == 0 or it == args.iters:
-            train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
-            train_loss /= steps * world_size
-            n_tokens = mx.distributed.all_sum(n_tokens, stream=mx.cpu).item()
-            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.get_peak_memory() / 1e9
-            if rank == 0:
-                print(
-                    f"Iter {it}: Train loss {train_loss:.3f}, "
-                    f"Learning Rate {learning_rate:.3e}, "
-                    f"It/sec {it_sec:.3f}, "
-                    f"Tokens/sec {tokens_sec:.3f}, "
-                    f"Trained Tokens {trained_tokens}, "
-                    f"Peak mem {peak_mem:.3f} GB",
-                    flush=True,
+                tic = time.perf_counter()
+                val_loss = evaluate(
+                    model=model,
+                    dataset=val_dataset,
+                    loss=loss,
+                    batch_size=args.batch_size,
+                    num_batches=args.val_batches,
+                    max_seq_length=args.max_seq_length,
+                    iterate_batches=iterate_batches,
+                    progress=False,
+                    batch_callback=_advance_val,
                 )
+                model.train()
+                val_time = time.perf_counter() - tic
+                if val_task is not None:
+                    progress.remove_task(val_task)
+                if rank == 0:
+                    progress.console.print(
+                        f"  [ui.muted]{it:>4}[/ui.muted]    "
+                        f"[ui.accent]val[/ui.accent] "
+                        f"[ui.strong]{val_loss:>5.3f}[/ui.strong]    "
+                        f"[ui.muted]({val_time:.2f}s)[/ui.muted]"
+                    )
 
-            if training_callback is not None:
-                train_info = {
-                    "iteration": it,
-                    "train_loss": train_loss,
-                    "learning_rate": learning_rate,
-                    "iterations_per_second": it_sec,
-                    "tokens_per_second": tokens_sec,
-                    "trained_tokens": trained_tokens,
-                    "peak_memory": peak_mem,
-                }
-                training_callback.on_train_loss_report(train_info)
+                if training_callback is not None:
+                    val_info = {
+                        "iteration": it - 1,
+                        "val_loss": val_loss,
+                        "val_time": val_time,
+                    }
+                    training_callback.on_val_loss_report(val_info)
 
-            losses = 0
-            n_tokens = 0
-            steps = 0
-            train_time = 0
+                tic = time.perf_counter()
 
-        # Save adapter weights
-        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 = (
-                Path(args.adapter_file).parent / f"{it:07d}_adapters.safetensors"
-            )
-            mx.save_safetensors(str(checkpoint), adapter_weights)
-            print(
-                f"Iter {it}: Saved adapter weights to "
-                f"{args.adapter_file} and {checkpoint}."
+            lvalue, toks, grad_accum = step(
+                batch,
+                grad_accum,
+                it % grad_accum_steps == 0,
             )
 
+            losses += lvalue
+            n_tokens += toks
+            steps += 1
+            mx.eval(state, losses, n_tokens, grad_accum)
+            _clear_cache(args.clear_cache_threshold)
+            train_time += time.perf_counter() - tic
+
+            progress.advance(task)
+
+            # Report training loss if needed
+            if it % args.steps_per_report == 0 or it == args.iters:
+                train_loss = mx.distributed.all_sum(losses, stream=mx.cpu).item()
+                train_loss /= steps * world_size
+                n_tokens = mx.distributed.all_sum(n_tokens, stream=mx.cpu).item()
+                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.get_peak_memory() / 1e9
+                if rank == 0:
+                    if prev_train_loss is None or train_loss <= prev_train_loss:
+                        arrow, arrow_style = "▼", "green"
+                    else:
+                        arrow, arrow_style = "▲", "yellow"
+                    prev_train_loss = train_loss
+                    progress.console.print(
+                        f"  [ui.muted]{it:>4}[/ui.muted]    "
+                        f"[bold {arrow_style}]{train_loss:>5.3f} {arrow}"
+                        f"[/bold {arrow_style}]    "
+                        f"[ui.strong]{tokens_sec:>5,.0f}[/ui.strong]    "
+                        f"[ui.muted]{trained_tokens / 1000:>5.1f}k[/ui.muted]"
+                    )
+
+                if training_callback is not None:
+                    train_info = {
+                        "iteration": it,
+                        "train_loss": train_loss,
+                        "learning_rate": learning_rate,
+                        "iterations_per_second": it_sec,
+                        "tokens_per_second": tokens_sec,
+                        "trained_tokens": trained_tokens,
+                        "peak_memory": peak_mem,
+                    }
+                    training_callback.on_train_loss_report(train_info)
+
+                losses = 0
+                n_tokens = 0
+                steps = 0
+                train_time = 0
+
+            # Save adapter weights
+            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 = (
+                    Path(args.adapter_file).parent / f"{it:07d}_adapters.safetensors"
+                )
+                mx.save_safetensors(str(checkpoint), adapter_weights)
+                progress.console.print(
+                    f"  [ui.good]save[/ui.good]  "
+                    f"[ui.muted]{checkpoint.name}[/ui.muted]"
+                )
+    finally:
+        progress.stop()
+
     # Save final weights
     if rank == 0:
         adapter_weights = dict(tree_flatten(model.trainable_parameters()))
         mx.save_safetensors(str(args.adapter_file), adapter_weights)
-        print(f"Saved final weights to {args.adapter_file}.")

mlx_lm/tuner/utils.py

@@ -15,6 +15,11 @@ from .dora import DoRAEmbedding, DoRALinear
 from .lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear
 
 
+def printf(*args, **kwargs):
+    if mx.distributed.init().rank() == 0:
+        print(*args, **kwargs)
+
+
 def build_schedule(schedule_config: Dict):
     """
     Build a learning rate schedule from the given config.
@@ -162,7 +167,7 @@ def print_trainable_parameters(model):
     trainable_p = (
         sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 1e6
     )
-    print(
+    printf(
         f"Trainable parameters: {(trainable_p * 100 / total_p):.3f}% "
         f"({trainable_p:.3f}M/{total_p:.3f}M)"
     )

mlx_lm/utils.py

@@ -342,7 +342,7 @@ def load_model(
 
     model = model_class(model_args)
 
-    if hasattr(model, "sanitize"):
+    if weights and hasattr(model, "sanitize"):
         weights = model.sanitize(weights)
 
     def _quantize(quantization):
@@ -507,6 +507,8 @@ def sharded_load(
     pipeline_group: Optional[mx.distributed.Group] = None,
     tensor_group: Optional[mx.distributed.Group] = None,
     return_config: bool = False,
+    *,
+    tokenizer_config: Optional[Dict[str, Any]] = None,
 ):
     # Get model path with everything but weight safetensors
     model_path = _download(
@@ -571,7 +573,7 @@ def sharded_load(
     # Load and shard the model, and load the weights
     tokenizer = load_tokenizer(
         model_path,
-        {"trust_remote_code": True},
+        tokenizer_config or {"trust_remote_code": True},
         eos_token_ids=config.get("eos_token_id", None),
     )
     model, _ = load_model(model_path, lazy=True, strict=False)

setup.py

@@ -10,7 +10,7 @@ sys.path.append(str(package_dir))
 
 from _version import __version__
 
-MIN_MLX_VERSION = "0.30.4"
+MIN_MLX_VERSION = "0.31.2"
 
 setup(
     name="mlx-lm",

tests/test_gguf.py

@@ -35,6 +35,9 @@ class TestConvertToGGUFWithoutMocks(unittest.TestCase):
         mock_tokenizer.get_vocab.return_value = {"<pad>": 0, "hello": 1, "world": 2}
         mock_tokenizer.all_special_tokens = ["<pad>"]
         mock_tokenizer.all_special_ids = [0]
+        mock_tokenizer.bos_token_id = None
+        mock_tokenizer.eos_token_id = None
+        mock_tokenizer.unk_token_id = None
         mock_from_pretrained.return_value = mock_tokenizer
 
         model_path = Path(self.test_dir)

tests/test_models.py

@@ -5,7 +5,7 @@ import unittest
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_map
+from mlx.utils import tree_flatten, tree_map
 
 from mlx_lm.models import rope_utils
 from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
@@ -1547,6 +1547,78 @@ class TestModels(unittest.TestCase):
             mx.allclose(logits, mx.ones((1, 1, 4), dtype=mx.float32) * 32.0)
         )
 
+    def test_gemma4_kv_shared_layers_omit_kv_projections(self):
+        """KV-shared layers must not create k_proj/v_proj/k_norm/v_norm so that
+        models saved without redundant weights (e.g. via transformers
+        save_pretrained) can be loaded with strict=True."""
+        from mlx_lm.models import gemma4_text
+
+        args = gemma4_text.ModelArgs(
+            model_type="gemma4_text",
+            hidden_size=128,
+            num_hidden_layers=10,
+            intermediate_size=256,
+            num_attention_heads=4,
+            head_dim=32,
+            global_head_dim=64,
+            rms_norm_eps=1e-6,
+            vocab_size=1000,
+            vocab_size_per_layer_input=1000,
+            num_key_value_heads=1,
+            num_kv_shared_layers=4,
+            hidden_size_per_layer_input=32,
+            sliding_window=8,
+            sliding_window_pattern=5,
+            final_logit_softcapping=30.0,
+            layer_types=[
+                "sliding_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            rope_parameters={
+                "full_attention": {
+                    "partial_rotary_factor": 0.25,
+                    "rope_theta": 1000000.0,
+                },
+                "sliding_attention": {
+                    "rope_theta": 10000.0,
+                },
+            },
+        )
+        model = gemma4_text.Model(args)
+
+        # Non-shared layers (0-5) should have KV projections
+        for i in range(6):
+            attn = model.model.layers[i].self_attn
+            self.assertTrue(attn.has_kv)
+            self.assertTrue(hasattr(attn, "k_proj"))
+            self.assertTrue(hasattr(attn, "k_norm"))
+
+        # Shared layers (6-9) should NOT have KV projections
+        for i in range(6, 10):
+            attn = model.model.layers[i].self_attn
+            self.assertFalse(attn.has_kv)
+            self.assertFalse(hasattr(attn, "k_proj"))
+            self.assertFalse(hasattr(attn, "k_norm"))
+            self.assertFalse(hasattr(attn, "v_proj"))
+
+        # Verify the model can load weights that omit shared-layer KV params
+        weights = dict(tree_flatten(model.parameters()))
+        kv_keys = [
+            k for k in weights if "k_proj" in k or "v_proj" in k or "k_norm" in k
+        ]
+        for k in kv_keys:
+            # All KV keys should belong to non-shared layers (0-5)
+            layer_idx = int(k.split("layers.")[1].split(".")[0])
+            self.assertLess(layer_idx, 6)
+
     def test_gemma4_input_embeddings_reconstruct_per_layer_inputs(self):
         from mlx_lm.models import gemma4_text
 

tests/test_prompt_cache.py

@@ -662,6 +662,102 @@ class TestPromptCache(unittest.TestCase):
         c_out = KVCache.merge((c1, c2))
         self.assertEqual(c_out.keys.shape, (2, 4, 4, 4))
 
+    def test_extend_with_empty_and_nonempty_batch_caches(self):
+        """Extending a batch cache when one side has keys=None should use the
+        correct batch size for the placeholder, not the batch size from the
+        non-None side. Regression test for broadcast error in dynamic_roll."""
+        H, D = 8, 64
+        max_size = 512
+
+        # -- BatchRotatingKVCache --
+        # Create 2 caches with content and 3 empty caches
+        c1 = RotatingKVCache(max_size=max_size)
+        c2 = RotatingKVCache(max_size=max_size)
+        c1.update_and_fetch(mx.ones((1, H, 5, D)), mx.ones((1, H, 5, D)))
+        c2.update_and_fetch(mx.ones((1, H, 3, D)), mx.ones((1, H, 3, D)))
+        batch_full = BatchRotatingKVCache.merge([c1, c2])
+
+        empty_caches = [RotatingKVCache(max_size=max_size) for _ in range(3)]
+        batch_empty = BatchRotatingKVCache.merge(empty_caches)
+
+        # Extend non-empty with empty (different batch sizes)
+        batch_full.extend(batch_empty)
+        self.assertEqual(batch_full.keys.shape[0], 5)
+        self.assertEqual(batch_full.offset.shape[0], 5)
+
+        # Prompt processing with right padding should not crash
+        batch_full.prepare(lengths=[10, 8, 12, 7, 11], right_padding=[2, 4, 0, 5, 1])
+        new_kv = mx.ones((5, H, 12, D))
+        batch_full.update_and_fetch(new_kv, new_kv)
+
+        # Also test empty extending non-empty
+        batch_full2 = BatchRotatingKVCache.merge(
+            [RotatingKVCache(max_size=max_size) for _ in range(3)]
+        )
+        c3 = RotatingKVCache(max_size=max_size)
+        c4 = RotatingKVCache(max_size=max_size)
+        c3.update_and_fetch(mx.ones((1, H, 4, D)), mx.ones((1, H, 4, D)))
+        c4.update_and_fetch(mx.ones((1, H, 6, D)), mx.ones((1, H, 6, D)))
+        batch_content = BatchRotatingKVCache.merge([c3, c4])
+        batch_full2.extend(batch_content)
+        self.assertEqual(batch_full2.keys.shape[0], 5)
+        self.assertEqual(batch_full2.offset.shape[0], 5)
+
+        # -- BatchKVCache --
+        c1 = KVCache()
+        c2 = KVCache()
+        c1.update_and_fetch(mx.ones((1, H, 5, D)), mx.ones((1, H, 5, D)))
+        c2.update_and_fetch(mx.ones((1, H, 3, D)), mx.ones((1, H, 3, D)))
+        batch_full = BatchKVCache.merge([c1, c2])
+
+        empty_caches = [KVCache() for _ in range(3)]
+        batch_empty = BatchKVCache.merge(empty_caches)
+
+        batch_full.extend(batch_empty)
+        self.assertEqual(batch_full.keys.shape[0], 5)
+        self.assertEqual(batch_full.offset.shape[0], 5)
+
+    def test_arrays_cache_extend_with_empty(self):
+        # test simple merge
+        c1 = ArraysCache(2)
+        c2 = ArraysCache(2)
+        c1[0] = mx.zeros((1, 4, 8))
+        c1[1] = mx.zeros((1, 4))
+        c2[0] = mx.zeros((1, 4, 8))
+        c2[1] = mx.zeros((1, 4))
+        full = ArraysCache.merge((c1, c2))
+        self.assertEqual(full[0].shape, (2, 4, 8))
+
+        # extend with empty
+        empty = ArraysCache.merge((ArraysCache(2),))
+        full.extend(empty)
+        self.assertEqual(full[0].shape, (3, 4, 8))
+        self.assertEqual(full[1].shape, (3, 4))
+        self.assertTrue(mx.all(full[0][2:] == 0))
+
+        # making an empty cache with 2 sequences and merging it with
+        # another one with 2 sequences
+        empty2 = ArraysCache.merge((ArraysCache(2), ArraysCache(2)))
+        content = ArraysCache.merge((c1, c2))
+        empty2.extend(content)
+        self.assertEqual(empty2[0].shape, (4, 4, 8))
+        self.assertEqual(empty2[1].shape, (4, 4))
+
+        # Extend content with empty
+        content = ArraysCache.merge((c1, c2))
+        empty2 = ArraysCache.merge((ArraysCache(2), ArraysCache(2)))
+        content.extend(empty2)
+        self.assertEqual(content[0].shape, (4, 4, 8))
+        self.assertEqual(content[1].shape, (4, 4))
+        self.assertEqual(content.make_mask(10).shape, (4, 10))
+
+        # multiple empty extensions accumulate correctly
+        stepwise = ArraysCache.merge((c1,))
+        stepwise.extend(ArraysCache(2))
+        stepwise.extend(ArraysCache.merge((ArraysCache(2), ArraysCache(2))))
+        self.assertEqual(stepwise[0].shape, (4, 4, 8))
+        self.assertEqual(stepwise[1].shape, (4, 4))
+
     def test_window_mask_with_full_kv_cache(self):
         c = KVCache()
         kv = mx.zeros((1, 1, 32, 128))

tests/test_server.py

@@ -4,13 +4,20 @@ import http
 import io
 import json
 import threading
+import types
 import unittest
 
 import mlx.core as mx
 import requests
 
 from mlx_lm.models.cache import KVCache
-from mlx_lm.server import APIHandler, LRUPromptCache, ResponseGenerator
+from mlx_lm.server import (
+    APIHandler,
+    LRUPromptCache,
+    Response,
+    ResponseGenerator,
+    _process_control_tokens,
+)
 from mlx_lm.utils import load
 
 
@@ -61,6 +68,9 @@ class DummyModelProvider:
         assert model in ["default_model", "chat_model"]
         return self.model, self.tokenizer
 
+    def load_default(self):
+        return self.load("default_model", None, "default_model")
+
 
 class MockCache:
     def __init__(self, value, is_trimmable: bool = True):
@@ -82,6 +92,71 @@ class MockCache:
         return n
 
 
+class TestProcessControlTokens(unittest.TestCase):
+    @staticmethod
+    def _r(text, state, match=None):
+        return Response(text, 0, state, match, 0.0, None, ())
+
+    def test_single_tool_call_passes_body_with_open_and_close_crossings(self):
+        r = self._r
+        stream = [
+            r("hi ", "normal"),
+            r("<tool_call>", "tool", match=(0,)),
+            r("body", "tool"),
+            r("</tool_call>", "normal", match=(1,)),
+            r(" bye", "normal"),
+        ]
+        ctx = types.SimpleNamespace(
+            sequences={(0,): "<tool_call>", (1,): "</tool_call>"}
+        )
+        out = list(_process_control_tokens(ctx, iter(stream)))
+
+        self.assertEqual("".join(t.text for t in out), "hi body bye")
+        states = [t.state for t in out]
+        self.assertEqual(sum(1 for a, b in zip(states, states[1:]) if a != b), 2)
+
+    def test_back_to_back_tool_calls_emit_state_crossings(self):
+        r = self._r
+        stream = [
+            r("<tool_call>", "tool", match=(0,)),
+            r("call1_body", "tool"),
+            r("</tool_call>", "normal", match=(1,)),
+            r("<tool_call>", "tool", match=(0,)),
+            r("call2_body", "tool"),
+            r("</tool_call>", "normal", match=(1,)),
+        ]
+        ctx = types.SimpleNamespace(
+            sequences={(0,): "<tool_call>", (1,): "</tool_call>"}
+        )
+        out = list(_process_control_tokens(ctx, iter(stream)))
+
+        self.assertEqual("".join(t.text for t in out), "call1_bodycall2_body")
+        states = [t.state for t in out]
+        crossings = sum(
+            1 for a, b in zip(states, states[1:]) if a == "tool" and b == "normal"
+        )
+        self.assertEqual(crossings, 2)
+
+    def test_multi_token_match_preserves_order(self):
+        r = self._r
+        match = (10, 11, 12)
+        stream = [
+            r("body", "tool"),
+            r("</", "tool"),
+            r("tool", "tool"),
+            r("_call>", "normal", match=match),
+            r(" ok", "normal"),
+        ]
+        ctx = types.SimpleNamespace(sequences={match: "</tool_call>"})
+        out = list(_process_control_tokens(ctx, iter(stream)))
+
+        self.assertEqual([t.text for t in out], ["body", "", "", "", " ok"])
+        self.assertEqual(
+            [t.state for t in out],
+            ["tool", "tool", "tool", "normal", "normal"],
+        )
+
+
 class TestServer(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -205,6 +280,33 @@ class TestServer(unittest.TestCase):
         self.assertIn("id", response_body)
         self.assertIn("choices", response_body)
 
+    def test_make_state_machine_empty_tool_call_end(self):
+        class FakeTokenizer:
+            has_thinking = False
+            has_tool_calling = True
+            tool_call_start = "[TOOL_CALLS]"
+            tool_call_end = ""
+            tool_call_start_tokens = (100,)
+            tool_call_end_tokens = ()
+            eos_token_ids = [2]
+
+            def convert_ids_to_tokens(self, t):
+                return f"<eos{t}>"
+
+        sm, _ = self.response_generator._make_state_machine(
+            ("fake-empty-end", None, None),
+            FakeTokenizer(),
+            stop_words=[],
+        )
+        state = sm.make_state()
+        state, _, s = sm.match(state, 100)
+        self.assertEqual(s, "tool")
+        for tok in [42, 43, 44]:
+            state, _, s = sm.match(state, tok)
+            self.assertEqual(s, "tool")
+        state, _, s = sm.match(state, 2)
+        self.assertIsNone(s)
+
     def test_handle_models(self):
         url = f"http://localhost:{self.port}/v1/models"
         response = requests.get(url)

tests/test_tokenizers.py

@@ -101,6 +101,14 @@ class TestTokenizers(unittest.TestCase):
         self.assertEqual(tokenizer.think_start, "<think>")
         self.assertEqual(tokenizer.think_end, "</think>")
 
+        tokenizer_repo = "mlx-community/Llama-3.2-1B-Instruct-4bit"
+        tokenizer = load_tokenizer(tokenizer_repo)
+        self.assertFalse(tokenizer.has_thinking)
+        self.assertIsNone(tokenizer.think_start)
+        self.assertIsNone(tokenizer.think_end)
+        self.assertIsNone(tokenizer.think_start_id)
+        self.assertIsNone(tokenizer.think_end_id)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_tool_parsing.py

@@ -254,6 +254,27 @@ class TestToolParsing(unittest.TestCase):
             {"settings": {"enabled": True, "name": "test"}},
         )
 
+        # Hyphenated function name (e.g. manim-video)
+        test_case = (
+            'call:manim-video{mode:<|"|>plan<|"|>,prompt:<|"|>explain KV caching<|"|>}'
+        )
+        tool_call = gemma4.parse_tool_call(test_case, None)
+        self.assertEqual(tool_call["name"], "manim-video")
+        self.assertEqual(
+            tool_call["arguments"],
+            {"mode": "plan", "prompt": "explain KV caching"},
+        )
+
+        # Braces inside a string argument (e.g. code snippets or markdown in content)
+        test_case = (
+            'call:skill_manage{action:<|"|>create<|"|>,'
+            'content:<|"|>use a dict like {key: value} in your code<|"|>}'
+        )
+        tool_call = gemma4.parse_tool_call(test_case, None)
+        self.assertEqual(tool_call["name"], "skill_manage")
+        self.assertEqual(tool_call["arguments"]["action"], "create")
+        self.assertIn("{", tool_call["arguments"]["content"])
+
     def test_kimi_k2(self):
         # Single tool call
         test_case = (
@@ -292,6 +313,22 @@ class TestToolParsing(unittest.TestCase):
         ]
         self.assertEqual(tool_calls, expected)
 
+    def test_minimax_m2(self):
+        test_case = (
+            '<invoke name="search">\n'
+            '<parameter name="query">weather</parameter>\n'
+            "</invoke>\n"
+            '<invoke name="read_file">\n'
+            '<parameter name="path">/tmp/test.txt</parameter>\n'
+            "</invoke>"
+        )
+        expected = [
+            {"name": "search", "arguments": {"query": "weather"}},
+            {"name": "read_file", "arguments": {"path": "/tmp/test.txt"}},
+        ]
+        tool_calls = minimax_m2.parse_tool_call(test_case, None)
+        self.assertEqual(expected, tool_calls)
+
 
 if __name__ == "__main__":
     unittest.main()