mlx-bun 0.0.4 → 0.0.5

package/docs/reference/orpo-quickstart.md

@@ -0,0 +1,91 @@
+# ORPO training quickstart — the full stack, preconfigured
+
+One launcher runs the complete new system — **flash-CCE Metal head** (steel fwd+bwd,
+`[M,V]`-free) + **segmented backward** (gradient-checkpointed layer activations) +
+**prefix-sharing** (one forward over `[prompt; chosen; rejected]`) + LoRA. Each feature
+**falls back and logs** if a precondition isn't met (e.g. a row whose chosen/rejected
+prompts aren't identical → two-forward for that row).
+
+## Command
+
+```bash
+MODEL=/path/to/snapshot DATA=/path/to/datadir bun scripts/train-orpo.ts
+```
+- `DATA` is a directory with `train.jsonl` (+ optional `valid.jsonl`); rows are
+  `{"prompt": "...", "chosen": "...", "rejected": "..."}`.
+- The script **auto-detects e4b/Gemma** and sets its required training env flags
+  (`MLX_BUN_PERF_KERNEL=0`, `MLX_BUN_FUSED_GELU=0`) before loading the model — you don't
+  set them yourself.
+
+### Preconfigured defaults (override via env)
+| env | default | meaning |
+|---|---|---|
+| `SEQ` | e4b 8192 · cpm 4096 | max sequence length |
+| `SEG` | 2 | segment size (layers/segment); the validated e4b@8192 sweet spot |
+| `ITERS` | 100 | iterations |
+| `LR` | 1e-5 | learning rate (cosine + short warmup) |
+| `RANK` / `SCALE` | 16 / 2.0 | LoRA rank (by_bits scaled) / scale |
+| `LAMBDA` | 0.1 | ORPO odds-ratio weight |
+| `FLASH=0` | (on) | use the MLX fused head instead of the flash Metal head |
+| `PREFIX=0` | (on) | disable prefix-sharing (plain two-forward) |
+| `SEGOFF=1` | (off) | disable the segmented backward (more memory) |
+| `ADAPTER` | `~/.cache/mlx-bun/mlx-bun-finetunes/orpo-<model>` | output dir (the finetunes cache, not the repo) |
+
+Example — e4b overnight at 8192:
+```bash
+MODEL=~/.cache/huggingface/hub/models--mlx-community--gemma-4-e4b-it-OptiQ-4bit/snapshots/<snap> \
+DATA=./data/mydpo ITERS=200 bun scripts/train-orpo.ts
+```
+
+## What to expect — performance (measured, M1 Max 32 GB)
+
+### e4b (Gemma4, H=2560, V=262144) — full stack
+| seq | config | peak GB | s/step |
+|---|---|---|---|
+| 8192 | flash + segmented(2) + **prefix-share** (prompt-dominant data) | **13.3** | **~70** |
+| 8192 | flash + segmented(2), no prefix | 16.1 | ~175 |
+| 4096 | flash + segmented(2), no prefix | 11.7 | ~81 |
+| 2048 | flash + segmented(2), no prefix | 10.2 | ~40 |
+| 1024 | flash + segmented(2), no prefix | 8.7 | ~21 |
+
+- **Fits comfortably in 32 GB at 8192** (the historical "e4b OOMs ≥2048" ceiling is gone).
+- **Prefix-sharing is a real win when the prompt dominates** — at 8192 with a long shared
+  prompt it cut memory 16→13 GB *and* step time 175→70 s (prompt encoded once, not twice).
+  On response-dominant data the saving shrinks toward zero (and it falls back per-row when
+  prompts differ); it never hurts correctness.
+- Memory is **linear in sequence length** (~+1.5 GB / 1024 tokens) — segmenting + flash
+  attention keep it linear, not quadratic. The head itself is **flat 0.93 GB** regardless
+  of vocab/seq.
+- An ~8 h run at 8192 ≈ 100–400 steps depending on the prompt/response split.
+
+### MiniCPM5 (1B, H=1536, V=130560) — full stack
+- Fits easily (~1–2 GB at short seq). ~0.4–1 s/step at seq 256; scales with seq.
+- Everything composes (flash / segmented / prefix-share, in any combination).
+
+### The flash-CCE head in isolation (the piece that made large-vocab affordable)
+- Forward 180 ms, backward **754 ms** (e4b) / 224 ms (cpm) — exact (dh 0.40% / 0.28%
+  bf16-class), `[M,V]`-free, peak **0.93 GB flat at M=8192**. The backward is **5× the old
+  hand-rolled kernel** (3687 ms) via the verbatim MLX steel GEMM + fused dequant.
+
+## Fallback behavior (logged)
+- **Prefix-share, per row**: if `chosen`/`rejected` prompts aren't byte-identical →
+  two-forward for that row. The first occurrence logs `prefix-share: row prompt mismatch
+  → two-forward fallback`; the saving is lost only on mismatched rows.
+- **`FLASH=0`** → the MLX fused linear-CE head (`quantizedMatmul` both ways; exact,
+  `[chunk,V]` transient) instead of the Metal kernel.
+- **`SEGOFF=1`** → all layer activations resident (only viable at short seq on e4b).
+- Models other than MiniCPM5 / Gemma4 with `orpoPrefixShared` error out (not silently).
+
+## Optional / experimental
+- The Apple-CCE backward skips (`MLX_BUN_CCE_BWD_BLOCK_EPS=1e-5`, the lossless vocab-block
+  skip; `MLX_BUN_CCE_BWD_FILTER_EPS=1e-5`, the coeff filter) are **off by default** — on
+  the now-fast steel kernel they give little and the coeff filter costs gradient accuracy.
+  The block skip *can* help on long genuinely-peaked real text; A/B it on your data and
+  watch that the loss doesn't move.
+
+## Caveats (honest)
+- The e4b numbers above are from **short probe runs** (2–3 steps); they confirm the
+  pipeline runs, fits memory, and produces finite/decreasing loss — **not** training
+  quality (that's your data + hyperparameters).
+- e4b gradients are **bf16-class** (~1–3% vs an exact reference) — expected for the fused
+  bf16 path; MiniCPM5 is tighter (forward bit-exact, grads ~1%).

package/docs/reference/server-api.md

@@ -30,7 +30,18 @@ Request body (OpenAI chat schema; unknown fields ignored):
   "tools": [ /* OpenAI function tools */ ],
   "tool_choice": "auto",         // "none" disables tools
   "chat_template_kwargs": {      // forwarded to the chat template
-    "enable_thinking": false     // MiniCPM5: <think> channel on/off
+    "enable_thinking": false     // MiniCPM5 / Qwen3.5: <think> channel on/off
+  },
+  "reasoning_effort": "medium",  // "none"|"minimal"|"low"|"medium"|"high"
+                                 // gates enable_thinking on Qwen3.5/MiniCPM5:
+                                 // "none" → off, any other level → on.
+                                 // Overrides chat_template_kwargs.enable_thinking.
+  "hlg": {                       // HLG tone-curve sampling (per request).
+    "enabled": true,             // merged over --hlg-sampling server defaults.
+    "width": 4,                  // logit-width of the tone plateau
+    "shoulder": 4,               // top shoulder width
+    "toe": 6,                    // bottom toe width
+    "pivot_offset": 6            // pivot point offset from top
   },
   "adapter": "id"                // LoRA: "id", stacked "a+b", or "none"
 }
