From 83c408e1e6586cc28e838efc7262ea7e1d22595d Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Tue, 26 May 2026 20:57:22 +0200 Subject: [PATCH] log only on rank 0 --- mlx_lm/cli_ui.py | 28 +++++++++++++++++++--------- mlx_lm/lora.py | 23 ++++++++++++++++------- mlx_lm/tuner/datasets.py | 12 +++++++++--- mlx_lm/tuner/trainer.py | 7 ++++++- mlx_lm/tuner/utils.py | 7 ++++++- mlx_lm/utils.py | 2 +- 6 files changed, 57 insertions(+), 22 deletions(-) diff --git a/mlx_lm/cli_ui.py b/mlx_lm/cli_ui.py index affcd1c..a9d3523 100644 --- a/mlx_lm/cli_ui.py +++ b/mlx_lm/cli_ui.py @@ -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, diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 43d480b..a24b8f5 100644 --- a/mlx_lm/lora.py +++ b/mlx_lm/lora.py @@ -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 diff --git a/mlx_lm/tuner/datasets.py b/mlx_lm/tuner/datasets.py index ac673f7..7b40a47 100644 --- a/mlx_lm/tuner/datasets.py +++ b/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: diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py index 07fefb5..0cbca52 100644 --- a/mlx_lm/tuner/trainer.py +++ b/mlx_lm/tuner/trainer.py @@ -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." diff --git a/mlx_lm/tuner/utils.py b/mlx_lm/tuner/utils.py index e40155a..3172d19 100644 --- a/mlx_lm/tuner/utils.py +++ b/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)" ) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index ef3d266..f3f6961 100644 --- a/mlx_lm/utils.py +++ b/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):