SPB Git

spb/zyquo-mlx Public MIT

The local MLX foundry for your Mac — run, fine-tune, quantize, and ship models. Nothing leaves your machine.

Swift 93.4% Python 3.8% Makefile 2.2% Shell 0.5%
19.8 KB

# Zyquo MLX — Fine-Tuning Research (MLX, Ground Truth 2026-07-30)

Phase 0.B research document. Verified against live sources: ml-explore/mlx-lm main branch (version 0.31.3, latest PyPI release 0.31.3, 2026-04-22) and ml-explore/mlx-swift-lm main. All code excerpts were fetched raw from GitHub on 2026-07-30. Items not verifiable against a primary source are marked UNVERIFIED.


# 1. LoRA / QLoRA Workflow (mlx-lm, Python)

# 1.1 Entry points

  • CLI: mlx_lm.lora … or python -m mlx_lm lora … (python -m mlx_lm.lora prints a deprecation notice but works).
  • Training extras: pip install "mlx-lm[train]".
  • Minimal run: mlx_lm.lora --model <hf-repo-or-local-dir> --train --data <dir> --iters 600.
  • YAML config: mlx_lm.lora -c config.yaml; CLI flags override config values, config overrides CONFIG_DEFAULTS.
  • Python API (what PythonRunner scripts will use): mlx_lm.tuner.trainer.{TrainingArgs, train, evaluate}, mlx_lm.tuner.utils.{linear_to_lora_layers, load_adapters, build_schedule}, mlx_lm.tuner.datasets.{load_dataset, CacheDataset}; mlx_lm.lora.run(args, training_callback) mirrors the CLI programmatically.

Source: mlx_lm/lora.py (main).

# 1.2 Complete hyperparameter table (exact names + defaults, CONFIG_DEFAULTS)

CLI flag Config key Default Notes
--model model "Qwen/Qwen3-0.6b" HF repo or local converted dir
--train train False
--data data "mlx-community/WikiSQL" dir with {train,valid,test}.jsonl or HF dataset
--fine-tune-type fine_tune_type "lora" lora | dora | full
--optimizer optimizer "adam" adam, adamw, muon, sgd, adafactor
(config only) optimizer_config {} per optimizer e.g. adamw: {betas:[0.9,0.98], eps:1e-6, weight_decay:0.05}
--mask-prompt mask_prompt False loss on completion only; chat & completions formats only
--num-layers num_layers 16 -1 = all layers; LoRA applies to the last N layers
--batch-size batch_size 4
--iters iters 1000
--val-batches val_batches 25 -1 = full valid set
--learning-rate learning_rate 1e-5
--steps-per-report steps_per_report 10
--steps-per-eval steps_per_eval 200 also evals at iter 1 and final iter
--grad-accumulation-steps grad_accumulation_steps 1 averaged accumulation
--resume-adapter-file resume_adapter_file None see §5.3 — warm start only
--adapter-path adapter_path "adapters"
--save-every save_every 100 checkpoint cadence
--test test False eval on test.jsonl
--test-batches test_batches 500 -1 = full test set
--max-seq-length max_seq_length 2048 longer sequences truncated with warning
--grad-checkpoint grad_checkpoint False per-layer mx.checkpoint
--clear-cache-threshold clear_cache_threshold 0 e.g. 4GB; calls mx.clear_cache() above it
--report-to report_to None wandb, swanlab, or both
--seed seed 0 mx.random.seed + np.random.seed
--trust-remote-code trust_remote_code False
(config only) lora_parameters {rank: 8, dropout: 0.0, scale: 20.0} plus optional keys: […]; not settable via CLI
(config only) lr_schedule None {name, warmup, warmup_init, arguments}
(config only) hf_dataset see §4.3

lora_parameters details: rank (8), scale (20.0 — MLX exposes a single scale factor, not alpha/rank), dropout (0.0), keys (target modules; example yaml uses ["self_attn.q_proj", "self_attn.v_proj"]). If keys is omitted, linear_to_lora_layers adapts all Linear/QuantizedLinear/ SwitchLinear/Embedding modules in each targeted layer. → Zyquo's hyperparameter form maps: rank, scale, dropout, keys, num-layers, LR, batch size, iters, seed, max-seq-length, grad-checkpoint, optimizer, mask-prompt, save-every, steps-per-eval — and must generate a YAML config for lora_parameters/lr_schedule since those have no CLI flags.