@@ -92,12 +103,15 @@ messages; multi-turn prompt prefixes reuse the KV prompt cache
 automatically.
 
 Tool-call parsing is per model family. Gemma 4 uses its native
-`<|tool_call>`…`<tool_call|>` sentinel tokens. MiniCPM5 emits XML in
-decoded text (`<function name="…"><param name="…">…`, CDATA-wrapped
-values supported); content before the tool markup still streams live,
-only the markup is withheld and converted to `tool_calls`. Argument
-values are decoded against the tool's JSON schema (string-typed params
-stay strings); markup that fails to parse falls back to plain content.
+`<|tool_call>`…`<tool_call|>` sentinel tokens. Qwen3.5 emits decoded
+text wrapped in `<tool_call><function=name><parameter=key>value` XML
+equals-style blocks. MiniCPM5 emits XML in decoded text
+(`<function name="…"><param name="…">…`, CDATA-wrapped values
+supported). For both Qwen3.5 and MiniCPM5 content before the tool
+markup still streams live, only the markup is withheld and converted to
+`tool_calls`. Argument values are decoded against the tool's JSON schema
+(string-typed params stay strings); markup that fails to parse falls
+back to plain content.
 
 ### Errors
 
@@ -191,6 +205,90 @@ OpenAI SDK speak this now). Oracle: optiq responses shim.
 }
 ```
 
+## GET /library
+
+Returns all models found in the local HuggingFace hub cache (via the
+registry scan), each annotated with a fit assessment for this machine.
+Response is cached for 30 seconds (registry scan + config reads; no
+tensor bytes are read).
+
+```jsonc
+{
+  "models": [{
+    "repo_id": "…",
+    "model_type": "gemma3" | "minicpm5" | "qwen3" | …,
+    "size_bytes": 0,
+    "quant_bits": 4,
+    "vision": false,             // has vision sidecar
+    "supported": true,           // recognized model family
+    "serving": false,            // currently loaded in this server
+    "assessment": {              // null if config unreadable
+      "fits": true,
+      "max_safe_context": 8192,
+      "predicted_decode_tps": 0.0
+    }
+  }]
+}
+```
+
+## GET /downloads
+
+Snapshot of the last 5 model downloads (active, completed, or errored)
+initiated via `mlx-bun download` or the web library panel.
+
+```jsonc
+{
+  "downloads": [{
+    "repo_id": "…",
+    "state": "active" | "done" | "error",
+    "current_file": "model.safetensors" | null,
+    "received_bytes": 0,
+    "total_bytes": 0,
+    "files_done": 0,
+    "files_total": 0,
+    "bytes_per_sec": 0,          // rolling ~5 s window
+    "started_at": 1760000000,
+    "finished_at": null,
+    "error": "…"                 // present on state "error"
+  }]
+}
+```
+
+## GET /fit
+
+Fit assessment for the loaded model on this machine, plus a capability
+matrix across Apple SKUs. Used by the status page. Experts bytes come
+from the registry so MoE active-parameter predictions match the serve
+banner and `mlx-bun fit`. When the eval DB has a real measured decode
+rate for this model snapshot, it is included and takes precedence over
+the prediction.
+
+```jsonc
+{
+  "machine": { "chip": "M4 Pro", "ram_bytes": 0, "bandwidth_gbs": 0.0 },
+  "context_tokens": 8192,          // current admission ceiling
+  "typical_context_tokens": 8192,  // min(8192, context_tokens)
+  "typical_decode_tps": 0.0,       // predicted at typical_context_tokens
+  "measured_decode_tps": null,     // real number from eval DB, or null
+  "measured_at": null,             // unix ms of measurement, or null
+  "report": {
+    "fits": true,
+    "weights_bytes": 0,
+    "kv_bytes": 0,
+    "transient_bytes": 0,
+    "total_bytes": 0,
+    "usable_bytes": 0,
+    "max_safe_context": 8192,
+    "predicted_decode_tps": 0.0
+  },
+  "sku_matrix_ctx": 32768,
+  "sku_matrix": [{
+    "sku": "M4 Pro 24 GB", "ram_gb": 24,
+    "fits": true, "max_context": 32768, "decode_tps": 0.0
+  }]
+}
+```
+
 ## Adapters (LoRA hot-swap)
 
 - `GET /v1/adapters` — `{ adapters: [{ id, path, rank, scale, size_bytes, mounted_layers }] }`

package/docs/reference/server-config.md

@@ -25,11 +25,18 @@ the design rationale is in
 | `--prompt-cache` | GB | `2` (≈2e9 bytes) | Byte cap for the prompt (prefix-reuse KV) cache. **Binary GiB (×2³⁰)** on input. `0`/unset ⇒ default; there is no flag to disable the cache outright. |
 | `--batch` | n | `1` (serial) | Max concurrent requests batched through the mlx-lm-parity engine. `>1` switches the **whole server** into bf16 continuous batching — a *mode*, not a load fallback. See [Execution modes](#execution-modes-serial-vs---batch-n). Alias: `--decode-concurrency` (mlx_lm.server drop-in). |
 | `--kv-quant` | `config`\|`off`\|`4`\|`8` | `config` serial / `off`(bf16) under `--batch N` | KV-cache quantization. `config` = per-layer `kv_config.json` (optiq parity); `off` = bf16; `4`/`8` = uniform bits (group 64, start 0). Under `--batch N`, an explicit value routes those requests to the **serial** lane (batched is bf16-only). |
-| `--thinking` | `true`\|`false` | model's own (false for CPM) | Server-wide default for the chat template's `enable_thinking` (MiniCPM5/CPM hybrid reasoning). A request's `chat_template_kwargs.enable_thinking` overrides it. |
+| `--thinking` | `true`\|`false` | model's own (false for CPM) | Server-wide default for the chat template's `enable_thinking` (MiniCPM5/CPM and Qwen3.5 hybrid reasoning). A request's `chat_template_kwargs.enable_thinking` overrides it. |
 | `--temperature` | n ∈ [0,5] | `generation_config.json` | Server-wide sampling default. Per-request `temperature` still wins; the browser chat (sends none) inherits this. |
 | `--top-p` | n ∈ [0,1] | `generation_config.json` | Server-wide top-p default (per-request `top_p` wins). |
 | `--top-k` | n ∈ [0,1e6] | `generation_config.json` | Server-wide top-k default (per-request `top_k` wins). |
-| `--compiled-decode` | on\|off | on | Replay the per-step decode graph in C++ (`MLX_BUN_COMPILED_DECODE`). Bit-exact A/B lever. **Serial lane only** (see note below). |
+| `--no-open` | (bool) | off | Skip the automatic browser open on start. By default an interactive terminal session opens `http://<host>:<port>/#/chat` once the server is ready; pass this flag to suppress it (e.g. headless or non-TTY environments already skip it). |
+| `--hlg-sampling` | `on`\|`off` | off | Piecewise tone-curve (HLG) sampling: rolls off the top-token region, boosts the mids, gentles the tail. The overall gain folds from `--temperature`. See [docs/design/hlg-sampling.md](../design/hlg-sampling.md). |
+| `--hlg-width` | nats | `4` | HLG mid-region half-width (nats). Only meaningful with `--hlg-sampling on`. |
+| `--hlg-shoulder` | nats | `4` | HLG highlight rolloff scale (nats). Only meaningful with `--hlg-sampling on`. |
+| `--hlg-toe` | nats | `6` | HLG shadow rolloff scale (nats). Only meaningful with `--hlg-sampling on`. |
+| `--hlg-pivot-offset` | nats | `6` | HLG pivot point: nats below the top token. Only meaningful with `--hlg-sampling on`. |
+| `--expert-offload` | (bool) | off | **MoE models only.** Serve experts from a page-aligned file mmap (built on first use). Keeps the model out of memory pressure — physical footprint ≈ active params. Ignored with a warning on dense models. Bit-exact with the resident path. |
+| `--compiled-decode` | on\|off | on | Replay the per-step decode graph in C++ (`MLX_BUN_COMPILED_DECODE`). Bit-exact A/B lever. **Serial lane only** (see note below). **Gemma4-dense only** — LoRA, MoE, and non-Gemma4 models (MiniCPM5 / Qwen3.5) run eager; an unsupported step falls back to eager for the rest of that generation. |
 | `--perf-kernel` | on\|off | **on** | Fused quantized-KV decode-SDPA Metal kernel (`MLX_BUN_PERF_KERNEL`), the perf side of the compat A/B. Engages on quantized caches at decode. **Serial lane only.** |
 | `--fused-decode` | on\|off | off | Experimental: tile the quantized decode SDPA (`MLX_BUN_FUSED_DECODE`). **Serial lane only.** |
 | `--fused-sdpa` | on\|off | on | Fused SDPA path for quantized prefill/continuation (inverted env `MLX_BUN_NO_FUSED_SDPA`). **Serial lane only.** |
