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

log only on rank 0

Browse Source
This commit is contained in:
Anastasiia Filippova
2026-05-26 20:57:22 +02:00
parent 5b56e28af8
commit 83c408e1e6
6 changed files with 57 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/cli_ui.py
+19 -9
View File
@@ -7,6 +7,7 @@ and training entry points. The theme is hardcoded for a light terminal
background.
"""
import os
import re
import shutil
import sys
@@ -18,13 +19,27 @@ from rich.progress import (
Progress,
ProgressColumn,
TextColumn,
TimeElapsedColumn,
TimeRemainingColumn,
)
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(
{
@@ -37,8 +52,6 @@ def _make_theme() -> Theme:
"ui.border": "blue",
"ui.good": "bold green",
"ui.warn": "yellow",
"progress.elapsed": "default",
"progress.remaining": "default",
"progress.percentage": "bold blue",
}
)
@@ -48,6 +61,7 @@ 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)
@@ -98,7 +112,7 @@ def make_corridor_prompt(console: Console):
return _ANSI_RE.sub(lambda m: f"\x01{m.group(0)}\x02", text)
def _draw() -> str:
width = shutil.get_terminal_size((80, 24)).columns
width = console.width
dashes = "─" * max(width - 1, 10)
with console.capture() as cap:
console.print(f"[ui.muted]{dashes}[/ui.muted]")
@@ -152,10 +166,6 @@ def make_train_progress(console: Console, *, disable: bool = False) -> Progress:
"[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
+16 -7
View File
@@ -24,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",
@@ -247,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)
@@ -336,17 +342,19 @@ def _print_run_header(args):
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):
@@ -358,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:
@@ -370,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)
@@ -384,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/tuner/datasets.py
+9 -3
View File
@@ -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
+6 -1
View File
@@ -18,6 +18,11 @@ 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()
@@ -148,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."
mlx_lm/tuner/utils.py
+6 -1
View File
@@ -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
+1 -1
View File
@@ -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):