LoRA fine-tuning

mlx-optiq ships a LoRA trainer that reads its own per-layer sensitivity assignments and gives sensitive layers proportionally more adapter capacity. Output is PEFT-compatible (adapter_config.json + adapters.safetensors) plus an mlx-optiq sidecar describing the per-layer rank distribution.

The same layers mlx-optiq kept at 8-bit during quantization also get more adapter rank during fine-tuning. The same sensitivity signal drives quantization and LoRA rank.

The basic recipe

# Defaults: all transformer blocks adapted, all 7 trainable linears per
# block (Unsloth-aligned), rank 8 with by_bits sensitivity overlay,
# alpha = rank, mask_prompt enabled, max_seq=512. LR + iterations are
# method-aware: SFT runs 3 epochs at 2e-4 (see "Iterations & learning
# rate" below). Pass --iters for an absolute step count instead.
$ optiq lora train mlx-community/Qwen3.5-4B-OptiQ-4bit \
    --data ./my_training_data \
    -o ./my_adapter

# Show the per-layer rank distribution
$ optiq lora info ./my_adapter

Preset bundles for quick rank selection: --preset small (r=8, α=16), default (r=8, α=8), medium (r=16, α=16), large (r=32, α=32), xl (r=64, α=64), xxl (r=128, α=128). Presets set the BASE rank; with --rank-scaling by_bits (default), per-layer rank still scales up on layers OptiQ kept at higher bits.

Data format

JSONL with either messages (chat format) or prompt/completion pairs. Use one of these formats, not bare text. The text format can't expose a prompt/response boundary, so prompt masking falls through to full-sequence loss and degrades quality on tasks where the base model is already competent.

{"messages": [{"role": "user", "content": "..."},
              {"role": "assistant", "content": "..."}]}
{"prompt": "...", "completion": "..."}

Chat template is applied automatically; do not pre-template the data. mask_prompt is on by default, so loss is computed only on the assistant's response tokens.

Layout on disk:

my_training_data/
├── train.jsonl
└── valid.jsonl   # optional, used for validation loss

SFT and DPO

The same optiq lora train command trains both objectives; pick with --method. SFT (--method sft, the default) is standard supervised fine-tuning on your responses. DPO (--method dpo, Direct Preference Optimization) aligns the model on preferences and is the usual second stage: SFT first to teach the format and task, then DPO to prefer better responses over worse ones.

# DPO on top of an SFT adapter (the standard two-stage flow):
# --mount-adapter starts DPO from your SFT weights; --dpo-beta is the
# KL constraint (default 0.1); LR defaults to 5e-5 for DPO.
$ optiq lora train mlx-community/Qwen3.5-4B-OptiQ-4bit \
--method dpo \
--data ./my_preference_data \
--mount-adapter ./my_sft_adapter \
--dpo-beta 0.1 \
--fused-dpo \                        # chunked logp for >4k context (see below)
--iters 200 \
-o ./my_dpo_adapter

DPO data is JSONL with a chosen and a rejected completion for the same context. Single-turn uses {prompt, chosen, rejected} strings; multi-turn / agentic uses full message lists:

// single-turn
{"prompt": "...", "chosen": "...", "rejected": "..."}
// multi-turn (chosen/rejected are full {messages} lists)
{"chosen": [{"role": "user", ...}, ...],
 "rejected": [{"role": "user", ...}, ...]}

Iterations & learning rate

--iters is optional. Omit it and OptiQ trains by epoch — num_epochs × ceil(examples / batch) — defaulting to 3 epochs for SFT and 1 for DPO (one epoch is a preference nudge; more invites the collapse pathology). Pass --num-epochs N to change the count, or --iters N for an absolute step count. An explicit value always wins.

Learning rate is method-aware: omit --learning-rate and SFT uses 2e-4, DPO uses 5e-5. These defaults resolve inside OptiqLoraConfig, so the CLI, OptiQ Lab, and direct construction all agree — a method="dpo" config no longer silently trains at the SFT rate.

# No --iters: 3 epochs (SFT default) at 2e-4 over the dataset
$ optiq lora train mlx-community/Qwen3.5-4B-OptiQ-4bit \
    --data ./my_training_data -o ./my_adapter

# Two epochs instead of the default three
$ optiq lora train mlx-community/Qwen3.5-4B-OptiQ-4bit \
    --data ./my_training_data --num-epochs 2 -o ./my_adapter

NEFTune noisy-embedding SFT

