From be13b3f735343f71e150a2e18faf5663ac4f2ff0 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Tue, 26 May 2026 13:49:35 +0200 Subject: [PATCH] UI for lora and chat --- mlx_lm/chat.py | 63 +++++++-- mlx_lm/cli_ui.py | 266 ++++++++++++++++++++++++++++++++++++++ mlx_lm/lora.py | 29 +++++ mlx_lm/tuner/trainer.py | 276 +++++++++++++++++++++++----------------- 4 files changed, 509 insertions(+), 125 deletions(-) create mode 100644 mlx_lm/cli_ui.py diff --git a/mlx_lm/chat.py b/mlx_lm/chat.py index e168e28..cf4c84d 100644 --- a/mlx_lm/chat.py +++ b/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__": diff --git a/mlx_lm/cli_ui.py b/mlx_lm/cli_ui.py new file mode 100644 index 0000000..fe243df --- /dev/null +++ b/mlx_lm/cli_ui.py @@ -0,0 +1,266 @@ +# 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 and exposes an adaptive theme so the same markup +reads well on both light and dark terminal backgrounds. +""" + +import os +import re +import shutil +import sys +import time + +from rich.box import ROUNDED +from rich.console import Console +from rich.panel import Panel +from rich.progress import ( + Progress, + ProgressColumn, + TextColumn, + TimeElapsedColumn, + TimeRemainingColumn, +) +from rich.text import Text +from rich.theme import Theme + + +def _osc11_to_rgb(timeout: float = 0.1): + """Ask the terminal for its background color via OSC 11. + + Returns an (r, g, b) tuple in the 0-255 range, or None if the terminal + does not respond (non-TTY, unsupported terminal, redirected stdio, ...). + """ + if not (sys.stdin.isatty() and sys.stdout.isatty()): + return None + try: + import select + import termios + import tty + except ImportError: + return None # Windows / restricted environments + + fd = sys.stdin.fileno() + try: + saved = termios.tcgetattr(fd) + except termios.error: + return None + + try: + tty.setraw(fd) + sys.stdout.write("\033]11;?\033\\") + sys.stdout.flush() + + deadline = time.monotonic() + timeout + buf = b"" + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + if not select.select([fd], [], [], remaining)[0]: + break + chunk = os.read(fd, 64) + if not chunk: + break + buf += chunk + if buf.endswith(b"\x07") or buf.endswith(b"\x1b\\"): + break + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, saved) + + match = re.search(rb"rgb:([0-9a-fA-F]+)/([0-9a-fA-F]+)/([0-9a-fA-F]+)", buf) + if not match: + return None + + def _to_byte(hex_bytes: bytes) -> int: + # OSC 11 components are typically 4 hex digits (16-bit) but some + # terminals reply with 2. Normalize to 8 bits by scaling. + digits = hex_bytes.decode("ascii") + value = int(digits, 16) + full = (1 << (4 * len(digits))) - 1 + return round(value * 255 / full) if full else 0 + + return tuple(_to_byte(g) for g in match.groups()) + + +def _detect_dark_background() -> bool: + override = os.environ.get("MLX_LM_THEME", "").strip().lower() + if override in ("dark", "light"): + return override == "dark" + + rgb = _osc11_to_rgb() + if rgb is not None: + r, g, b = rgb + # Perceived luminance (Rec. 601). < 128 ≈ dark background. + return (0.299 * r + 0.587 * g + 0.114 * b) < 128 + + # COLORFGBG is "fg;bg" or "fg;default;bg" with ANSI color indices. + cfb = os.environ.get("COLORFGBG", "") + if cfb: + last = cfb.split(";")[-1].strip() + try: + bg = int(last) + # 0-6 are the dim base colors and 8 is dark grey; 7 and 9-15 + # are the bright/light variants. + return bg in (0, 1, 2, 3, 4, 5, 6, 8) + except ValueError: + pass + + # No signal from the terminal — assume dark, which is the modern default. + return True + + +IS_DARK_BACKGROUND = _detect_dark_background() + + +def _make_theme() -> Theme: + if IS_DARK_BACKGROUND: + styles = { + "ui.strong": "bold white", + "ui.label": "grey70", + "ui.muted": "grey62", + "ui.heading": "bold grey62", + "ui.dim": "grey50", + } + else: + styles = { + "ui.strong": "bold #000000", + "ui.label": "#2a2a2a", + "ui.muted": "grey42", + "ui.heading": "bold #1a1a1a", + "ui.dim": "grey62", + } + styles.update( + { + "ui.accent": "bold purple", + "ui.border": "blue", + "ui.good": "bold green", + "ui.warn": "yellow", + "progress.elapsed": "default", + "progress.remaining": "default", + "progress.percentage": "bold blue", + } + ) + return Theme(styles) + + +def make_console(**kwargs) -> Console: + """Return a rich Console pre-loaded with the adaptive mlx_lm theme.""" + kwargs.setdefault("highlight", False) + # Force truecolor so hex values in the theme survive instead of being + # downgraded to ANSI colors that the terminal may remap. + kwargs.setdefault("color_system", "truecolor") + 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 = shutil.get_terminal_size((80, 24)).columns + 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]" + ), + TextColumn("[ui.muted]·[/ui.muted]"), + TimeElapsedColumn(), + TextColumn("[ui.muted]<[/ui.muted]"), + TimeRemainingColumn(), + console=console, + transient=False, + disable=disable, + ) diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index b5fcce3..8c3ff14 100644 --- a/mlx_lm/lora.py +++ b/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 @@ -291,6 +292,8 @@ def train_model( opt = opt_class(learning_rate=lr, **optimizer_config) + _print_run_header(args) + # Train model train( model=model, @@ -302,6 +305,32 @@ def train_model( ) +def _print_run_header(args): + 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): test_loss = evaluate( model=model, diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py index 3bd5e5c..07fefb5 100644 --- a/mlx_lm/tuner/trainer.py +++ b/mlx_lm/tuner/trainer.py @@ -13,6 +13,7 @@ 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 @@ -182,6 +183,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 +192,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 +236,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 +279,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}.")