Fix tokenizer test failure

mlx_lm/chat.py

@@ -1,16 +1,9 @@
 # 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
@@ -25,19 +18,6 @@ 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")
@@ -116,8 +96,6 @@ 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)
@@ -137,38 +115,24 @@ def main():
             },
         )
 
-    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: ""
+    def print_help():
+        rprint("The command list:")
+        rprint("- 'q' to exit")
+        rprint("- 'r' to reset the chat")
+        rprint("- 'h' to display these commands")
 
+    rprint(f"[INFO] Starting chat session with {args.model}.")
+    print_help()
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
-        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()
+        query = input(">> " if rank == 0 else "")
         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":
-            if rank == 0:
-                print_chat_help(console)
+            print_help()
             continue
         messages = []
         if args.system_prompt is not None:
@@ -178,7 +142,6 @@ def main():
             messages,
             add_generation_prompt=True,
         )
-        last_response = None
         for response in stream_generate(
             model,
             tokenizer,
@@ -196,15 +159,7 @@ 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

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

@@ -12,7 +12,6 @@ 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
@@ -24,12 +23,6 @@ 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",
@@ -253,7 +246,7 @@ def train_model(
 
     # Resume from weights if provided
     if args.resume_adapter_file is not None:
-        printf(f"Loading fine-tuned weights from {args.resume_adapter_file}")
+        print(f"Loading fine-tuned weights from {args.resume_adapter_file}")
         model.load_weights(args.resume_adapter_file, strict=False)
 
     print_trainable_parameters(model)
@@ -298,8 +291,6 @@ def train_model(
 
     opt = opt_class(learning_rate=lr, **optimizer_config)
 
-    _print_run_header(args)
-
     # Train model
     train(
         model=model,
@@ -311,50 +302,18 @@ 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)
 
-    printf(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
+    print(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
 
 
 def run(args, training_callback: TrainingCallback = None):
@@ -366,10 +325,10 @@ def run(args, training_callback: TrainingCallback = None):
         config=vars(args),
     )
 
-    printf("Loading pretrained model")
+    print("Loading pretrained model")
     model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})
 
-    printf("Loading datasets")
+    print("Loading datasets")
     train_set, valid_set, test_set = load_dataset(args, tokenizer)
 
     if args.test and not args.train:
@@ -378,13 +337,13 @@ def run(args, training_callback: TrainingCallback = None):
             load_adapters(model, args.adapter_path)
 
     elif args.train:
-        printf("Training")
+        print("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:
-        printf("Testing")
+        print("Testing")
         evaluate_model(args, model, test_set)
 
 
@@ -392,11 +351,10 @@ def main():
     os.environ["TOKENIZERS_PARALLELISM"] = "true"
     parser = build_parser()
     args = parser.parse_args()
-
     config = args.config
     args = vars(args)
     if config:
-        printf("Loading configuration file", config)
+        print("Loading configuration file", config)
         with open(config, "r") as file:
             config = yaml.load(file, yaml_loader)
         # Prefer parameters from command-line arguments

mlx_lm/tokenizer_utils.py

@@ -452,12 +452,12 @@ class TokenizerWrapper:
         """
         Get a stateful streaming detokenizer.
         """
-        return self._detokenizer_class(self)
+        if not hasattr(self, "_detokenizer"):
+            self._detokenizer = self._detokenizer_class(self)
+        return self._detokenizer
 
     def __getattr__(self, attr):
-        if attr == "detokenizer":
-            return self._detokenizer
-        elif attr == "eos_token_ids":
+        if attr == "eos_token_ids":
             return self._eos_token_ids
         elif attr.startswith("_"):
             return self.__getattribute__(attr)

mlx_lm/tuner/datasets.py

@@ -5,15 +5,9 @@ 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.
@@ -270,7 +264,7 @@ def load_custom_hf_dataset(args, tokenizer: PreTrainedTokenizer):
     collection = []
     for ds in dataset_collection:
         ds_path = ds["path"]
-        printf(f"Loading Hugging Face dataset {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", {})
@@ -320,7 +314,7 @@ def load_dataset(args, tokenizer: PreTrainedTokenizer):
         if data_path.exists():
             train, valid, test = load_local_dataset(data_path, tokenizer, args)
         else:
-            printf(f"Loading Hugging Face dataset {args.data}.")
+            print(f"Loading Hugging Face dataset {args.data}.")
             train, valid, test = load_hf_dataset(args.data, tokenizer, args)
 
     if args.train and len(train) == 0:
@@ -328,7 +322,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:
-        printf(
+        print(
             "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,16 +13,10 @@ 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()
@@ -153,7 +147,7 @@ def iterate_batches(
                 offsets = [0] * len(batch)
             lengths = [len(x) for x in batch]
             if max(lengths) > max_seq_length:
-                printf(
+                print(
                     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."
@@ -188,8 +182,6 @@ 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)
@@ -197,30 +189,24 @@ def evaluate(
 
     index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
 
-    batch_iter = zip(
-        index_iterator,
-        iterate_batches(
-            dataset=dataset,
-            batch_size=batch_size,
-            max_seq_length=max_seq_length,
-            comm_group=mx.distributed.init(),
+    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(),
+            ),
         ),
-    )
-    if progress:
-        batch_iter = tqdm(
-            batch_iter,
-            desc="Calculating loss...",
-            total=min(len(dataset) // batch_size, num_batches),
-        )
-
-    for _, batch in batch_iter:
+        desc="Calculating loss...",
+        total=min(len(dataset) // batch_size, num_batches),
+    ):
         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)
@@ -241,13 +227,12 @@ 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()
-
-    console = make_console()
-    if rank == 0 and world_size > 1:
-        console.print(f"[ui.muted]node {rank} of {world_size}[/ui.muted]")
+    if world_size > 1:
+        print(f"Node {rank} of {world_size}")
 
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
@@ -284,153 +269,119 @@ def train(
     train_time = 0
     grad_accum = None
 
-    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,
-            ),
+    # 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
         ):
             tic = time.perf_counter()
-            # 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
-                )
-
-                def _advance_val():
-                    if val_task is not None:
-                        progress.advance(val_task)
-
-                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:
-                    val_info = {
-                        "iteration": it - 1,
-                        "val_loss": val_loss,
-                        "val_time": val_time,
-                    }
-                    training_callback.on_val_loss_report(val_info)
-
-                tic = time.perf_counter()
-
-            lvalue, toks, grad_accum = step(
-                batch,
-                grad_accum,
-                it % grad_accum_steps == 0,
+            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,
             )
-
-            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"
+            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,
                 )
-                mx.save_safetensors(str(checkpoint), adapter_weights)
-                progress.console.print(
-                    f"  [ui.good]save[/ui.good]  "
-                    f"[ui.muted]{checkpoint.name}[/ui.muted]"
+
+            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)
+
+            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,
                 )
-    finally:
-        progress.stop()
+
+            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)
+            print(
+                f"Iter {it}: Saved adapter weights to "
+                f"{args.adapter_file} and {checkpoint}."
+            )
 
     # 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,11 +15,6 @@ 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.
@@ -167,7 +162,7 @@ def print_trainable_parameters(model):
     trainable_p = (
         sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 1e6
     )
-    printf(
+    print(
         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 weights and hasattr(model, "sanitize"):
+    if hasattr(model, "sanitize"):
         weights = model.sanitize(weights)
 
     def _quantize(quantization):

tests/test_server.py

@@ -203,8 +203,10 @@ class TestServer(unittest.TestCase):
         self.assertIn("choices", response_body)
         first_text = response_body["choices"][0]["text"]
         self.assertEqual(
-            first_text,
-            json.loads(requests.post(url, json=post_data).text)["choices"][0]["text"],
+            first_text.strip(),
+            json.loads(requests.post(url, json=post_data).text)["choices"][0][
+                "text"
+            ].strip(),
         )
 
     def test_handle_chat_completions(self):

tests/test_tokenizers.py

@@ -65,6 +65,8 @@ class TestTokenizers(unittest.TestCase):
                 tokenizer = load_tokenizer(tokenizer_repo)
                 tokenizer.decode([0, 1, 2])
                 self.assertTrue(isinstance(tokenizer.detokenizer, expected_detokenizer))
+                if expected_detokenizer is BPEStreamingDetokenizer:
+                    tokenizer.detokenizer.clean_spaces = False
                 self.check_tokenizer(tokenizer)
 
         # Try one with a naive detokenizer