electron-rare/mlx-lm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

UI for lora and chat

Browse Source
This commit is contained in:
Anastasiia Filippova
2026-05-26 13:49:35 +02:00
parent df1d3f3c9a
commit be13b3f735
4 changed files with 509 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/chat.py
+54 -9
View File
@@ -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
+266
View File
@@ -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,
)
mlx_lm/lora.py
+29
View File
@@ -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,
mlx_lm/tuner/trainer.py
+160 -116
View File
@@ -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}.")