@@ -40,7 +47,7 @@ winners; flip them to compare. They are set as `MLX_BUN_*` env vars
 before the model loads, so they apply to `mlx-bun pi` too. They affect
 the **serial** decode path (`generate()`); the batched scheduler calls
 `model.forwardHidden` directly and so is unaffected by all of them — see
-[Levers that don't reach the batched lane](#levers-that-dont-reach-the-batched-lane).
+[Levers that don't reach the batched lane](#--batch-n-is-compat-mode--perf-flags-dont-apply-by-design).
 
 > **Note — `--perf-kernel` default.** The code defaults it **on**
 > (`perfKernelEnabled()` returns true unless `MLX_BUN_PERF_KERNEL=0`, and
@@ -129,7 +136,7 @@ concurrently with each other.
 | `--thinking` / `enable_thinking` | ✅ batches (template-render concern, lane-independent) |
 | multi-turn / long prompt | ✅ batches, but **no prompt-cache reuse** (`cached_tokens=0`) |
 
-Both full-attention (CPM) and sliding-window (Gemma) models batch; the
+All three model families — full-attention (CPM), sliding-window (Gemma), and hybrid gated-DeltaNet (Qwen3.5) — batch; the
 scheduler assembles each layer's cache by attention type.
 
 ## Compatibility matrix
@@ -216,7 +223,7 @@ The live config and batch state:
 
 ```jsonc
 {
-  "server":  { "owner": "serve", "model": "...", "started_at": 0 },
+  "server":  { "owner": "serve" | "pi-session" | "embedded", "model": "...", "started_at": 0 },
   "prompt_cache":  { "entries": 0, "bytes": 0, "max_bytes": 2000000000, "hits": 0, "misses": 0 },
   "response_store": { "entries": 0, "bytes": 0, "max_bytes": 33554432, "ttl_ms": 3600000 },
   "kv_quant": {

package/docs/reference/training.md

@@ -0,0 +1,292 @@
+# Training / fine-tuning reference
+
+How to fine-tune a served model with LoRA adapters: the entry points, the
+data formats, every config flag, and the methodology behind the knobs.
+
+mlx-bun trains **LoRA adapters** (the base quantized weights stay frozen),
+supports **SFT**, **DPO**, and **ORPO**, and runs on a single Apple-Silicon GPU. The
+output is a PEFT-compatible adapter you hot-swap into the server — see
+[adapters-end-to-end](../design/adapters-end-to-end.md) for the serving side
+and [segmented-backward-training](../design/segmented-backward-training.md)
+for the long-context memory mechanism.
+
+> Source of truth: the config schema is `FinetuneSubmit` in
+> [`src/train/job.ts`](../../src/train/job.ts); defaults are
+> `DEFAULT_TRAIN_CONFIG` in [`src/train/trainer.ts`](../../src/train/trainer.ts).
+> This doc is generated against those — if they drift, the code wins.
+
+## Entry points
+
+There is **no `mlx-bun train` CLI verb.** Training runs as a subprocess job,
+reachable four ways:
+
+| Path | How | Use when |
+|---|---|---|
+| **Web UI** | `mlx-bun serve`, open `/finetune` — pick model → dataset → hyperparameters → train; watch live train/val loss; merge/export the adapter | Interactive, the default |
+| **HTTP API** | `POST /api/finetune/submit` (job id + SSE events); `POST /api/finetune/inspect-dataset` to probe a file; `POST /api/finetune/merge` to fold an adapter into base weights | Scripted / remote |
+| **Script** | [`scripts/chunk-finetune.ts`](../../scripts/chunk-finetune.ts) — an env-driven wrapper that calls the runner directly | Repeatable CLI runs |
+| **Shell recipe** | [`scripts/ft-e4b-v2.sh`](../../scripts/ft-e4b-v2.sh) `probe`\|`train` — the actual e4b run we use; sets the required env (see [What we actually run](#what-we-actually-run-the-e4b-recipe)) | Reproducing our e4b fine-tune |
+| **ORPO launcher** | [`scripts/train-orpo.ts`](../../scripts/train-orpo.ts) — preconfigured ORPO with the **full stack on by default** (flash-CCE head + segmented backward + prefix-sharing); auto-detects e4b and sets its env; `RESUME=<dir>` warm-starts. See [orpo-quickstart](./orpo-quickstart.md) | Running an ORPO fine-tune |
+| **Library** | `import { finetuneRunner } from "./src/train/job"` and call it with a config + emitter | Embedding training in your own TS |
+
+### Quick start (HTTP API)
+
+```bash
+curl -s localhost:8090/api/finetune/submit -X POST -H 'content-type: application/json' -d '{
+  "model_dir": "/path/to/snapshot",
+  "data_dir":  "/path/to/dataset",        // dir with train.jsonl (+ optional valid.jsonl)
+  "adapter_path": "/path/to/output-adapter",
+  "method": "sft",
+  "rank": 16, "iters": 300, "learning_rate": 2e-4, "max_seq_length": 2048
+}'
+# → { "jobId": "..." }   then stream events from the jobs SSE endpoint
+```
+
+`model_dir`, `data_dir`, and `adapter_path` are **required**; everything else
+falls back to the defaults below.
+
+### Quick start (script)
+
+[`scripts/chunk-finetune.ts`](../../scripts/chunk-finetune.ts) is the worked
+example (MiniCPM5 on the chunking task). It calls `finetuneRunner` directly,
+driven by env vars:
+
+```bash
+# MODEL unset → defaults to the MiniCPM5-1B-OptiQ-4bit snapshot
+DATA=/path/to/chunk ITERS=300 RANK=16 SEQ=2048 SEG=4 \
+  bun scripts/chunk-finetune.ts
+```
+
+Env knobs — note the script applies its **own task-tuned defaults**, which
+differ from the trainer/API defaults in the table below:
+
+| Env | Maps to | Script default |
+|---|---|---|
+| `MODEL` | `model_dir` | MiniCPM5-1B-OptiQ-4bit snapshot (a **path** if set, not a name) |
+| `DATA` | `data_dir` | the lucien chunk dataset path |
+| `SEQ` | `max_seq_length` | `8192` |
+| `ITERS` | `iters` | `2` (probe; use 300 for a real run) |
+| `RANK` | `rank` | `16` |
+| `LR` | `learning_rate` | `1e-5` |
+| `SCALE` | `scale` | `20` |
+| `SEG` | `segment_size` | `0` (off) |
+| `EVAL_EVERY` | `steps_per_eval` | auto from `ITERS` |
+| `ADAPTER` | `adapter_path` | `~/.cache/mlx-bun-finetunes/minicpm5-chunk-seq<SEQ>` |
+| `CKPT` | `save_checkpoints` | on (`CKPT=0` disables) |
+| `GRAD_CKPT` | `grad_checkpoint` | off (`GRAD_CKPT=1` enables) |
+
+The script hard-codes `method=sft`, `batch_size=1`, `steps_per_report=1`, and
+uses the default `ops.sdpa` training attention (set `MLX_BUN_TRAIN_ATTN=flash`
+to override — but flash crashes e4b at multi-K; see Methodology).
+
+## Data formats
+
+Each row of `train.jsonl` (and optional `valid.jsonl`) is auto-detected by its
+keys ([`src/train/dataset.ts`](../../src/train/dataset.ts)):
+
+| Format | Shape | Loss boundary |
+|---|---|---|
+| `messages` | `{"messages": [{"role","content"}, …]}` | response-only — loss on the final turn, prompt = chat-template render of all prior turns |
+| `prompt-completion` | `{"prompt": "...", "completion": "..."}` | loss on the completion only |
+| `text` | `{"text": "..."}` | full-sequence (no prompt mask) |
+| `dpo` *(method=dpo or orpo)* | `{"prompt", "chosen", "rejected"}` | preference loss on chosen vs rejected (the same preference format serves both DPO and ORPO) |
+
+Probe a file before submitting: `POST /api/finetune/inspect-dataset` with
+`{"path": "..."}` returns `{ ok, n_train, n_valid, format }`.
+
+## SFT vs DPO vs ORPO
+
+- **SFT** (`method: "sft"`, default) — supervised fine-tune; response-only
+  cross-entropy. Default LR `2e-4`. For instruction-following, formatting,
+  task adaptation.
+- **DPO** (`method: "dpo"`) — Direct Preference Optimization on
+  chosen/rejected pairs; loss `-log σ(β·((π_c − ref_c) − (π_r − ref_r)))`
+  with reference log-probs computed at LoRA scale 0. Default LR `5e-5`. Tune
+  with `dpo_beta`, `dpo_warmup_iters`, `dpo_lr_schedule`.
+- **ORPO** (`method: "orpo"`) — Odds Ratio Preference Optimization (Hong et
+  al. 2024): a **reference-free** monolithic objective combining the SFT term
+  and a preference term, `L = L_NLL(chosen) + λ·L_OR`, with
+  `L_OR = -log σ(log_odds)` and
+  `log_odds = (ℓ_w − ℓ_r) − (log1mexp(ℓ_w) − log1mexp(ℓ_r))`, where `ℓ` is the
+  **length-normalized** (mean over response tokens) log-prob. Uses the same
+  `{prompt, chosen, rejected}` preference data as DPO but needs **no reference
+  model** (2 forwards/step vs DPO's 4). Default LR `1e-5` (lower than DPO — the
+  loss carries a full SFT NLL term a high LR destabilizes). Tune with
+  `orpo_lambda`, `orpo_warmup_iters`, `orpo_lr_schedule`. Design + optimization
+  roadmap: [orpo-training](../design/orpo-training.md).
+
+The ORPO head + long-context machinery compose into one stack (all B=1), each
+piece independently toggled and each falling back cleanly:
+
+- **`orpo_flash_ce`** — the [M,vocab]-free flash-CCE Metal head (steel GEMM fwd
+  **and** bwd; e4b backward 754 ms, peak 0.93 GB flat at M=8192). Implies the
+  fused head. Falls back to the MLX fused head when the tiling isn't clean.
+- **`segment_size`** — segmented backward (gradient-checkpointed layer
+  activations) for long context; the head is ~free in memory, the activations
+  are the wall.
+- **`orpo_prefix_shared`** — encode the shared prompt once; composes with both
+  the flash head and the segmented backward (MiniCPM5 **and** e4b/Gemma4). Per-row
+  two-forward fallback when chosen/rejected prompts differ.
+- **warm-start** — `scripts/train-orpo.ts RESUME=<adapter-dir>` continues from a
+  checkpoint's weights (optimizer + LR schedule restart).
+
+The one-command way to run all of it is the **ORPO launcher**
+([`scripts/train-orpo.ts`](../../scripts/train-orpo.ts)) — see
+[orpo-quickstart](./orpo-quickstart.md) for the command, the per-model defaults,
+and the measured memory/throughput (e.g. e4b full stack @ 8192 ≈ 13 GB,
+~70 s/step on the M1 Max dev box). Run long jobs detached from your shell
+(`nohup … &`), not as a session-spawned background task.
+
+## Configuration reference
+
+All fields optional except `model_dir` / `data_dir` / `adapter_path`. Defaults
+are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
+
+| Field (API) | Type | Default | Effect |
+|---|---|---|---|
+| `method` | `sft` \| `dpo` \| `orpo` | `sft` | Training objective (see above) |
+| `rank` | int ≥2 | `8` | LoRA rank per adapted linear |
+| `scale` | float >0 | `1.0` | LoRA α (effective update = α·BA) |
+| `rank_scaling` | `constant` \| `by_bits` \| `by_kl` | `by_bits` | Per-layer rank policy (see Methodology) |
+| `target_modules` | string[] | `q,k,v,o,gate,up,down _proj` | Which linears get adapters |
+| `num_layers` | int | `-1` | `-1` = all layers; `N` = last N only |
+| `iters` | int >0 | `100` | Total training steps |
+| `learning_rate` | float >0 | `2e-4` (sft) / `5e-5` (dpo) | AdamW LR |
+| `max_seq_length` | int >0 | `512` | Truncate/pad sequences to this |
+| `batch_size` | int ≥1 | `1` | Rows per step (B=1 is the safe path; B>1 length-sorts + pads to 32) |
+| `grad_accumulation_steps` | int ≥1 | `1` | Accumulate grads over N micro-steps |
+| `seed` | int | `0` | RNG for shuffling + LoRA init |
+| `steps_per_report` | int >0 | `10` | Emit a train-loss metric every N steps |
+| `steps_per_eval` | int >0 | `50` | Eval on `valid.jsonl` every N steps |
+| `weight_decay` | float ≥0 | `0.01` | AdamW weight decay (β = `[0.9, 0.999]`, fixed) |
+| `lora_dropout` | float [0,1) | `0.0` | LoRA-input dropout (PEFT-style), training-only regularizer for small sets. Recompute-safe (mask keyed by step+layer, so segmented/checkpointed backward reproduces it) |
+| `rs_lora` | bool | `false` | rsLoRA — scale the update by α/√rank instead of α, so `rank_scaling` changes capacity not step size. Recorded in the adapter config; inference applies the same per-layer scale |
+| `lora_plus_ratio` | float ≥1 | `1.0` | LoRA+ — LR multiplier for the B leaves (A stays at base LR). >1 speeds the B-driven early learning (B is zero-init). 1 = off. Wired in the ORPO loop |
+| `grad_checkpoint` | bool | `false` | Recompute layer activations in backward (memory↔compute; bit-identical) |
+| `segment_size` | int | `0` (off) | `>0` enables segmented backward — layers per segment (see below) |
+| `save_checkpoints` | bool | `false` | Save every eval-step checkpoint + write `metrics.json` |
+| `dpo_beta` | float >0 | `0.1` | DPO strength (dpo only) |
+| `dpo_warmup_iters` | int ≥0 | `0` | DPO LR warmup (dpo only) |
+| `dpo_lr_schedule` | `constant` \| `cosine` | `cosine` | DPO LR schedule (dpo only) |
+| `orpo_lambda` | float >0 | `0.1` | ORPO λ — weights **only** the odds-ratio term; the SFT-NLL term stays unweighted (orpo only) |
+| `orpo_warmup_iters` | int ≥0 | `0` | ORPO LR warmup (orpo only) |
+| `orpo_lr_schedule` | `constant` \| `cosine` | `cosine` | ORPO LR schedule (orpo only) |
+| `orpo_chunk_size` | int | `0` (off) | Token-chunk the response head, bounding the `[M,vocab]` logits to `[chunk,vocab]` (defaults to 512 when a fused head is on). Exact (orpo only) |
+| `orpo_fused_ce` | bool | `false` | Fused linear-CE head: one CustomVjp with an analytic softmax−onehot backward (MLX `quantizedMatmul` both ways) — no autograd through the head, no retained `[M,vocab]`. Exact; `[chunk,vocab]` transient (orpo only) |
+| `orpo_flash_ce` | bool | `false` | Route the fused head through the **flash-CCE Metal kernel** (verbatim MLX steel GEMM + ORPO epilogue, fwd **and** bwd): neither `[M,vocab]` nor a dequantized head touches HBM → `[M,vocab]`-free, fastest on large vocab (e4b bwd 754 ms, 0.93 GB flat @ M=8192). Implies the fused head. The Apple-CCE coeff filter / blockMax skip are opt-in via `MLX_BUN_CCE_BWD_FILTER_EPS` / `_BLOCK_EPS` (default off — exact). B=1 (orpo only) |
+| `orpo_prefix_shared` | bool | `false` | Shared prompt-prefix: one forward over `[prompt; chosen; rejected]` with a block-sparse mask + block-wise RoPE, so the shared prompt is encoded **once** (a big win when the prompt dominates). Falls back to two-forward per row on prompt mismatch. Composes with the flash head and the segmented backward (MiniCPM5 + e4b). B=1 (orpo only) |
+
+### Environment variables (training)
+
+| Env var | Set for training | Default | Why |
+|---|---|---|---|
+| `MLX_BUN_FUSED_GELU` | **`0` (required for Gemma)** | on | The fused GeGLU is a **CustomKernel with no gradient (vjp)** ([`fused-geglu-kernel.ts`](../../src/model/fused-geglu-kernel.ts)); the Gemma forward uses it ([`gemma4.ts:277`](../../src/model/gemma4.ts)), so a Gemma (e4b/12B/26B) backward fails unless it's off. MiniCPM5 (Llama-arch SwiGLU) never hits it, so the `.ts` script's MiniCPM5 default doesn't need it. |
+| `MLX_BUN_PERF_KERNEL` | **`0` for training** | on | The fused quantized-decode kernel likewise has no vjp. It only fires at decode L=1 (rare in the L>1 training forward), but the e4b recipe sets `0` to be safe. |
+| `MLX_BUN_TRAIN_ATTN` | **leave unset** | unset → `ops.sdpa` | Default `ops.sdpa` **is** mlx's fused flash-attention kernel — the correct, working path. `flash` selects a *different* hand-rolled custom kernel that crashes e4b at multi-K; do not set it. |
+| `MLX_BUN_MEM_LOG` | `1` to profile | off | Print per-step peak/active/cache memory |
+
+> Important: the trainer itself ([`trainer.ts`](../../src/train/trainer.ts))
+> only reads `MLX_BUN_TRAIN_ATTN` and `MLX_BUN_MEM_LOG` — but the **model
+> forward** it runs reads `MLX_BUN_FUSED_GELU` / `MLX_BUN_PERF_KERNEL`. The
+> trainer does **not** disable those itself, so **the caller must** export
+> them. The [e4b recipe](#what-we-actually-run-the-e4b-recipe) does this; if
+> you train a Gemma model by hand, set `MLX_BUN_FUSED_GELU=0` yourself.
+
+## What we actually run (the e4b recipe)
+
+Everything above is the full surface (*what you can do*). In practice the
+fine-tune we run is [`scripts/ft-e4b-v2.sh`](../../scripts/ft-e4b-v2.sh):
+e4b (gemma-4-e4b-it-OptiQ-4bit, pinned snapshot) on the lucien `chunk-v2-500`
+curated set (450 train convs) through the segmented-backward trainer. It wraps
+`chunk-finetune.ts` with the e4b-required env and a two-step workflow:
+
+```bash
+scripts/ft-e4b-v2.sh probe   # 2-iter memory/stability check (~1 min) — RUN FIRST
+scripts/ft-e4b-v2.sh train   # the real run (~900 iters ≈ 2 epochs, batch_size 1)
+ITERS=750 SEQ=4096 SEG=1 scripts/ft-e4b-v2.sh train   # override any knob inline
+```
+
+What the recipe pins (and why it differs from the bare defaults):
+
+| Knob | Recipe value | Why |
+|---|---|---|
+| model | e4b OptiQ-4bit, pinned snapshot `fcdb12d7…` | the validated e4b snapshot |
+| data | `chunk-v2-500` (450 train convs) | the curated chunk set |
+| `SEQ` | `8192` | long context |
+| `SEG` | `4` | **segmented backward**, 4 layers/segment — so 8K-ctx activations fit |
+| `RANK` / `SCALE` / `LR` | `16` / `20` / `1e-5` | task-tuned |
+| `ITERS` | `2` (probe) / `900` (train) | ~2 epochs over 450 examples |
+| `MLX_BUN_PERF_KERNEL` / `MLX_BUN_FUSED_GELU` | `0` / `0` | **required** — the fused kernels have no vjp (see env table) |
+| attention | default `ops.sdpa` | mlx's fused flash kernel; **not** `MLX_BUN_TRAIN_ATTN=flash` (that one crashes e4b) |
+
+The two non-negotiables for e4b: **segmented backward** (`SEG>0`, so the
+long-context activations fit) and the **fused kernels off** (so the backward
+has gradients). Always run `probe` before `train`.
+
+## Methodology
+
+### LoRA
+Adapters attach to the target linears; A is initialized uniform, B is zeros,
+so the adapted model equals the base model at step 0. Only A/B are
+differentiated — base quantized weights are frozen. Default targets are the 7
+attention+MLP projections per block (`q/k/v/o_proj`, `gate/up/down_proj`),
+following Unsloth. See [`src/train/lora-params.ts`](../../src/train/lora-params.ts).
+
+### Rank scaling (`rank_scaling`)
+- `constant` — every target gets `rank`.
+- `by_bits` *(default)* — `rank × (bits / 4)`, clamped ≥2; gives **wider
+  adapters to the more-sensitive (higher-bit) layers** (mixed precision: a
+  4-bit layer gets the base rank, an 8-bit layer 2×). This is optiq's "one
+  sensitivity signal, two optimizations" — the same KL pass that assigns the
+  bits also sets the rank. The per-layer bits are read straight from the loaded
+  model's quant specs (auto-wired; no metadata file needed).
+- `by_kl` — scales by per-layer KL importance, clamped to [0.5×, 2×]; reads the
+  KL values from `optiq_metadata.json` and falls back to `by_bits` if none are
+  recorded.
+
+### Long-context memory: segmented backward vs gradient checkpointing
+At long `max_seq_length`, activation memory dominates. Two levers:
+
+- `grad_checkpoint: true` — recompute each layer's activations during
+  backward. Bit-identical; trades compute for memory.
+- `segment_size: N` — **segmented backward**: run the layer stack forward
+  detaching the residual stream into graph-free boundary leaves every N
+  layers, then backprop segment-by-segment via `mlx_vjp` (cotangent passed
+  directly, *not* a surrogate-loss `value_and_grad`, which leaked). Only one
+  segment's activations live at a time. This is the path to multi-K context;
+  full mechanism, proofs, and measured peaks (e.g. MiniCPM5 10.91→3.29 GB
+  @2048) are in
+  [segmented-backward-training](../design/segmented-backward-training.md).
+
+### Training attention kernel (`MLX_BUN_TRAIN_ATTN`)
+- **default `ops.sdpa`** — mlx's fused SDPA; correct (0.00% vs autograd),
+  O(L²) backward memory. Use this.
+- **`flash`** — opt-in O(L) memory path, but the hand-rolled dK kernel is
+  slow and **crashes e4b at multi-K (≥2K)**; do not use it for e4b LoRA
+  training. Detail in segmented-backward-training §6.
+
+## Outputs
+
+A finished run writes a PEFT-compatible adapter directory:
+
+- `adapters.safetensors` — the `lora_a` / `lora_b` tensors
+- `optiq_lora_config.json` — mlx-bun/optiq adapter metadata (per-layer ranks)
+- `adapter_config.json` — PEFT-compatible config
+
+When `save_checkpoints: true`, each eval step also writes
+`checkpoints/step-<NNNNN>-val<loss>/` and a durable `metrics.json` (config,
+wall seconds, peak GB, final/best train+val loss, full val trajectory).
+
+**Serving the adapter:** hot-swap it into a running server via the adapter
+API and select it per-request — see
+[adapters-end-to-end](../design/adapters-end-to-end.md) and the adapter
+endpoints in [server-api](server-api.md). Or fold it into the base weights
+with `POST /api/finetune/merge`.
+
+## Memory & performance tips
+
+- Start at `batch_size: 1` (the no-padding path); raise only with headroom.
+- OOM → lower `max_seq_length`, set `segment_size` (e.g. 2–4), or reduce
+  `rank` / `num_layers`.
+- Set `MLX_BUN_MEM_LOG=1` to watch per-step peak memory.
+- `grad_accumulation_steps` raises effective batch without the memory cost.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mlx-bun",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "description": "Native MLX inference for Bun on Apple Silicon — local LLM server (OpenAI/Anthropic-compatible) and TypeScript library. No Python.",
   "type": "module",
   "license": "MIT",