spb/zyquo-mlx Public MIT

1<!--2  TRAINING-RESEARCH.md3  Zyquo MLX45  Author: Simon-Pierre Boucher6  Mail: contact@spboucher.ai7-->89# Zyquo MLX — Fine-Tuning Research (MLX, Ground Truth 2026-07-30)1011> Phase 0.B research document. Verified against live sources: `ml-explore/mlx-lm`12> main branch (version `0.31.3`, latest PyPI release 0.31.3, 2026-04-22) and13> `ml-explore/mlx-swift-lm` main. All code excerpts were fetched raw from GitHub14> on 2026-07-30. Items not verifiable against a primary source are marked15> *UNVERIFIED*.1617---1819## 1. LoRA / QLoRA Workflow (mlx-lm, Python)2021### 1.1 Entry points2223- CLI: `mlx_lm.lora …` or `python -m mlx_lm lora …`24  (`python -m mlx_lm.lora` prints a deprecation notice but works).25- Training extras: `pip install "mlx-lm[train]"`.26- Minimal run: `mlx_lm.lora --model <hf-repo-or-local-dir> --train --data <dir> --iters 600`.27- YAML config: `mlx_lm.lora -c config.yaml`; **CLI flags override config values**,28  config overrides `CONFIG_DEFAULTS`.29- Python API (what `PythonRunner` scripts will use):30  `mlx_lm.tuner.trainer.{TrainingArgs, train, evaluate}`,31  `mlx_lm.tuner.utils.{linear_to_lora_layers, load_adapters, build_schedule}`,32  `mlx_lm.tuner.datasets.{load_dataset, CacheDataset}`;33  `mlx_lm.lora.run(args, training_callback)` mirrors the CLI programmatically.3435Source: `mlx_lm/lora.py` (main).3637### 1.2 Complete hyperparameter table (exact names + defaults, `CONFIG_DEFAULTS`)3839| CLI flag | Config key | Default | Notes |40|---|---|---|---|41| `--model` | `model` | `"Qwen/Qwen3-0.6b"` | HF repo or local converted dir |42| `--train` | `train` | `False` | |43| `--data` | `data` | `"mlx-community/WikiSQL"` | dir with `{train,valid,test}.jsonl` or HF dataset |44| `--fine-tune-type` | `fine_tune_type` | `"lora"` | `lora` \| `dora` \| `full` |45| `--optimizer` | `optimizer` | `"adam"` | `adam`, `adamw`, `muon`, `sgd`, `adafactor` |46| *(config only)* | `optimizer_config` | `{}` per optimizer | e.g. `adamw: {betas:[0.9,0.98], eps:1e-6, weight_decay:0.05}` |47| `--mask-prompt` | `mask_prompt` | `False` | loss on completion only; chat & completions formats only |48| `--num-layers` | `num_layers` | `16` | `-1` = all layers; LoRA applies to the **last** N layers |49| `--batch-size` | `batch_size` | `4` | |50| `--iters` | `iters` | `1000` | |51| `--val-batches` | `val_batches` | `25` | `-1` = full valid set |52| `--learning-rate` | `learning_rate` | `1e-5` | |53| `--steps-per-report` | `steps_per_report` | `10` | |54| `--steps-per-eval` | `steps_per_eval` | `200` | also evals at iter 1 and final iter |55| `--grad-accumulation-steps` | `grad_accumulation_steps` | `1` | averaged accumulation |56| `--resume-adapter-file` | `resume_adapter_file` | `None` | see §5.3 — warm start only |57| `--adapter-path` | `adapter_path` | `"adapters"` | |58| `--save-every` | `save_every` | `100` | checkpoint cadence |59| `--test` | `test` | `False` | eval on `test.jsonl` |60| `--test-batches` | `test_batches` | `500` | `-1` = full test set |61| `--max-seq-length` | `max_seq_length` | `2048` | longer sequences truncated with warning |62| `--grad-checkpoint` | `grad_checkpoint` | `False` | per-layer `mx.checkpoint` |63| `--clear-cache-threshold` | `clear_cache_threshold` | `0` | e.g. `4GB`; calls `mx.clear_cache()` above it |64| `--report-to` | `report_to` | `None` | `wandb`, `swanlab`, or both |65| `--seed` | `seed` | `0` | `mx.random.seed` + `np.random.seed` |66| `--trust-remote-code` | `trust_remote_code` | `False` | |67| *(config only)* | `lora_parameters` | `{rank: 8, dropout: 0.0, scale: 20.0}` | plus optional `keys: […]`; **not settable via CLI** |68| *(config only)* | `lr_schedule` | `None` | `{name, warmup, warmup_init, arguments}` |69| *(config only)* | `hf_dataset` | — | see §4.3 |7071**`lora_parameters` details:** `rank` (8), `scale` (20.0 — MLX exposes a single72`scale` factor, *not* alpha/rank), `dropout` (0.0), `keys` (target modules;73example yaml uses `["self_attn.q_proj", "self_attn.v_proj"]`). If `keys` is74omitted, `linear_to_lora_layers` adapts **all** Linear/QuantizedLinear/75SwitchLinear/Embedding modules in each targeted layer.76→ Zyquo's hyperparameter form maps: rank, scale, dropout, keys, num-layers,77LR, batch size, iters, seed, max-seq-length, grad-checkpoint, optimizer,78mask-prompt, save-every, steps-per-eval — and must generate a **YAML config**79for `lora_parameters`/`lr_schedule` since those have no CLI flags.8081**`lr_schedule`** (`build_schedule`): `name` = any `mlx.optimizers.schedulers`82function (e.g. `cosine_decay`), `arguments` positional (first = initial LR),83optional linear `warmup` steps with `warmup_init`.84Example: `{name: cosine_decay, warmup: 100, warmup_init: 1e-7, arguments: [1e-5, 1000, 1e-7]}`.8586### 1.3 QLoRA8788No flag: **if `--model` points to a quantized model, training is QLoRA**89(LORA.md). Quantized checkpoints load as `nn.QuantizedLinear`;90`linear_to_lora_layers` wraps them keeping the frozen base quantized. Create a91quantized base with `mlx_lm.convert --hf-path <repo> -q` or use any92`mlx-community/*-4bit` repo.9394### 1.4 DoRA and full9596- **DoRA**: `--fine-tune-type dora` (`tuner/dora.py`: `DoRALinear`,97  `DoRAEmbedding`). MoE `SwitchLinear` layers raise "doesn't support DoRA yet".98- **Full**: `--fine-tune-type full` unfreezes the last `num_layers` layers99  (`-1` = true full fine-tune); `lora_parameters` becomes `None`.100101---102103## 2. Adapters & Fusing104105### 2.1 Contents of `--adapter-path` (verified)106107- `adapter_config.json` — the **entire resolved training config** written at108  start (`save_config(vars(args))`); `load_adapters()` reads `fine_tune_type`,109  `num_layers`, `lora_parameters` from it to rebuild the LoRA layers.110- `adapters.safetensors` — latest weights, overwritten every `save_every`111  iterations and at the final iteration.112- Numbered checkpoints `{it:07d}_adapters.safetensors`113  (e.g. `0000100_adapters.safetensors`) — one per cadence, 7-digit zero-padded.114- Contents = `tree_flatten(model.trainable_parameters())` — for full115  fine-tuning this contains all unfrozen weights (same naming).116117### 2.2 `mlx_lm.fuse` (verified argparse)118119| Flag | Default | Notes |120|---|---|---|121| `--model` | `"mlx_model"` | base model path/repo |122| `--save-path` | `"fused_model"` | |123| `--adapter-path` | `"adapters"` | |124| `--upload-repo` | `None` | push to HF Hub |125| `--dequantize` | off | de-quantize a QLoRA base to fp16 while fusing; strips `quantization` from config |126| `--export-gguf` | off | `llama`/`mixtral`/`mistral` model types only, fp16 |127| `--gguf-path` | `"ggml-model-f16.gguf"` | written inside save-path |128129Mechanics: loads model + adapters, calls `.fuse(dequantize:)` on every capable130module, saves a complete standalone MLX model dir (config.json, sharded131safetensors + index, tokenizer files) — directly loadable afterwards.132(LORA.md mentions `--hf-path` for fuse but it is absent from current argparse —133docs stale; UNVERIFIED if intentional. The pinned 0.31.3 flag is `--dequantize`,134not `--de-quantize`.)135136**⚠️ EMPIRICALLY VERIFIED (2026-07-30, this Mac):** fusing a QLoRA adapter into137a **quantized** base without `--dequantize` re-quantizes the merged weights and138**rounds small LoRA deltas away entirely** — our rank-8 adapter's behavior139vanished from the fused 4-bit model while `--adapter-path` inference kept it,140and a `--dequantize` fuse preserved it verbatim. Re-quantizing the fp16 fused141model to 4-bit wiped the behavior again — for lightly-trained adapters the142deltas are simply below the 4-bit quantization step wherever they land.143Zyquo therefore: (1) **defaults to de-quantize when fusing onto a quantized144base**, (2) recommends **adapter-attached inference** (Python `--adapter-path`145/ Swift `LoRAContainer.load(into:)`) as the lossless default for QLoRA146results, and (3) warns before re-quantizing a fused model that adapter effects147may not survive unless training was substantial. The UI must explain this.148149---150151## 3. Full Fine-Tuning Feasibility152153- Memory with Adam/AdamW: weights + grads + 2 moments, all bf16 ≈154  **8 bytes/param ≈ 4× bf16 weights**, plus activations155  (∝ batch × seq; `--grad-checkpoint` cuts these substantially).156- Realistic tiers (*engineering estimates from the 8 B/param rule — UNVERIFIED,157  no published Apple table exists*):158  - 32 GB: full FT ≤ ~1.5–3B (7B with small `--num-layers` + grad-checkpoint + batch 1)159  - 48 GB (**this Mac**): ~3–4B full, 7B partial160  - 64 GB: ~7B full (tight; SGD/Adafactor shrink optimizer state)161  - 128 GB: ~13–15B full162- Community consensus: full FT > 7B belongs on other hardware — use LoRA/QLoRA.163- Trainer machinery (checkpoints, metrics, resume file) is identical to LoRA;164  the "adapter" file simply holds all trainable weights.165166---167168## 4. Dataset Formats (truth from `mlx_lm/tuner/datasets.py`)169170### 4.1 Directory layout (`--data <dir>`)171172`train.jsonl` (required for `--train`), `valid.jsonl` (**optional** — absent →173"Warning: Validation set not found or empty. Training will proceed without174validation."), `test.jsonl` (required for `--test`). One JSON object per line.175176### 4.2 Auto-detected formats (detection order, first sample)1771781. **completions** — both `prompt` and `completion` keys → `CompletionsDataset`.179   Keys renamable via `prompt_feature`/`completion_feature`. Internally180   converted to a 2-message chat and run through181   `tokenizer.apply_chat_template` — **completions data IS chat-templated**,182   incl. optional per-row `tools`.1832. **chat** — `messages` key (renamable via `chat_feature`) → `ChatDataset`,184   tokenized via `apply_chat_template(messages, tools=d.get("tools"))`.185   A `tools` key (OpenAI function-calling schema) is passed to the template.1863. **text** — `text` key (renamable via `text_feature`) → raw encode, EOS187   appended if missing. `mask_prompt` raises `ValueError` for text datasets.188189### 4.3 HF datasets190191- `--data <hf-dataset-id>` directly if pre-formatted (`train`/`valid`/`test`192  splits; needs `pip install datasets`).193- Config `hf_dataset:` (dict or **list** — concatenated): `path`,194  `train_split` (default `"train[:80%]"`), `valid_split` (default195  `"train[-10%:]"`), `test_split`, feature-name overrides, `config:` dict196  forwarded to `datasets.load_dataset`.197198### 4.4 `mask_prompt` (exact behavior)199200Each sample is `(tokens, offset)`. With `mask_prompt`, offset = token length of201the template applied to all-but-last message (`add_generation_prompt=True` for202completions; for chat, set when last role is `assistant`). Loss mask:203`steps >= offset AND steps <= length` — only completion tokens count. Without204it, offset = 0.205206### 4.5 Batching / truncation (`trainer.py iterate_batches`)207208- Samples sorted by token length, batched, batch order shuffled per epoch;209  dataset must have ≥ `batch_size` examples or `ValueError`.210- Padding to `1 + 32·ceil(maxlen/32)` capped at `max_seq_length`; over-long211  sequences truncated with: `[WARNING] Some sequences are longer than212  {max_seq_length} tokens. The longest sentence {n} will be truncated to213  {max_seq_length}. Consider pre-splitting your data to save memory.`214- Tokenization lazy + memoized (`CacheDataset`).215- **No auto train/valid split for local JSONL** — Zyquo's `DatasetService`216  must produce the split files itself (this is a feature we own).217218---219220## 5. Observability221222### 5.1 ⚠️ Two stdout regimes exist right now — do not scrape stdout223224- **PyPI ≤ 0.31.3** (plain prints, flush=True):225  ```226  Iter {it}: Val loss {val_loss:.3f}, Val took {val_time:.3f}s227  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} GB228  Iter {it}: Saved adapter weights to {adapter_file} and {checkpoint}.229  Saved final weights to {adapter_file}.230  ```231- **main (post-0.31.3)**: `rich`-based `TrainUI` (`mlx_lm/cli_ui.py`) — ANSI232  panels, progress bar, columnar rows; **learning rate and peak memory no233  longer printed** (still in the callback dict).234235**Decision for Zyquo:** never parse trainer stdout. Drive training via a236pinned-version Python helper that registers a custom `TrainingCallback` and237emits a **JSON-lines protocol** on stdout (our own, stable).238239### 5.2 The stable channel: `TrainingCallback` (`tuner/callbacks.py`)240241- `on_train_loss_report({iteration, train_loss, learning_rate,242  iterations_per_second, tokens_per_second, trained_tokens, peak_memory})`243  (peak_memory in GB via `mx.get_peak_memory()/1e9`)244- `on_val_loss_report({iteration, val_loss, val_time})`245- Built-ins: `--report-to wandb|swanlab`. No built-in JSON stdout mode — we246  write our own driver script.247248Other startup lines (main): `Loading pretrained model`, `Loading datasets`,249`Training`, `Trainable parameters: {pct:.3f}% ({M}M/{T}M)`; `--test` prints250`Test loss {:.3f}, Test ppl {:.3f}.`251252### 5.3 Resume semantics (`--resume-adapter-file`)253254`model.load_weights(path, strict=False)` **only** — does NOT restore optimizer255state, LR-schedule position, iteration counter, or RNG. Training restarts at256iter 1 as a **warm start**. Prints `Loading fine-tuned weights from {path}`.257→ Zyquo's `RunStore` must persist its own notion of completed iterations and258present resume honestly (remaining-iters warm start), or drive the Python API259directly to keep optimizer state within one process (pause = in-process, not260cross-process).261262---263264## 6. Memory & Speed (LoRA/QLoRA)265266### 6.1 Official guidance (LORA.md "Memory Issues")267268QLoRA (quantized base); reduce `--batch-size` (4→2→1);269`--grad-accumulation-steps` for effective batch; reduce `--num-layers`270(16→8→4); pre-split long examples; `--grad-checkpoint` ("more helpful for271larger batch sizes or sequence lengths with smaller or quantized models").272Reference datum: Mistral-7B, batch 1, num-layers 4, M1 Max 32 GB →273**~250 tokens/sec** training.274275### 6.2 Community data (*indicative, UNVERIFIED individually*)276277- QLoRA 7–8B ≈ 7 GB peak (fits 16 GB); LoRA fp16 8B ≈ 14 GB; 14B QLoRA ≈ 12 GB /278  LoRA ≈ 24 GB; 70B QLoRA ≈ 45 GB (64 GB+, comfortable at 96–128).279- QLoRA tiers: 8 GB→≤3B, 16 GB→7–8B, 24 GB→8–14B, 32 GB→14B, 48 GB→32B,280  64 GB+→32–70B. LoRA fp16 ≈ one tier down.281- Training throughput: M1 Max ~250 tok/s (matches official), M3 Max ~320,282  M4 Max ~380, M2 Ultra ~475. `--grad-checkpoint` ≈ ~30% slower for large283  activation savings.284- Peak-memory drivers in order: base weights (dominant), activations285  (∝ batch × seq² attention + hidden), trainable-layer count (LoRA adapter286  optimizer state is tiny — rank-8 on 7B ≈ 10–20 MB).287288These feed `MemoryAdvisor`'s training-side verdicts; calibrate in Phase 7.289290---291292## 7. Evaluation293294- **Held-out loss/perplexity**: `mlx_lm.lora --model <m> --adapter-path <a>295  --data <dir> --test [--test-batches N]` → `Test loss {:.3f}, Test ppl {:.3f}.`296  (ppl = exp(loss)). Pass `--adapter-path ""` to test the **base** — that plus297  a run with the adapter is the canonical base-vs-tuned scorecard. Python:298  `tuner.trainer.evaluate(model, dataset, batch_size, num_batches,299  max_seq_length)` → avg loss.300- **`mlx_lm.evaluate`** = lm-evaluation-harness integration301  (`lm_eval.simple_evaluate`). Flags: `--model` (req), `--tasks` (req),302  `--output-dir`, `--batch-size` (16), `--num-shots`, `--max-tokens`,303  `--limit`, `--seed` (123), `--fewshot-as-multiturn`, `--apply-chat-template`,304  `--chat-template-args`, `--temp`/`--top-p`/`--top-k`. **No `--adapter-path`**305  — harness eval of an adapter requires fusing first (or API loading).306- **Qualitative**: `mlx_lm.generate --model <m> --adapter-path <a> --prompt …`307  runs the adapted model directly, no fuse needed → powers the Playground308  base-vs-tuned compare.309310---311312## 8. Swift-Native Training — State Today313314**Repo shift (important):** `MLXLMCommon`, `MLXLLM`, `MLXVLM`, `MLXEmbedders`315moved from `mlx-swift-examples` into **`ml-explore/mlx-swift-lm`** ("all316updates … in the other repository"). `mlx-swift-examples` now hosts only317apps/tools + `MLXMNIST`/`StableDiffusion`, depending on mlx-swift ≥ 0.31.4 and318swift-transformers ≥ 1.3.0. **Zyquo MLX must depend on `mlx-swift-lm`.**319320### 8.1 What exists natively in Swift (verified in mlx-swift-lm main)321322- **`LoRATrain`** (`Libraries/MLXLLM/LoraTrain.swift`): `Parameters`323  (batchSize=4, iterations=1000, stepsPerReport=10, stepsPerEval=100,324  validationBatches=10 [0=all], saveEvery=100, adapterURL);325  `train(model:train:validate:optimizer:loss:tokenizer:parameters:progress:)`326  with `Progress` enum — `.train(iteration, trainingLoss,327  iterationsPerSecond, tokensPerSecond)`, `.validation(…)`, `.save(…)` — and328  `ProgressDisposition .stop/.more` → **built-in cooperative cancellation**;329  `evaluate(…) -> Float`; `saveLoRAWeights(model:url:)`.330- **Adapters** (`Libraries/MLXLMCommon/Adapters/LoRA/`): `LoRAConfiguration`331  (Codable; CodingKeys `num_layers`/`fine_tune_type`/`lora_parameters` —332  **documented as compatible with mlx-lm's `adapter_config.json`**);333  `LoRAContainer.from(directory:)` loads `adapter_config.json` +334  `adapters.safetensors`; `.load(into:)`, `.fuse(with:)`, `.unload(from:)`.335  Layers: `LoRALinear`, **`QLoRALinear`** (→ **QLoRA training on quantized336  bases works natively in Swift**; `fused()` re-quantizes), `DoRALinear`,337  `QDoRALinear`; plus `PEFTAdapter.swift` (HF-PEFT format) and an adapter338  factory registry.339- **CLI reference**: `mlx-swift-examples/Tools/llm-tool/LoraCommands.swift` —340  train/fuse/test/eval subcommands (`--adapter`, `--layers`, `--resume`,341  `--data`, `--learning-rate`, `--batch-size`, `--iterations`,342  `--steps-per-report`, `--steps-per-eval`, `--validation-batches`,343  `--save-every`; fuse: `--de-quantize`).344- **GUI example**: `LoRATrainingExample` — QLoRA on345  `mlx-community/Mistral-7B-v0.1-hf-4bit-mlx`, ~4 GB model memory / ~6 GB346  physical RAM.347348### 8.2 Swift gaps vs Python (verified)349350- Dataset loader: only `{"text": …}` jsonl / plain `.txt` — no chat-template/351  messages/completions/`mask_prompt` handling (would require applying the chat352  template ourselves via swift-transformers).353- No gradient checkpointing, no grad accumulation, no `lr_schedule` yaml (but354  MLXOptimizers schedules can be attached to Adam), no dropout in LoRA params355  (rank/scale/keys only), no LR/peak-mem in progress payload, no numbered356  checkpoints (single adapterURL overwritten), no wandb, **no full fine-tuning**.357- Adapter interchange: Swift uses `lora_a`/`lora_b` keys + mlx-lm-style358  `adapter_config.json` — **round-trip with Python adapters supported by359  design** (`LoRAContainer` doc).360361### 8.3 Execution-strategy conclusion for `TrainingService`362363Python (`mlx_lm.lora` via `PythonRunner` + custom `TrainingCallback` JSON364protocol) is the primary training backend: chat/completions datasets,365`mask_prompt`, grad-checkpoint, DoRA, full FT, numbered checkpoints, LR366schedules — all Python-only today. Swift-native QLoRA (`LoRATrain` +367`QLoRALinear`) is real but dataset- and feature-limited; keep it as a368possible fast path later, not the Phase 3 core. Adapters interchange both ways.369370---371372## Key Implications for Zyquo MLX3733741. **Never parse trainer stdout** — format changed between 0.31.3 and main.375   Pin mlx-lm and drive `mlx_lm.lora.run(args, callback)` from our own helper376   script emitting JSON lines from the stable callback dicts.3772. **Resume ≠ true resume** upstream. `RunStore` persists iteration counts and378   presents warm-start resume honestly; in-process pause/resume can be added by379   owning the training loop via the Python API.3803. QLoRA works on quantized bases in both stacks; DoRA is Python-only for MoE.3814. Depend on `ml-explore/mlx-swift-lm` (not mlx-swift-examples) for Swift LM382   libraries.3835. The hyperparameter form must serialize a YAML config (some keys are384   config-only), and `DatasetService` owns train/valid splitting (mlx-lm does385   not auto-split local JSONL).386387### Primary sources388389`mlx-lm`: lora.py · tuner/trainer.py · tuner/datasets.py · tuner/utils.py ·390tuner/callbacks.py · cli_ui.py · fuse.py · evaluate.py · LORA.md ·391examples/lora_config.yaml (main + v0.31.3 tags).392`mlx-swift-lm`: Libraries/MLXLLM/LoraTrain.swift ·393Libraries/MLXLMCommon/Adapters/LoRA/*.394`mlx-swift-examples`: README · Tools/llm-tool/LoraCommands.swift ·395Applications/LoRATrainingExample. Community memory/speed figures:396insiderllm.com (*UNVERIFIED*).397