ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/tests/studio/run_real_mlx_smoke.py
Daniel Han 02ea6c9233
 tests: unblock three stale assertions broken on main (MLX CI + Backend CI) (#5803) 
* tests: unblock three stale assertions broken on main

MLX CI on Mac M1 + Backend CI (both Repo tests CPU and Python 3.10/11/12/13)
have been red on every push to main for days. None of the underlying code
is wrong; three test files have stale anchors / assertions left behind by
PR #5537 (max_steps bump) and PR #5775 (composer + provision-desktop-auth).

1. tests/studio/run_real_mlx_smoke.py:393
   PR #5537 bumped max_steps from 7 to 30 for seed-robust convergence but
   left `assert len(losses_per_step) == 7`. With logging_steps=1 the
   callback fires once per step; 30 entries, not 7. Track config.max_steps
   so the gate auto-follows future bumps.

2. tests/studio/test_composer_rtl_bidi_attribute.py:29
   PR #5775 changed the composer aria-label from the literal
   `aria-label="Message input"` to a JSX ternary
   `aria-label={overlay ? "Image edit instructions" : "Message input"}`.
   Anchor on the inner string literal `"Message input"` instead.

3. studio/backend/tests/test_desktop_auth.py:487
   The guarded_import in test_provision_desktop_auth_writes_secret_and_creates_db_without_backend_deps
   blocks any import whose name == "utils", including the relative
   `from .utils import echo` inside typer._click.decorators (typer 0.25+).
   Gate the block on level == 0 so only absolute imports of `utils` /
   `auth` / `fastapi` / `structlog` are rejected; relative imports
   inside third-party packages pass through.

All three tests pass locally; the MLX one is a mechanical 7->config.max_steps
swap and will be exercised by MLX CI on this PR.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-27 00:30:30 -07:00

724 lines
28 KiB
Python
Raw Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""
End-to-end MLX smoke test on real Apple Silicon -- multi-process driver.
Two subcommands so the workflow can drive cold-start reloads in fresh
Python processes (the way real users hit the load path):
python run_real_mlx_smoke.py train --workdir DIR
python run_real_mlx_smoke.py reload --format {lora|merged|gguf} --dir D
The `train` subcommand:
1. Loads `unsloth/gemma-3-270m-it` via FastMLXModel.from_pretrained.
2. Applies LoRA r=8 on q/k/v/o.
3. Computes pre-training loss + grad norm via mx.nn.value_and_grad.
4. Trains 7 deterministic steps on a dataset of the SAME row repeated
("<<HELLO!!>> My name is Unsloth!"), with batch_size=2 and
gradient_accumulation_steps=3 so each step processes 6 sequences
and the run sees 42 sequences total.
5. Computes post-training loss + grad norm.
6. Generates from "<<HELLO!!>> My name is " and asserts "Unsloth"
appears in the in-memory completion.
7. Saves the trained model in three formats:
- LoRA adapter (save_pretrained_merged save_method="lora")
- Merged 16-bit (save_pretrained_merged save_method="merged_16bit")
- GGUF (save_pretrained_gguf, best-effort -- skipped with a
clear reason if save raises; e.g. llama.cpp's
convert_hf_to_gguf currently asserts on Gemma-3-270m's
tokenizer vocab. Soft-skipped so the LoRA + merged checks
continue to gate the PR.)
8. Emits `train_metrics.json` with per-phase timing / peak GPU /
peak RSS / per-step losses / pre+post grad norms / generations
/ gguf_supported flag, for regression detection across CI runs.
Reloads run as separate workflow steps so each is a fresh Python
process. For lora / merged the reload uses
FastMLXModel.from_pretrained directly. For gguf the reload spawns
the llama-cli binary built by save_pretrained_gguf and parses
stdout. Each subcommand emits `<format>_reload_metrics.json` next
to the saved dir.
The two upstream unsloth_zoo bugs the earlier draft of this script
worked around are fixed in unslothai/unsloth-zoo#627: GGUF export
no longer raises NotImplementedError on Apple Silicon (llama_cpp.py
catches it from the device_type module-level call) and LoRA reload
via FastMLXModel.from_pretrained(lora_dir) works without an external
config.json copy (mlx_loader.py preserves local_path when config.json
is missing so the adapter_config.json branch can run).
Determinism: seeds Python `random`, `numpy`, and `mlx.core.random` in
every process before any MLX operation. Forwards `random_state=SEED`
to FastMLXModel.from_pretrained / get_peft_model and `seed=SEED` to
MLXTrainingConfig. Metal still has minor reduction-order
nondeterminism, so loss assertions are bounds rather than exact.
Only runnable on a real Apple Silicon host; invoked from
.github/workflows/mlx-ci.yml on the macos-14 runner.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import random as _random
import resource
import subprocess
import sys
import time
from pathlib import Path
import numpy as np
SEED = 3407
TRAIN_TEXT = "<<HELLO!!>> My name is Unsloth!"
PROMPT = "<<HELLO!!>> My name is "
EXPECT_IN_OUTPUT = "Unsloth"
MODEL_NAME = "unsloth/gemma-3-270m-it"
# ---------------------------------------------------------------------------
# Determinism + telemetry helpers
# ---------------------------------------------------------------------------
def _seed_everything() -> None:
_random.seed(SEED)
np.random.seed(SEED)
import mlx.core as mx
mx.random.seed(SEED)
def _peak_gpu_gb() -> float:
import mlx.core as mx
if not mx.metal.is_available():
return 0.0
# Newer MLX deprecates mx.metal.get_peak_memory in favour of the
# top-level mx.get_peak_memory; fall back to the old API for
# compatibility with older MLX versions still present in the
# environment.
getter = getattr(mx, "get_peak_memory", None) or getattr(
mx.metal, "get_peak_memory", None
)
if getter is None:
return 0.0
try:
return float(getter()) / (1024**3)
except Exception:
return 0.0
def _peak_rss_gb() -> float:
"""Peak resident set size for this process. macOS getrusage returns
bytes; Linux returns kilobytes."""
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
if sys.platform == "darwin":
return float(rss) / (1024**3)
return float(rss) / (1024**2)
class Phase:
"""Wall-clock + memory tracker for a named phase. Records into a
metrics dict so we can later JSON-dump for regression detection."""
def __init__(self, name: str, metrics: dict):
self.name = name
self.metrics = metrics
def __enter__(self):
self._t0 = time.perf_counter()
print(f"\n=== phase:{self.name} START ===", flush = True)
return self
def __exit__(self, exc_type, exc, tb):
elapsed = time.perf_counter() - self._t0
peak_gpu = _peak_gpu_gb()
peak_rss = _peak_rss_gb()
self.metrics.setdefault("phases", {})[self.name] = {
"elapsed_seconds": round(elapsed, 3),
"peak_gpu_gb": round(peak_gpu, 3),
"peak_rss_gb": round(peak_rss, 3),
"ok": exc_type is None,
}
status = "OK" if exc_type is None else f"FAIL ({exc_type.__name__})"
print(
f"=== phase:{self.name} {status} elapsed={elapsed:.2f}s "
f"peak_gpu={peak_gpu:.2f}GB peak_rss={peak_rss:.2f}GB ===",
flush = True,
)
return False # don't swallow exceptions
def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, float]:
"""One forward+backward of next-token cross-entropy on `text`.
Returns (loss, ||grad||_2)."""
import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten
ids = list(tokenizer.encode(text))
eos_id = getattr(tokenizer, "eos_token_id", None)
if eos_id is not None:
ids.append(int(eos_id))
if len(ids) < 2:
raise RuntimeError(f"text too short to compute loss: {len(ids)} tokens")
inputs = mx.array([ids[:-1]], dtype = mx.int32)
targets = mx.array([ids[1:]], dtype = mx.int32)
def loss_fn(m):
logits = m(inputs)
return nn.losses.cross_entropy(logits, targets, reduction = "mean")
loss_and_grad = nn.value_and_grad(model, loss_fn)
loss_val, grad = loss_and_grad(model)
norm_sq = mx.array(0.0, dtype = mx.float32)
for _name, value in tree_flatten(grad):
v = value.astype(mx.float32)
norm_sq = norm_sq + mx.sum(v * v)
return float(loss_val.item()), float(mx.sqrt(norm_sq).item())
def _teacher_forced_completion_loss(
model, tokenizer, prompt: str, completion: str
) -> float:
"""Mean next-token CE loss on `completion` tokens given `prompt` (teacher
forced -- no decoding, no sampling, no greedy argmax).
Decouples the memorisation check from greedy-decode geometry. A 47-round,
13-seed sweep on this fixture showed greedy `completion in output` lands
in the 46-77% range across MLX configs (config-fragile), while
post_train_loss is < 0.1 in 100% of configs that reach the basin. Teacher-
forced completion loss is a subset of post_train_loss so it inherits the
same reliability AND is more specific: it asserts *what* the model
memorised, not just *that* it reached low loss on the full row.
Args:
model: the LoRA-trained MLX model
tokenizer: the tokenizer used during training (must match)
prompt: the conditioning text (e.g. PROMPT)
completion: the substring the model should have learnt to emit
after `prompt` (e.g. EXPECT_IN_OUTPUT + "!")
Returns mean cross-entropy over the completion's tokens.
"""
import mlx.core as mx
import mlx.nn as nn
prompt_ids = list(tokenizer.encode(prompt))
full_ids = list(tokenizer.encode(prompt + completion))
if len(full_ids) <= len(prompt_ids):
raise RuntimeError(
f"completion {completion!r} tokenises to zero new tokens after "
f"{prompt!r}; check tokenizer / chat template."
)
inputs = mx.array([full_ids[:-1]], dtype = mx.int32)
targets = mx.array([full_ids[1:]], dtype = mx.int32)
logits = model(inputs)
# logits at position i predict targets[i]; completion tokens occupy
# target positions [len(prompt_ids)-1 ... len(full_ids)-2].
start = len(prompt_ids) - 1
completion_logits = logits[:, start:, :]
completion_targets = targets[:, start:]
loss = nn.losses.cross_entropy(
completion_logits, completion_targets, reduction = "mean"
)
return float(loss.item())
def _write_metrics(path: Path, metrics: dict) -> None:
path.write_text(json.dumps(metrics, indent = 2, default = str))
print(f"\n[metrics] wrote {path}", flush = True)
print(json.dumps(metrics, indent = 2, default = str), flush = True)
# ---------------------------------------------------------------------------
# `train` subcommand
# ---------------------------------------------------------------------------
def cmd_train(args) -> int:
_seed_everything()
metrics: dict = {
"subcommand": "train",
"seed": SEED,
"model": MODEL_NAME,
"train_text": TRAIN_TEXT,
"prompt": PROMPT,
"phases": {},
}
workdir = Path(args.workdir).resolve()
workdir.mkdir(parents = True, exist_ok = True)
import mlx.core as mx
from unsloth_zoo.mlx.loader import FastMLXModel
from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig
hf_token = os.environ.get("HF_TOKEN") or None
with Phase("load_base", metrics):
model, tokenizer = FastMLXModel.from_pretrained(
MODEL_NAME,
load_in_4bit = False,
dtype = "float16",
text_only = True,
max_seq_length = 128,
random_state = SEED,
token = hf_token,
trust_remote_code = False,
)
metrics["base_src_path"] = str(getattr(model, "_src_path", "") or "")
mx.random.seed(SEED)
with Phase("apply_lora", metrics):
# Standard unsloth LoRA target set (q/k/v/o + gate/up/down).
# With bs=2 grad_accum=3 (effective batch 6) the q/k/v/o-only
# LoRA collapsed in 7 steps -- training loss kept dropping but
# inference output the structural skeleton ("My name") without
# recovering the specific "Unsloth" token. Including the MLP
# projections gives the LoRA enough capacity to memorize the
# training row at the larger effective batch.
model = FastMLXModel.get_peft_model(
model,
r = 8,
lora_alpha = 16,
lora_dropout = 0.0,
target_modules = [
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
],
use_gradient_checkpointing = False,
random_state = SEED,
finetune_language_layers = True,
finetune_attention_modules = True,
finetune_mlp_modules = True,
)
with Phase("pre_train_grad_probe", metrics):
pre_loss, pre_norm = _compute_loss_and_grad_norm(model, tokenizer, TRAIN_TEXT)
metrics["pre_train_loss"] = round(pre_loss, 4)
metrics["pre_train_grad_norm"] = round(pre_norm, 4)
assert math.isfinite(pre_loss) and math.isfinite(pre_norm) and pre_norm > 0
losses_per_step: list[float] = []
with Phase("train", metrics):
config = MLXTrainingConfig(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 3,
# 47-round mlx-parity-probes sweep (PR #5498 / staging-2#119)
# found 7 steps is below the convergence horizon at any clip
# setting -- the trainer hasn't memorized the train row yet
# when the smoke probes loss/generation. At 30 steps every
# seed tested hits post_train_loss=0 across all clip
# configurations, so 30 is the seed-robust gate.
max_steps = 30,
learning_rate = 1e-3,
warmup_steps = 0,
lr_scheduler_type = "constant",
optim = "adamw",
weight_decay = 0.0,
# max_grad_value (elementwise) is materially cheaper than
# max_grad_norm on MLX -- norm clip needs a cross-tree
# reduction + materializing all grad tensors at full
# precision, value clip is tree_map(mx.clip) per leaf.
# MLXTrainingConfig defaults to max_grad_value=1.0 for
# exactly this reason; pin both explicitly here so the
# configured clip matches what runs (the trainer prints a
# notice when both > 0 and value wins, so disable norm).
# Empirical 13-seed pass rate at this fixture: value=1.0
# 62%, norm=1.0 46%, value=5.0 33%, value=0.5 77% -- the
# cheaper default is also the higher-pass-rate default.
max_grad_norm = 0.0,
max_grad_value = 1.0,
logging_steps = 1,
max_seq_length = 64,
seed = SEED,
use_cce = False,
compile = False,
gradient_checkpointing = False,
output_dir = str(workdir / "trainer_outputs"),
save_steps = 0,
eval_steps = 0,
dataset_text_field = "text",
)
trainer = MLXTrainer(
model = model,
tokenizer = tokenizer,
train_dataset = [{"text": TRAIN_TEXT}] * 64,
args = config,
)
def _on_step(
step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens, grad_norm = None
):
losses_per_step.append(round(float(loss), 4))
grad_text = f" grad={grad_norm:.4f}" if grad_norm is not None else ""
print(
f" step {step}/{total} loss={loss:.4f} lr={lr:.2e} "
f"tok/s={tok_s:.0f} peak={peak_gb:.2f}GB{grad_text}",
flush = True,
)
trainer.add_step_callback(_on_step)
train_result = trainer.train()
metrics["losses_per_step"] = losses_per_step
metrics["train_summary"] = {
k: train_result[k]
for k in (
"train_loss",
"train_runtime",
"train_steps",
"trained_tokens",
"train_samples_per_second",
"compile_enabled",
"patch_mode",
)
if k in train_result
}
# logging_steps=1 + max_steps=N -> N callbacks; track config so the
# gate auto-follows if max_steps is bumped again.
assert (
len(losses_per_step) == config.max_steps
), f"expected {config.max_steps} logged steps, got {losses_per_step}"
for i, l in enumerate(losses_per_step):
# Allow exact 0.0: fp16 per-step loss underflows to 0.0 after
# the LoRA reaches loss=0 around step ~10 with this fixture +
# max_steps=30. That's the memorization success signal, not a
# bug. Lower bound is "finite and >= 0" not "strictly > 0".
assert math.isfinite(l) and 0 <= l < 50, f"step {i+1} loss bad: {l}"
assert (
losses_per_step[-1] < losses_per_step[0] * 1.1
), f"loss diverged: {losses_per_step[0]} -> {losses_per_step[-1]}"
with Phase("post_train_grad_probe", metrics):
post_loss, post_norm = _compute_loss_and_grad_norm(model, tokenizer, TRAIN_TEXT)
metrics["post_train_loss"] = round(post_loss, 4)
metrics["post_train_grad_norm"] = round(post_norm, 4)
assert post_loss < pre_loss, f"post {post_loss} >= pre {pre_loss}"
# Memorisation gate: teacher-forced loss on the training row must
# be very low after 30 steps of overfit-on-one-example. This is
# the robust signal that the model learned the trained
# continuation, regardless of MLX's autoregressive-generation
# numerics. Empirical 47-round, 13-seed sweep: every (clip, bc,
# seed) configuration that converges hits post_train_loss <= 0.05.
# Tighten gate to 0.1.
assert post_loss < 0.1, (
f"post_train_loss={post_loss:.4f} >= 0.1 -- training did not "
"memorise the single training row in 30 steps. Trainer "
"regression suspected."
)
from mlx_lm import generate
with Phase("inference_in_memory", metrics):
model.eval()
in_mem_out = generate(
model,
tokenizer,
prompt = PROMPT,
max_tokens = 48,
verbose = False,
)
metrics["in_memory_generation"] = in_mem_out
# Soft greedy-decode visibility (metric only). Empirically this lands in
# 46-77% of seeds depending on clip config (47-round, 13-seed sweep) --
# fp16 + MLX attention/generate path puts noticeable noise on the first
# token even after near-zero teacher-forced loss. Surface the mismatch
# for regression tracking, but the next assertion is the load-bearing
# one.
metrics["in_memory_generation_has_expected"] = EXPECT_IN_OUTPUT in in_mem_out
if EXPECT_IN_OUTPUT not in in_mem_out:
print(
f" [INFO] greedy decode did not contain {EXPECT_IN_OUTPUT!r} "
f"(post_train_loss={post_loss:.4f}, completion={in_mem_out!r}). "
"Hard gate is the teacher-forced completion-loss check below.",
flush = True,
)
# Hard check: teacher-forced loss on the completion the model was trained
# to emit. Bypasses greedy-decode fp16 fragility -- if the LoRA actually
# memorised the row, the probability mass on `EXPECT_IN_OUTPUT` after
# `PROMPT` is essentially 1.0 (and the loss essentially 0). 13/13 of the
# MLX configs we measured reached post_train_loss < 1e-3, so this gate
# is deterministic on every (seed, clip, bc) combination tested.
completion_loss = _teacher_forced_completion_loss(
model, tokenizer, PROMPT, EXPECT_IN_OUTPUT + "!"
)
metrics["in_memory_completion_teacher_forced_loss"] = round(completion_loss, 6)
assert completion_loss < 0.5, (
f"teacher-forced completion loss {completion_loss:.4f} >= 0.5: "
f"the LoRA did not memorise {EXPECT_IN_OUTPUT + '!'!r} after "
f"{PROMPT!r} (post_train_loss={post_loss:.4f}). Trainer regression "
"suspected -- check unsloth_zoo MLX trainer gradient clipping / "
"optimizer defaults vs torch.optim.AdamW."
)
# Save LoRA. unsloth-zoo#627 fixed FastMLXModel.from_pretrained(lora_dir)
# so the cold-start reload below works on the saved adapter dir directly.
lora_dir = workdir / "lora"
with Phase("save_lora", metrics):
model.save_pretrained_merged(
str(lora_dir),
tokenizer = tokenizer,
save_method = "lora",
)
metrics["lora_dir"] = str(lora_dir)
assert (lora_dir / "adapters.safetensors").exists()
assert (lora_dir / "adapter_config.json").exists()
# Save merged_16bit (full HF directory)
merged_dir = workdir / "merged_16bit"
with Phase("save_merged_16bit", metrics):
model.save_pretrained_merged(
str(merged_dir),
tokenizer = tokenizer,
save_method = "merged_16bit",
)
metrics["merged_dir"] = str(merged_dir)
assert any(merged_dir.glob("*.safetensors"))
# Save GGUF (best-effort). save_pretrained_gguf clones llama.cpp,
# builds it with cmake (Metal=ON), then runs convert_hf_to_gguf.
# For some models -- including unsloth/gemma-3-270m-it as of
# 2026-05-07 -- llama.cpp's converter asserts on the tokenizer vocab
# (`assert max(tokenizer.vocab.values()) < vocab_size`) because the
# tokenizer carries reserved IDs beyond the embedding matrix size.
# That's an llama.cpp / convert_hf_to_gguf limitation, not an
# unsloth_zoo bug. Soft-skip with a recorded reason so the LoRA +
# merged_16bit assertions still gate the PR.
gguf_dir = workdir / "gguf"
metrics["gguf_supported"] = False
metrics["gguf_skip_reason"] = None
metrics["gguf_dir"] = str(gguf_dir)
with Phase("save_gguf", metrics):
try:
model.save_pretrained_gguf(
str(gguf_dir),
tokenizer = tokenizer,
quantization_method = "not_quantized",
)
gguf_files = sorted(gguf_dir.glob("*.gguf"))
if not gguf_files:
raise RuntimeError(f"no .gguf produced in {gguf_dir}")
metrics["gguf_supported"] = True
metrics["gguf_files"] = [p.name for p in gguf_files]
except Exception as e:
err_text = f"{type(e).__name__}: {e}"
if "AssertionError" in err_text or "tokenizer.vocab" in err_text:
metrics["gguf_skip_reason"] = (
f"llama.cpp convert_hf_to_gguf asserted on tokenizer "
f"vocab for {MODEL_NAME} (max(vocab IDs) >= "
f"vocab_size). Downstream llama.cpp limitation, not "
f"unsloth_zoo. Underlying error: {err_text}"
)
else:
metrics["gguf_skip_reason"] = err_text
print(f" GGUF SKIPPED: {metrics['gguf_skip_reason']}", flush = True)
metrics["final_peak_gpu_gb"] = round(_peak_gpu_gb(), 3)
metrics["final_peak_rss_gb"] = round(_peak_rss_gb(), 3)
_write_metrics(workdir / "train_metrics.json", metrics)
return 0
# ---------------------------------------------------------------------------
# `reload` subcommand (fresh process per format)
# ---------------------------------------------------------------------------
def cmd_reload(args) -> int:
_seed_everything()
save_dir = Path(args.dir).resolve()
if not save_dir.exists():
raise SystemExit(f"reload dir not found: {save_dir}")
metrics: dict = {
"subcommand": "reload",
"format": args.format,
"dir": str(save_dir),
"phases": {},
}
if args.format == "gguf":
return _reload_gguf(save_dir, metrics)
import mlx.core as mx
from unsloth_zoo.mlx.loader import FastMLXModel
from mlx_lm import generate
hf_token = os.environ.get("HF_TOKEN") or None
with Phase(f"reload_{args.format}", metrics):
mx.random.seed(SEED)
m, t = FastMLXModel.from_pretrained(
str(save_dir),
load_in_4bit = False,
dtype = "float16",
text_only = True,
max_seq_length = 128,
random_state = SEED,
token = hf_token,
)
m.eval()
with Phase(f"generate_{args.format}", metrics):
out = generate(m, t, prompt = PROMPT, max_tokens = 48, verbose = False)
metrics["generation"] = out
print(f" [reload:{args.format}] output: {out!r}", flush = True)
# Verify save/reload preserved the trained weights via teacher-
# forced loss on the training row: the reloaded model should have
# approximately the same loss on TRAIN_TEXT as the in-memory model
# had at post_train_loss. This is the real save/reload invariant
# and is robust to MLX's known near-zero-loss adamw greedy-decode
# perturbation (step-7 grad spike at seed=3407, see
# scripts/cuda_mlx_step7_*) which can flip the first generated
# token while leaving teacher-forced loss essentially identical.
train_metrics_path = save_dir.parent / "train_metrics.json"
in_mem_loss = None
in_mem_out = None
if train_metrics_path.exists():
try:
tm = json.loads(train_metrics_path.read_text())
in_mem_loss = tm.get("post_train_loss")
in_mem_out = tm.get("in_memory_generation")
except Exception:
in_mem_loss = None
metrics["in_memory_generation_ref"] = in_mem_out
metrics["in_memory_post_train_loss"] = in_mem_loss
metrics["reload_completion_matches_in_memory"] = (
in_mem_out is not None and out == in_mem_out
)
if isinstance(in_mem_loss, (int, float)) and math.isfinite(in_mem_loss):
reload_loss, _ = _compute_loss_and_grad_norm(m, t, TRAIN_TEXT)
metrics["reload_post_train_loss"] = round(reload_loss, 4)
# float16 round-trip should be near-exact for LoRA + merged;
# 0.2 tolerates the dequant noise we have seen empirically.
assert abs(reload_loss - float(in_mem_loss)) < 0.2, (
f"reload {args.format!r} loss diverged from in-memory: "
f"reload={reload_loss:.4f}, in-memory={in_mem_loss:.4f}"
)
else:
# Fallback when train_metrics.json wasn't found (older
# workdir layouts): keep a non-empty-completion gate.
body = out.replace(PROMPT, "", 1).strip()
assert len(body) >= 4, (
f"reload {args.format!r} produced no usable output for "
f"{PROMPT!r}: {out!r}"
)
metrics["final_peak_gpu_gb"] = round(_peak_gpu_gb(), 3)
metrics["final_peak_rss_gb"] = round(_peak_rss_gb(), 3)
_write_metrics(save_dir.parent / f"{args.format}_reload_metrics.json", metrics)
return 0
def _reload_gguf(save_dir: Path, metrics: dict) -> int:
candidates = [
Path("llama.cpp/llama-cli"),
Path("llama.cpp/build/bin/llama-cli"),
]
llama_cli = next((c for c in candidates if c.exists()), None)
if llama_cli is None:
raise SystemExit(f"llama-cli not found; checked {candidates}")
gguf_files = sorted(save_dir.glob("*.gguf"))
if not gguf_files:
raise SystemExit(f"no .gguf files in {save_dir}")
gguf_path = gguf_files[0]
with Phase("reload_gguf", metrics):
proc = subprocess.run(
[
str(llama_cli),
"-m",
str(gguf_path),
"-p",
PROMPT,
"-n",
"24",
"--temp",
"0",
"--seed",
str(SEED),
"-no-cnv",
"--no-warmup",
],
capture_output = True,
text = True,
timeout = 300,
)
metrics["llama_cli_returncode"] = proc.returncode
metrics["generation"] = (proc.stdout or "")[:1500]
metrics["stderr_head"] = (proc.stderr or "")[:600]
print(f" [reload:gguf] stdout (head):\n{proc.stdout[:800]}", flush = True)
if proc.returncode != 0:
raise SystemExit(
f"llama-cli exit {proc.returncode}; stderr head: {proc.stderr[:400]}"
)
# llama.cpp uses different tokenisation + sampling internals than
# mlx_lm, so the GGUF reload completion does not have to match the
# in-memory completion exactly. Require non-empty, non-prompt-only
# output to catch real save/reload corruption (zero-weight model,
# tokenizer mismatch). Surface whether EXPECT_IN_OUTPUT appears in
# the metrics for visibility without gating on it.
body = (proc.stdout or "").replace(PROMPT, "", 1).strip()
metrics["gguf_has_expected"] = EXPECT_IN_OUTPUT in (proc.stdout or "")
assert len(body) >= 4, (
f"GGUF reload produced no usable output for {PROMPT!r}: "
f"{proc.stdout[:400]!r}"
)
metrics["final_peak_rss_gb"] = round(_peak_rss_gb(), 3)
_write_metrics(save_dir.parent / "gguf_reload_metrics.json", metrics)
return 0
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main() -> int:
parser = argparse.ArgumentParser()
sub = parser.add_subparsers(dest = "cmd", required = True)
p_train = sub.add_parser("train")
p_train.add_argument("--workdir", required = True)
p_reload = sub.add_parser("reload")
p_reload.add_argument(
"--format",
required = True,
choices = ["lora", "merged", "gguf"],
)
p_reload.add_argument("--dir", required = True)
args = parser.parse_args()
if args.cmd == "train":
return cmd_train(args)
if args.cmd == "reload":
return cmd_reload(args)
return 1
if __name__ == "__main__":
sys.exit(main())
Reference in a new issue View git blame Copy permalink