lr_schedule (build_schedule): name = any mlx.optimizers.schedulers function (e.g. cosine_decay), arguments positional (first = initial LR), optional linear warmup steps with warmup_init. Example: {name: cosine_decay, warmup: 100, warmup_init: 1e-7, arguments: [1e-5, 1000, 1e-7]}.

# 1.3 QLoRA

No flag: if --model points to a quantized model, training is QLoRA (LORA.md). Quantized checkpoints load as nn.QuantizedLinear; linear_to_lora_layers wraps them keeping the frozen base quantized. Create a quantized base with mlx_lm.convert --hf-path <repo> -q or use any mlx-community/*-4bit repo.

# 1.4 DoRA and full

  • DoRA: --fine-tune-type dora (tuner/dora.py: DoRALinear, DoRAEmbedding). MoE SwitchLinear layers raise "doesn't support DoRA yet".
  • Full: --fine-tune-type full unfreezes the last num_layers layers (-1 = true full fine-tune); lora_parameters becomes None.

# 2. Adapters & Fusing

# 2.1 Contents of --adapter-path (verified)

  • adapter_config.json — the entire resolved training config written at start (save_config(vars(args))); load_adapters() reads fine_tune_type, num_layers, lora_parameters from it to rebuild the LoRA layers.
  • adapters.safetensors — latest weights, overwritten every save_every iterations and at the final iteration.
  • Numbered checkpoints {it:07d}_adapters.safetensors (e.g. 0000100_adapters.safetensors) — one per cadence, 7-digit zero-padded.
  • Contents = tree_flatten(model.trainable_parameters()) — for full fine-tuning this contains all unfrozen weights (same naming).

# 2.2 mlx_lm.fuse (verified argparse)

Flag Default Notes
--model "mlx_model" base model path/repo
--save-path "fused_model"
--adapter-path "adapters"
--upload-repo None push to HF Hub
--dequantize off de-quantize a QLoRA base to fp16 while fusing; strips quantization from config
--export-gguf off llama/mixtral/mistral model types only, fp16
--gguf-path "ggml-model-f16.gguf" written inside save-path

Mechanics: loads model + adapters, calls .fuse(dequantize:) on every capable module, saves a complete standalone MLX model dir (config.json, sharded safetensors + index, tokenizer files) — directly loadable afterwards. (LORA.md mentions --hf-path for fuse but it is absent from current argparse — docs stale; UNVERIFIED if intentional. The pinned 0.31.3 flag is --dequantize, not --de-quantize.)

⚠️ EMPIRICALLY VERIFIED (2026-07-30, this Mac): fusing a QLoRA adapter into a quantized base without --dequantize re-quantizes the merged weights and rounds small LoRA deltas away entirely — our rank-8 adapter's behavior vanished from the fused 4-bit model while --adapter-path inference kept it, and a --dequantize fuse preserved it verbatim. Re-quantizing the fp16 fused model to 4-bit wiped the behavior again — for lightly-trained adapters the deltas are simply below the 4-bit quantization step wherever they land. Zyquo therefore: (1) defaults to de-quantize when fusing onto a quantized base, (2) recommends adapter-attached inference (Python --adapter-path / Swift LoRAContainer.load(into:)) as the lossless default for QLoRA results, and (3) warns before re-quantizing a fused model that adapter effects may not survive unless training was substantial. The UI must explain this.


# 3. Full Fine-Tuning Feasibility

  • Memory with Adam/AdamW: weights + grads + 2 moments, all bf16 ≈ 8 bytes/param ≈ 4× bf16 weights, plus activations (∝ batch × seq; --grad-checkpoint cuts these substantially).
  • Realistic tiers (engineering estimates from the 8 B/param rule — UNVERIFIED, no published Apple table exists):
    • 32 GB: full FT ≤ ~1.5–3B (7B with small --num-layers + grad-checkpoint + batch 1)
    • 48 GB (this Mac): ~3–4B full, 7B partial
    • 64 GB: ~7B full (tight; SGD/Adafactor shrink optimizer state)
    • 128 GB: ~13–15B full
  • Community consensus: full FT > 7B belongs on other hardware — use LoRA/QLoRA.
  • Trainer machinery (checkpoints, metrics, resume file) is identical to LoRA; the "adapter" file simply holds all trainable weights.

# 4. Dataset Formats (truth from mlx_lm/tuner/datasets.py)

# 4.1 Directory layout (--data <dir>)

train.jsonl (required for --train), valid.jsonl (optional — absent → "Warning: Validation set not found or empty. Training will proceed without validation."), test.jsonl (required for --test). One JSON object per line.

# 4.2 Auto-detected formats (detection order, first sample)

  1. completions — both prompt and completion keys → CompletionsDataset. Keys renamable via prompt_feature/completion_feature. Internally converted to a 2-message chat and run through tokenizer.apply_chat_templatecompletions data IS chat-templated, incl. optional per-row tools.
  2. chatmessages key (renamable via chat_feature) → ChatDataset, tokenized via apply_chat_template(messages, tools=d.get("tools")). A tools key (OpenAI function-calling schema) is passed to the template.
  3. texttext key (renamable via text_feature) → raw encode, EOS appended if missing. mask_prompt raises ValueError for text datasets.

# 4.3 HF datasets

  • --data <hf-dataset-id> directly if pre-formatted (train/valid/test splits; needs pip install datasets).
  • Config hf_dataset: (dict or list — concatenated): path, train_split (default "train[:80%]"), valid_split (default "train[-10%:]"), test_split, feature-name overrides, config: dict forwarded to datasets.load_dataset.

# 4.4 mask_prompt (exact behavior)

Each sample is (tokens, offset). With mask_prompt, offset = token length of the template applied to all-but-last message (add_generation_prompt=True for completions; for chat, set when last role is assistant). Loss mask: steps >= offset AND steps <= length — only completion tokens count. Without it, offset = 0.

# 4.5 Batching / truncation (trainer.py iterate_batches)

  • Samples sorted by token length, batched, batch order shuffled per epoch; dataset must have ≥ batch_size examples or ValueError.
  • Padding to 1 + 32·ceil(maxlen/32) capped at max_seq_length; over-long sequences truncated with: [WARNING] Some sequences are longer than {max_seq_length} tokens. The longest sentence {n} will be truncated to {max_seq_length}. Consider pre-splitting your data to save memory.
  • Tokenization lazy + memoized (CacheDataset).
  • No auto train/valid split for local JSONL — Zyquo's DatasetService must produce the split files itself (this is a feature we own).

# 5. Observability

# 5.1 ⚠️ Two stdout regimes exist right now — do not scrape stdout

  • PyPI ≤ 0.31.3 (plain prints, flush=True):
    text
    Iter {it}: Val loss {val_loss:.3f}, Val took {val_time:.3f}s
    Iter {it}: Train loss {t:.3f}, Learning Rate {lr:.3e}, It/sec {i:.3f}, Tokens/sec {tk:.3f}, Trained Tokens {n}, Peak mem {p:.3f} GB
    Iter {it}: Saved adapter weights to {adapter_file} and {checkpoint}.
    Saved final weights to {adapter_file}.
  • main (post-0.31.3): rich-based TrainUI (mlx_lm/cli_ui.py) — ANSI panels, progress bar, columnar rows; learning rate and peak memory no longer printed (still in the callback dict).

Decision for Zyquo: never parse trainer stdout. Drive training via a pinned-version Python helper that registers a custom TrainingCallback and emits a JSON-lines protocol on stdout (our own, stable).

# 5.2 The stable channel: TrainingCallback (tuner/callbacks.py)

  • on_train_loss_report({iteration, train_loss, learning_rate, iterations_per_second, tokens_per_second, trained_tokens, peak_memory}) (peak_memory in GB via mx.get_peak_memory()/1e9)
  • on_val_loss_report({iteration, val_loss, val_time})
  • Built-ins: --report-to wandb|swanlab. No built-in JSON stdout mode — we write our own driver script.

Other startup lines (main): Loading pretrained model, Loading datasets, Training, Trainable parameters: {pct:.3f}% ({M}M/{T}M); --test prints Test loss {:.3f}, Test ppl {:.3f}.

# 5.3 Resume semantics (--resume-adapter-file)

model.load_weights(path, strict=False) only — does NOT restore optimizer state, LR-schedule position, iteration counter, or RNG. Training restarts at iter 1 as a warm start. Prints Loading fine-tuned weights from {path}. → Zyquo's RunStore must persist its own notion of completed iterations and present resume honestly (remaining-iters warm start), or drive the Python API directly to keep optimizer state within one process (pause = in-process, not cross-process).


# 6. Memory & Speed (LoRA/QLoRA)

# 6.1 Official guidance (LORA.md "Memory Issues")

QLoRA (quantized base); reduce --batch-size (4→2→1); --grad-accumulation-steps for effective batch; reduce --num-layers (16→8→4); pre-split long examples; --grad-checkpoint ("more helpful for larger batch sizes or sequence lengths with smaller or quantized models"). Reference datum: Mistral-7B, batch 1, num-layers 4, M1 Max 32 GB → ~250 tokens/sec training.

# 6.2 Community data (indicative, UNVERIFIED individually)

  • QLoRA 7–8B ≈ 7 GB peak (fits 16 GB); LoRA fp16 8B ≈ 14 GB; 14B QLoRA ≈ 12 GB / LoRA ≈ 24 GB; 70B QLoRA ≈ 45 GB (64 GB+, comfortable at 96–128).
  • QLoRA tiers: 8 GB→≤3B, 16 GB→7–8B, 24 GB→8–14B, 32 GB→14B, 48 GB→32B, 64 GB+→32–70B. LoRA fp16 ≈ one tier down.
  • Training throughput: M1 Max ~250 tok/s (matches official), M3 Max ~320, M4 Max ~380, M2 Ultra ~475. --grad-checkpoint ≈ ~30% slower for large activation savings.
  • Peak-memory drivers in order: base weights (dominant), activations (∝ batch × seq² attention + hidden), trainable-layer count (LoRA adapter optimizer state is tiny — rank-8 on 7B ≈ 10–20 MB).

These feed MemoryAdvisor's training-side verdicts; calibrate in Phase 7.


# 7. Evaluation

  • Held-out loss/perplexity: mlx_lm.lora --model <m> --adapter-path <a> --data <dir> --test [--test-batches N]Test loss {:.3f}, Test ppl {:.3f}. (ppl = exp(loss)). Pass --adapter-path "" to test the base — that plus a run with the adapter is the canonical base-vs-tuned scorecard. Python: tuner.trainer.evaluate(model, dataset, batch_size, num_batches, max_seq_length) → avg loss.
  • mlx_lm.evaluate = lm-evaluation-harness integration (lm_eval.simple_evaluate). Flags: --model (req), --tasks (req), --output-dir, --batch-size (16), --num-shots, --max-tokens, --limit, --seed (123), --fewshot-as-multiturn, --apply-chat-template, --chat-template-args, --temp/--top-p/--top-k. No --adapter-path — harness eval of an adapter requires fusing first (or API loading).
  • Qualitative: mlx_lm.generate --model <m> --adapter-path <a> --prompt … runs the adapted model directly, no fuse needed → powers the Playground base-vs-tuned compare.

# 8. Swift-Native Training — State Today

Repo shift (important): MLXLMCommon, MLXLLM, MLXVLM, MLXEmbedders moved from mlx-swift-examples into ml-explore/mlx-swift-lm ("all updates … in the other repository"). mlx-swift-examples now hosts only apps/tools + MLXMNIST/StableDiffusion, depending on mlx-swift ≥ 0.31.4 and swift-transformers ≥ 1.3.0. Zyquo MLX must depend on mlx-swift-lm.

# 8.1 What exists natively in Swift (verified in mlx-swift-lm main)

  • LoRATrain (Libraries/MLXLLM/LoraTrain.swift): Parameters (batchSize=4, iterations=1000, stepsPerReport=10, stepsPerEval=100, validationBatches=10 [0=all], saveEvery=100, adapterURL); train(model:train:validate:optimizer:loss:tokenizer:parameters:progress:) with Progress enum — .train(iteration, trainingLoss, iterationsPerSecond, tokensPerSecond), .validation(…), .save(…) — and ProgressDisposition .stop/.morebuilt-in cooperative cancellation; evaluate(…) -> Float; saveLoRAWeights(model:url:).
  • Adapters (Libraries/MLXLMCommon/Adapters/LoRA/): LoRAConfiguration (Codable; CodingKeys num_layers/fine_tune_type/lora_parametersdocumented as compatible with mlx-lm's adapter_config.json); LoRAContainer.from(directory:) loads adapter_config.json + adapters.safetensors; .load(into:), .fuse(with:), .unload(from:). Layers: LoRALinear, QLoRALinear (→ QLoRA training on quantized bases works natively in Swift; fused() re-quantizes), DoRALinear, QDoRALinear; plus PEFTAdapter.swift (HF-PEFT format) and an adapter factory registry.
  • CLI reference: mlx-swift-examples/Tools/llm-tool/LoraCommands.swift — train/fuse/test/eval subcommands (--adapter, --layers, --resume, --data, --learning-rate, --batch-size, --iterations, --steps-per-report, --steps-per-eval, --validation-batches, --save-every; fuse: --de-quantize).
  • GUI example: LoRATrainingExample — QLoRA on mlx-community/Mistral-7B-v0.1-hf-4bit-mlx, ~4 GB model memory / ~6 GB physical RAM.

# 8.2 Swift gaps vs Python (verified)

  • Dataset loader: only {"text": …} jsonl / plain .txt — no chat-template/ messages/completions/mask_prompt handling (would require applying the chat template ourselves via swift-transformers).
  • No gradient checkpointing, no grad accumulation, no lr_schedule yaml (but MLXOptimizers schedules can be attached to Adam), no dropout in LoRA params (rank/scale/keys only), no LR/peak-mem in progress payload, no numbered checkpoints (single adapterURL overwritten), no wandb, no full fine-tuning.
  • Adapter interchange: Swift uses lora_a/lora_b keys + mlx-lm-style adapter_config.jsonround-trip with Python adapters supported by design (LoRAContainer doc).

# 8.3 Execution-strategy conclusion for TrainingService

Python (mlx_lm.lora via PythonRunner + custom TrainingCallback JSON protocol) is the primary training backend: chat/completions datasets, mask_prompt, grad-checkpoint, DoRA, full FT, numbered checkpoints, LR schedules — all Python-only today. Swift-native QLoRA (LoRATrain + QLoRALinear) is real but dataset- and feature-limited; keep it as a possible fast path later, not the Phase 3 core. Adapters interchange both ways.


# Key Implications for Zyquo MLX

  1. Never parse trainer stdout — format changed between 0.31.3 and main. Pin mlx-lm and drive mlx_lm.lora.run(args, callback) from our own helper script emitting JSON lines from the stable callback dicts.
  2. Resume ≠ true resume upstream. RunStore persists iteration counts and presents warm-start resume honestly; in-process pause/resume can be added by owning the training loop via the Python API.
  3. QLoRA works on quantized bases in both stacks; DoRA is Python-only for MoE.
  4. Depend on ml-explore/mlx-swift-lm (not mlx-swift-examples) for Swift LM libraries.
  5. The hyperparameter form must serialize a YAML config (some keys are config-only), and DatasetService owns train/valid splitting (mlx-lm does not auto-split local JSONL).

# Primary sources

mlx-lm: lora.py · tuner/trainer.py · tuner/datasets.py · tuner/utils.py · tuner/callbacks.py · cli_ui.py · fuse.py · evaluate.py · LORA.md · examples/lora_config.yaml (main + v0.31.3 tags). mlx-swift-lm: Libraries/MLXLLM/LoraTrain.swift · Libraries/MLXLMCommon/Adapters/LoRA/*. mlx-swift-examples: README · Tools/llm-tool/LoraCommands.swift · Applications/LoRATrainingExample. Community memory/speed figures: insiderllm.com (UNVERIFIED).