# 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 --train --data --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 -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 `) `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_template` — **completions data IS chat-templated**, incl. optional per-row `tools`. 2. **chat** — `messages` 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. **text** — `text` key (renamable via `text_feature`) → raw encode, EOS appended if missing. `mask_prompt` raises `ValueError` for text datasets. ### 4.3 HF datasets - `--data ` 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): ``` 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 --adapter-path --data --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 --adapter-path --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/.more` → **built-in cooperative cancellation**; `evaluate(…) -> Float`; `saveLoRAWeights(model:url:)`. - **Adapters** (`Libraries/MLXLMCommon/Adapters/LoRA/`): `LoRAConfiguration` (Codable; CodingKeys `num_layers`/`fine_tune_type`/`lora_parameters` — **documented 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.json` — **round-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*).