NEFTune (Jain et al. 2023) adds uniform noise — scaled by alpha / sqrt(seq_len × embed_dim) — to the token embeddings during the SFT forward pass, and nothing at inference. It's a near-free regularizer that stops the adapter from over-memorizing the exact surface form of a small dataset, which measurably improves instruction-following. Pass --neftune-noise-alpha N (the paper suggests 5–15); it's off by default (matching TRL/Unsloth), and SFT only (mirroring TRL, which applies it in the SFTTrainer, not DPO). The noise is gated on training mode, so your validation loss is still measured on clean embeddings.

$ optiq lora train mlx-community/Qwen3.5-4B-OptiQ-4bit \
    --data ./my_training_data -o ./my_adapter \
    --neftune-noise-alpha 5

Early stopping & experiment logging

Both are optional and off by default, and both work the same in the CLI and in OptiQ Lab (the Lab surfaces them through the same config, and its live chart keeps updating alongside them).

Early stopping. Pass --early-stopping-patience N to halt once the validation loss hasn't improved for N consecutive evaluations (evaluations fire every --steps-per-eval steps). It needs a validation set (valid.jsonl); without one it's a no-op with a warning. Use --early-stopping-min-delta D to require an improvement of at least D to count as progress (so tiny wiggles don't keep training alive). When it fires, the best-validation adapter is promoted to the returned adapter — the trailing steps that tripped the counter are worse than the best, so you get the good checkpoint, not the last one. (OptiQ always snapshots the best under <adapter>/best/ regardless; early stopping just also makes it the top-level adapter.)

Experiment logging. Pass --report-to wandb (needs pip install wandb + a WANDB_API_KEY or wandb login) to stream train/val metrics to Weights & Biases; swanlab is also supported, and you can pass several comma-separated. Set the project with --wandb-project. The tracker runs alongside the normal CLI log and the Lab chart, not instead of them.

# Stop after 3 evals with no meaningful val improvement; log to W&B
$ optiq lora train mlx-community/Qwen3.5-4B-OptiQ-4bit \
    --data ./my_training_data -o ./my_adapter \
    --steps-per-eval 50 \
    --early-stopping-patience 3 --early-stopping-min-delta 0.01 \
    --report-to wandb --wandb-project my-optiq-runs

Sensitivity-aware rank scaling

--rank-scaling by_bits (default) gives each layer an adapter rank proportional to its quantization bit-width. With --rank 8:

• Layers mlx-optiq quantized at 4-bit get rank 8.
• Layers mlx-optiq quantized at 8-bit get rank 16.

Head-to-head on a 6-category logical-puzzles reasoning dataset (Qwen3.5-4B-OptiQ-4bit, 1 epoch over 200 training samples, 100-sample test split):

by_bits matches constant rank-16 on accuracy (35% vs 36%, within noise at n=100) using 39% fewer trainable parameters. Versus constant rank-8 at almost matched param budget, by_bits is +8 absolute accuracy points. Full per-category breakdown in the sensitivity-aware LoRA blog post.

Other scaling modes:

# Constant rank (matches Unsloth / PEFT default behaviour)
$ optiq lora train ... --rank-scaling constant

# Scale by raw KL sensitivity (more aggressive than by_bits)
$ optiq lora train ... --rank-scaling by_kl

Training-ceiling map (36 GB Mac)

Empirical sequence-length and peak-memory ceilings at the system-default iogpu.wired_limit_mb=0. Measured under the conservative recipe: num_layers=16 (only the last 16 transformer blocks adapted), target_modules=q_proj,v_proj, rank=8. The current default (num_layers=-1, all 7 target modules) adapts ~3× more LoRA modules so pushes seq-length proportionally lower for the same model on the same machine; drop num_layers or max_seq_length to land within these ceilings if you need the full capacity headroom for very long contexts.

my_adapter/
├── adapter_config.json       # PEFT-compatible config
├── adapters.safetensors      # PEFT-compatible weights
└── optiq_lora_config.json    # mlx-optiq sidecar with per-layer ranks

$ optiq lora info ./my_adapter
# OptiQ LoRA adapter
#   base: mlx-community/Qwen3.5-4B-OptiQ-4bit
#   rank: 8 (scaling: by_bits)
#   scale: 1.0  dropout: 0.0
#   rank distribution: {8: 101, 16: 27} across 128 adapted modules

from mlx_lm import load, generate

model, tok = load(
    "mlx-community/Qwen3.5-9B-OptiQ-4bit",
    adapter_path="./my_adapter",
)
print(generate(model, tok, prompt="...", max_tokens=200))