mlx-bun 0.0.5 → 0.0.6

package/docs/investigations/orpo-base-uf-experiment-and-directions.md

@@ -0,0 +1,120 @@
+# ORPO from-base on UltraFeedback (mixed-5bpw) — results + research directions
+
+**Status:** experiment complete (2026-06-23). This is the validation run **and** the
+research program it surfaced. Detailed research notes for the session of 2026-06-21…23.
+
+**Companion docs:**
+- [orpo-dynamic-lambda.md](../design/orpo-dynamic-lambda.md) — the adaptive-λ / PID controller (extended below).
+- [orpo-future-enhancements.md](../design/orpo-future-enhancements.md) — SimPO + objective variants.
+- [orpo-uf-testing-handoff.md](orpo-uf-testing-handoff.md) — the **earlier** from-*instruct* UF run (4-bit, seq 2048, no-segment). **Distinct** from this run.
+- [segmented-backward-training.md](../design/segmented-backward-training.md), [steel-flash-cce-handoff.md](steel-flash-cce-handoff.md) — the memory/kernel stack this run exercised.
+
+---
+
+## 1. The run
+
+From-**base** (true base, not instruct) ORPO on UltraFeedback, MiniCPM5-1B.
+
+- **Base:** `openbmb/MiniCPM5-1B-Base-OptiQ-mixed-5bpw` — the **true base**, mixed-precision quant (unblocked by the quantization fix). NOT the instruct model the earlier run used.
+- **Data:** `uf-binarized-chat` — UltraFeedback as `{prompt:[{role,content}], chosen, rejected}`, **60,630 train pairs**. Prompt is a structured message array; chosen/rejected are raw response strings.
+- **Config:** ORPO, **1 epoch (3,790 optimizer steps)**, B=1, **seq 4096**, rank 16 / scale 2, **λ=0.1**, lr **1e-5**, **segmented backward (2 layers/seg)**, **flash-CCE head**, **grad-accum 16** (eff. batch 16), coeff filter ON (`MLX_BUN_CCE_BWD_FILTER_EPS`, real-data ~0.16%).
+- **Template:** the base quant has **no `tokenizer_config.chat_template`**; `ChatTemplate.load` falls back to **`chat_template.jinja`** (which *we* dropped into the quant when building it). It's **ChatML** (`<|im_start|>{role}\n…<|im_end|>`) **+ MiniCPM's XML tool format** (`<function><param>…</param></function>`, CDATA) **+ the hybrid `<think>` block**. We rendered `enableThinking:false` (injects an empty `<think>\n</think>\n\n` before the response).
+- **Recipe fidelity:** lr 1e-5 + λ 0.1 (NOT the saturating 5e-5); **mean-logp verified length-normalized** over response tokens (`responseOnlyLogpMean`: `sum(logp)/M`, M = response-token count) — this is the single most common ORPO parity break and we are on the right side of it, including in the flash-CCE head (parity-validated to 0.21% vs the full-logits mean path).
+
+## 2. Memory validation (the headline systems result)
+
+- **Peak 4.82 GB** — the legitimate high-watermark, driven by the longest example (~2,354 tok chosen+rejected, idx 5582, **+ the prompt**), **not** the 4096 cap (all UF < 4096; 0 examples exceed it — that's *why* 4096 was chosen). My earlier ~4 GB estimate was ~0.5 GB low (slope under-call), not a leak.
+- **Active flat at 0.892 GB** across ~3,300+ steps — **bit-identical early-to-recent, zero drift → NO leak.** Active (the baseline), not peak (example-length-driven), is the leak signal. The ~0.4 MB/step leak we caught earlier would've surfaced as **~1.3 GB of drift** over the run; we saw **0**. So any residual leak is bounded *below detectability*.
+- The whole stack — flash-CCE, steel GEMM, segmented backward, AdamW, grad-accum, the coeff filter + block-skip — ran **dispose-balanced end-to-end** for a full epoch. **A 1B ORPO epoch in < 5 GB, holding flat.** This retires the memory risk for the bigger runs below.
+
+**Predictive memory model (validated):** peak ≈ resting + (2-layer-segment activation) × (longest sequence). Deterministic, bounded; plateaus once the longest example runs. Watch `active_gb` for drift = the leak detector.
+
+## 3. Results — the model
+
+- **Val preference accuracy FLAT ~55.5%** across every eval (500→3000: 55.1, 55.1, 56.3, 55.9, 56.3, 55.9). Plateaued by step 500, **never strengthened**. The 55→56 wobble is 1–3 of 256 examples — inside the ±3.1% SE. n=256.
+- **Val loss** crept down (1.356 → 1.305 — the NLL/format half); **train margin** rose slightly (mean 0.19 → 0.25). A **weak, stable preference** reached early and held.
+- This is the **two-ORPO-halves-decouple** result again: the SFT/NLL half learns the format fast; the preference half barely generalizes on a 1B with subtle UF pairs.
+
+## 4. Vibe-check — 20 questions through base+UF
+
+**Single-turn (fresh context each):**
+- **Good:** simple **code** (prime function correct, sqrt-optimized), **arithmetic** (17×23=391 ✓, 3+5−2=6 ✓), **creative persona** (the pirate plays the role).
+- **Weak/wrong:** **factual recall** ("capital of Australia" → *"there is no specific capital"*; "most recent election" → *"Trump won 2020"*), **translation** ("I love programming" → *"Je love à programmer"*), **identity/privacy/weather** (hedging non-answers), **instruction constraints** ("list exactly three… and nothing else" → preambled anyway).
+
+**Two failure MODES (the actually-useful findings):**
+1. **Multi-turn copying loop** — in a running conversation, once a generic disclaimer enters the history, the model **parrots its own previous assistant turn** for nearly every subsequent question (Q4–9, 15–20 were byte-identical). **Root cause: UF is single-turn preference data** — it never learned to hold a conversation, so the path of least resistance is copy-the-last-turn.
+2. **Hedging-disclaimer boilerplate** ("I'm here to help… I am a language model… respectful and safe… I will never fab[ricate]") — over-learned from UF's RLHF-safety filler.
+
+## 5. Conclusions
+
+1. **VALIDATION SUCCESS.** The chain works end-to-end: quant **base** loads → trains clean a full epoch → **zero leak** → **chats**. ORPO replicated at ~1/60th scale, **on a laptop**, single-stage base→assistant.
+2. **Modest-as-predicted.** base + 1 epoch UF + λ=0.1 = a weak preference (flat 55.5%) and modest quality. We forecast this; it is the expected outcome, **not a miss**.
+3. **The failure modes are DATA artifacts, not method artifacts.** Multi-turn collapse ← single-turn UF; hedging ← UF's disclaimer filler. Both fixable by **changing the data** → this *reinforces* the specialist-curation direction (we control the data). The vibe-check effectively handed us the **curation spec**: add multi-turn examples, filter the disclaimer filler, use our tools in our format.
+4. **CONFOUND CORRECTION (load-bearing).** The gap vs `MiniCPM5-1B-SFT` is **NOT** evidence that single-stage can't replace SFT. We changed **method AND data simultaneously** — 60k UF pairs vs ~400B-token SFT (~1000×). The data difference dominates. The only clean claim is the near-tautological **"60k UF pairs ≠ 400B-token SFT"** (a *data* statement). The **method** question (single-stage vs multi-stage) is **untested** — it requires holding data constant (see §6.5).
+
+## 6. Research directions (crystallized this session)
+
+These form a build: each removes a confound or fixes a measured weakness of the prior.
+
+### 6.1 SimPO, and why it's nearly free for us
+SimPO = reference-free, length-normalized, `−log σ(β·(ℓw−ℓl) − γ)` with a **target margin γ** and **no SFT term** (pure preference; applied *after* SFT). We already compute the length-normalized `ℓw`, `ℓl` (`branchLogpMean`), so SimPO is a **one-term loss swap** → a `--method simpo` + `--beta`/`--gamma`. Natural fit for **SimPO-on-`MiniCPM5-1B-SFT`** (the staged comparison).
+
+### 6.2 ORPO + SimPO's margin (the γ-hybrid) — the fix aimed at *our* result
+For our regime (ℓw≈ℓl), ORPO's odds-ratio ≈ the logp difference already; what SimPO *adds* is the **target margin γ**. ORPO's `−log σ(separation)` gradient **vanishes once chosen ≥ rejected** (separation ≥ 0) → the model plateaus at tiny separation = **our flat margin**. The γ margin **moves the saturation point from 0 to γ**, keeping gradient alive. So: `L = NLL(chosen) + λ·[−log σ(β·(ℓw−ℓl) − γ)]` — ORPO's SFT anchor + SimPO's non-saturating push. Aimed directly at the flat-55.5% we measured.
+
+### 6.3 SFT + SimPO single-stage (the SimPO analog of ORPO)
+"SFT + SimPO at the same time," the way ORPO is "SFT + DPO at the same time." Composes cleanly (both NLL and SimPO push chosen-logp **up** → reinforce). **Watch-out:** SimPO also pushes `logp(rejected)` **down**, and the NLL anchor grounds chosen but **not** rejected — over-aggressive β/γ over-suppresses coherent rejected text → fluency erosion. Tune so the push doesn't outrun the anchor.
+
+### 6.4 Dynamic preference controller (PID/AIMD) — see [orpo-dynamic-lambda.md](../design/orpo-dynamic-lambda.md)
+The flat margin looks like **saturation**, and the *optimal* preference pressure is **non-stationary** (fragile early, robust late) → **no fixed λ is optimal**; a controller tracks the moving target and **beats any fixed value** (a sweep finds the best *constant*; this finds the best *trajectory*). Reframed this session as a **control problem**: NLL = the process variable, the preference weight = the actuator, AIMD/PID = the law. The integral term (the "missing letter") kills chronic steady-state degradation; the derivative (slope) is a **leading** indicator that catches the helping→hurting inflection *before* the level rises. Build **P → PD → PID**, each term earning its place against an observed failure of the prior. Folds the λ/β/γ sweep into one run. Detail + the control-theory framing now in orpo-dynamic-lambda.md §"Control-theory framing (PID/AIMD)".
+
+### 6.5 SFT-as-preference / SPIN comparison — the confound-free method test
+**The interesting one.** Take a **chunk of `UltraData-SFT-2605`** (CPM5's *own* SFT data), make preference pairs **SPIN-style** (chosen = the gold answer, **rejected = on-policy: the model's own generation** for that prompt), train single-stage. Removes the data confound (same data CPM5-SFT saw). Two comparisons:
+- **Method-isolating (rigorous):** plain SFT vs single-stage-preference on the **same chunk** → proves "preference-formatted SFT ≥/= SFT" with no confound. *Why it could be ≥:* ORPO on (gold vs self-gen) = NLL does the SFT **+** the preference term adds "commit to gold over your own output."
+- **"Approach CPM5-SFT" (aspirational, data-gated):** needs a large fraction of 400B tokens → the data gap re-enters; don't let it quietly re-import the confound.
+- Smart first cut: a **single-domain chunk** (knowledge or instruction-following — where base+UF was weakest) — small, legible.
+- **NOT an ORPO-paper replication** — a *different*, arguably more novel comparison.
+
+### 6.6 Scale-up to a standard 7B, on the laptop
+Memory math (validated): 7B resting ~4.4 GB + 2-layer-segment activations → **~8–12 GB peak at seg 2** on UF-length data; **seg 1** halves the activation term (more recompute = slower) for headroom. **Fits 32 GB with room** — the segmented backward was *built* for exactly this (e4b targeted ~10 GB @ 8192; UF < 4096 is the easy end). **Time, not memory, is the cost** (~7× FLOPs/token → days on a laptop). Prefer a **standard dense 7B (Mistral-7B-base** — the ORPO paper's flagship) over e4b for direct literature comparison; it's standard Llama-arch (≈ MiniCPM5's family) so the handler is ~a day. **Strategic flip:** "7B preference-aligned on a 32 GB consumer laptop" *is* the local-to-you thesis at scale — scale-up strengthens the accessibility pitch instead of diluting it.
+
+### 6.7 Diffusion-LM frontier (high-novelty, low *marginal* cost)
+Apply ORPO/SimPO/the controller to a **diffusion language model** (e.g. a Gemma-diffusion). Likely **open territory** (DPO-for-diffusion exists for *images* — Diffusion-DPO — not for diffusion *LMs* + single-stage/dynamic). The diffusion **forward is product-justified anyway** (mlx-optiq has it, mlx-lm doesn't, mlx-vlm does) → once it exists, the preference objective is the same **loss-head swap**. Technical crux: the likelihood is a **stochastic ELBO** (not exact logp), so `ℓw`/`ℓl` are noisy MC estimates → **the dynamic controller's noise-robustness is a *synergy*, not a mismatch.** Length normalization must be rederived for the diffusion loss. Base quality lags AR (24B-diff ≈ 12B-AR), so the contribution is *"does single-stage preference-opt work on diffusion LMs,"* novel regardless.
+
+## 7. Template note (for any future run)
+
+- For **portability / paper-grade** runs, impose a **clean plain-ChatML** template (no MiniCPM XML-tool / `<think>` extensions) at data-prep instead of falling back to the model's custom one. Trivial (drop a clean `chat_template.jinja` / override `renderDpoPrompt`).
+- The model's **native tool format is the XML `<function>/<param>`** one — any tool-use training data must use **that**, not OpenAI JSON `tool_calls`.
+
+## 8. The data — `UltraData-SFT-2605` (the substrate for §6.5)
+
+[openbmb/UltraData-SFT-2605](https://huggingface.co/datasets/openbmb/UltraData-SFT-2605) — the actual SFT data behind `MiniCPM5-1B-SFT`:
+- **15,036,178 examples · 319 GB · ~400B tokens** (200B deep-thinking + 200B hybrid-thinking). **Apache 2.0.**
+- **Single-response (gold only) — NOT preference pairs** → the rejected must be manufactured (SPIN-style on-policy).
+- **Single-turn** (no multi-turn) → won't, by itself, fix the multi-turn copying loop.
+- Domains: math / code / knowledge / instruction-following. Thinking split **6.5M deep / 8.5M non-thinking**. **Chunked by domain** → take a slice, never the full 319 GB. (Multi-GB download is the user's to kick off, not a session's.)
+- The three published models: `MiniCPM5-1B-Base` → `MiniCPM5-1B-SFT` → `MiniCPM5-1B` (instruct, RL+OPD on top of SFT).
+
+## 9. The framing that ties it together
+
+- **What SFT builds vs what preference-opt aligns:** our failure list (multi-turn, instruction discipline, breadth, `<think>`) is ~the SFT stage's job description — but per §5.4 that's mostly a **data** statement (400B curated tokens), not proof the SFT *method* is irreplaceable.
+- **Where single-stage genuinely wins:** specialists (narrow curated data — no foundation breadth needed) and as a **final alignment stage on a good SFT**. NOT a foundation-SFT replacement — *and our own run is the evidence for that boundary, honestly scoped.*
+- **The narrow specialist is the product** (welcome-agent / memory-agent on our closed tool surface); the general UF model was always the **validation**, and it handed us the curation spec for free.
+
+## 10. Choosing what's next — the three lenses (NO decision made yet)
+
+Apply all three to any §6 direction; they point different ways:
+
+1. **Will it work + be useful to US?** (the *product* lens — primary.) Does it improve the specialists we actually ship? Usually gated by a **cheap diagnostic first** — e.g. for the controller: does a *higher fixed λ* or the **γ-margin** even unstick the flat preference? If a fixed push does nothing, *no controller can.* Run the cheap fixed-knob version before the fancy one.
+2. **Interesting paper?** (read / replicated / adds to the ORPO-style discourse.) Strong hooks (control-theory/PID, SFT-as-preference, diffusion-LM preference-opt), but it's a **crowded, competitive space** → the result has to be the non-stationarity / confound-free kind to be more than "auto-tuning," and credibility likely needs **scale** (7B+, standard model, standard benches).
+3. **Fundable / by whom?** The *technique* funds itself via the **product** (cheaper specialist training), independent of any paper. As a paper/grant: compute-cost-sensitive trainers + the local-AI community + project credibility. It's a *method* contribution, not a product.
+
+**The trap + the rule:** don't let the *paper* lens drive the work. **Let the product need drive it** (lens 1), run the cheap diagnostics first, and treat the **paper as a harvest** if the result turns out strong — the experiments that make it a good feature are the same ones that'd make it a paper, so the lens-2/3 decision can wait until lens 1 reports.
+
+**Rough risk / novelty / cost of the §6 directions (for when we pick):**
+- **§6.2 γ-hybrid / higher-λ diagnostic** — lowest cost, directly targets the measured flat preference. *The cheap first cut.*
+- **§6.1/6.3 SimPO-on-SFT / SFT+SimPO** — low (loss-swap), needs the SFT checkpoint pulled.
+- **§6.4 PID controller** — medium (build + tune), highest method-novelty of the AR options.
+- **§6.5 SPIN-on-UltraData** — medium (chunk + on-policy generation), the **confound-free method test**, distinct from ORPO replication.
+- **§6.6 7B-Mistral-on-laptop** — ~a day handler + days of compute; the credibility + thesis-at-scale move.
+- **§6.7 diffusion** — highest novelty, lift mostly paid by the product (the forward exists anyway); the ambitious swing.

package/docs/investigations/orpo-flash-cce-pin-leak.md

@@ -0,0 +1,167 @@
+# flash-CCE host-buffer pin leak — post-mortem + L3 verification approach
+
+**Date:** 2026-06-21 · **Status:** fixed (`fix(train): stop host-buffer pin leak…`,
+commit `2592959`) + hardened (`harden(train): anchor ptr()'d FFI buffers…`, `6491696`).
+
+## Symptom
+
+ORPO LoRA training (MiniCPM5-1B, flash-CCE head + prefix-share + segmented
+backward) crashed **natively** (segfault / "panic: A C++ exception occurred",
+SIGTRAP — *not* an MLX-reported error) after a few hundred steps:
+
+| config | crashed at |
+|---|---|
+| accum-4 + seg-2 (most ops/step) | step **363** |
+| accum-4 + no-seg | step **783** |
+| accum-1 + no-seg (fewest ops/step) | step **≥915** |
+
+Tells: **deterministic** (same step across runs), **monotonic in ops/step**
+(more kernel ops per step → earlier crash), and MLX `active` memory **flat**
+(~1.3 GB) the whole time. Crucially, **it ran the entire training with zero
+crashes on a 32 GB M1 Max**, and crashed on a 24 GB M4 Pro.
+
+## Wrong turns (recorded so we don't repeat them)
+
+1. **Agent sandbox** — the first crash was on a session-launched run, so the
+   sandbox looked guilty. Ruled out: a user-launched run crashed at the *same*
+   step 783.
+2. **GPU-allocator fragmentation / OOM** — `peak` sat at 16.7 GB on a 24 GB
+   box. Ruled out: with segmentation `peak` was only 2.9 GB and it crashed
+   **earlier** (363), and the same 16.7 GB peak ran fine post-fix. Not memory.
+3. **prefix-share fallback / templated-data seam** — a real new behavior (the
+   `</think>\n\n` seam trips the byte-identical prompt check → two-forward
+   fallback) but not the cause.
+4. **grad-accum / grad-clip (the new settings this session)** — the right
+   *thread* (look at what changed), wrong leaf. Both were exonerated.
+
+## Root cause
+
+`u32()` in `flash-cce.ts` (and the twin in `model/flash-attention.ts`) built the
+tiny `shape` / `targets` kernel-argument arrays with **`MlxArray.fromView`** —
+the **zero-copy pinning** path. `fromView` inserts the host `Uint8Array` into a
+module-level `pinned` Map and hands MLX a raw pointer with an **async**
+`unpinCallback` ("mlx may release buffers from its eval thread"). These args feed
+a **lazy** kernel and are `dispose()`d in a `finally` *before the kernel
+evaluates* — so the pin is never reclaimed in time. The pins **leaked**.
+
+**Proof (the smoking gun).** Instrumenting `pinnedBufferCount()` per step on the
+buggy code:
+
+```
+step 1: pinned 32   active 1.087GB
+step 2: pinned 64   active 1.088GB
+...
+step 32: pinned 1024  active 1.099GB     →  +32 per step, never released
+```
+
+Perfectly linear `+32/step`, `active` dead flat. The leaked buffers live in the
+JS heap + a threadsafe-JSCallback dtor queue, **outside MLX's `active`
+accounting** — which is exactly why every symptom held: per-op accumulation (more
+kernel calls/step → faster growth → earlier crash), `active` flat, native crash
+(host allocator / callback queue corruption), and **machine-dependence** (a
+24 GB box reuses the freed-but-pinned buffer sooner → a latent use-after-free
+becomes a fatal read; 32 GB had the headroom to never hit it).
+
+## Fix
+
+`u32()` → **`fromBytesCopy`** (`mlx_array_new_data`, which **copies**). These
+few-int arrays feed a lazy kernel and are disposed before it runs, so they must
+be **MLX-owned**, not host-pinned. A few-int copy is free.
+
+**Verification:** `pinned` goes `+32/step → flat 0`; the previously-fatal
+seg-2/accum-4 config ran to **step 548** (past 363) clean; numerics unchanged
+(copy vs pin = same bytes).
+
+## Hardening (defense-in-depth)
+
+- **`MetalKernel.apply` keepAlive** — anchor every `ptr()`'d JS buffer (output
+  shapes, template-arg names, the input-handle array, the result slot) in a
+  `keepAlive` list referenced after the FFI calls. `bun:ffi` `ptr()` does not
+  retain its argument; this closes the same lifetime class. Pure-additive.
+- Added `pinned: pinnedBufferCount()` to the per-step train metric as a permanent
+  **canary** for this bug class (visible in `train-watch`).
+
+## Lesson for the codebase
+
+`fromView` (zero-copy pin) is only safe for buffers that **outlive the MLX
+array** (mmap'd weights — `fromPointer`/`fromView` with process-lifetime memory).
+For any **transient** host buffer handed to a **lazy** op and disposed before
+eval, use the **copying** constructor (`fromBytesCopy`/`fromInt32`/`fromFloat32`).
+Grep for `fromView` in hot paths when adding kernels.
+
+---
+
+# L3 verification approach (no oracle exists)
+
+This bug also clarified how we *should* verify the flash-CCE ORPO head, because
+"bf16-class error" was the wrong framing.
+
+## The parity tiers
+
+- **L1** — bit-for-bit drop-in for **mlx-lm** (proven: 8192/8192 greedy tokens
+  identical).
+- **L2** — bit-for-bit drop-in for **mlx-optiq**.
+- **L3** — mlx-bun originals (the flash-CCE ORPO training head). **Neither
+  mlx-lm nor optiq ships ORPO+CCE training — there is no oracle to be bit-exact
+  against.**
+
+So the "0.28% dh" figure is **not an error rate or a degradation.** It is the
+**floating-point reassociation gap between two *correct* implementations** — the
+flash kernel (tiled, online-softmax, fp32) vs the full-logits MLX head (a
+different reduction order) — and the full-logits head is itself only a *proxy*,
+not an oracle. Two correct reductions in different orders cannot be bit-identical;
+0.28% is the inherent floor, not slop.
+
+**Correctness in L3 is proven by what does NOT need a reference implementation:**
+
+1. **Finite-difference gradient check** (`flash-fd-check.ts`) — the kernel's `dh`
+   vs numerical ground truth to ~4 decimals. This is the real proof: it says the
+   kernel computes the *right gradient*, independent of any other implementation.
+2. **By-hand math audit** of the dequant, online-LSE recurrence, `coeff =
+   onehot − softmax`, dh accumulation, and the `1/M` cotangent (done; clean).
+3. **Agreement with the full-logits proxy** to fp-reassociation precision
+   (logp 0.21%, dh 0.28%) — corroboration, not the proof.
+
+## The coeff filter — measure on REAL data, not synthetic
+
+`flash-cce-parity.ts` uses a **synthetic random hidden**, which gives a
+**flat-ish** softmax (high entropy) → gradient mass spread across many tokens →
+filtering the near-zero coeffs costs accuracy (the recorded 0.66%→2.7%). **That
+is the filter's worst case and not what it's for.** On real trained-model
+outputs the next-token distribution is **sharply peaked** (low entropy):
+`coeff_v = c·(onehot_v − softmax_v)`, and for almost every non-target token
+`softmax_v ≈ 0`, so its coeff ≈ 0 and dropping it costs ~nothing — the gradient
+mass is in the target + the handful of high-softmax tokens. **The filter is
+designed exactly for the peaky real-data regime; the synthetic benchmark
+maximizes its apparent cost.**
+
+**MEASURED (2026-06-22, `scripts/experiments/flash-cce-filter-realdata.ts`) — the
+hypothesis holds.** On 8629 real response tokens from real examples (real model
+forward, **real next-token targets**; mean P(true token) = 53.9% → genuinely
+peaked), the filter's dh error vs exact:
+
+| eps | dh error | backward speedup |
+|---|---|---|
+| 1e-6 | 0.024% | 1.15× |
+| **1e-5** | **0.162%** (below the 0.33% bf16 floor → indistinguishable) | **1.35×** |
+| 1e-4 | 0.681% | 1.47× |
+| 1e-3 | 1.943% | 1.52× |
+
+The synthetic `REAL=1` test (random targets) showed **21.4% @ 1e-4** — a ~30×
+overstatement, an artifact of the onehot landing on a ~0-softmax random token.
+**Verdict: enable the filter at eps ~1e-5** — free speed (error under the
+reassociation floor) on real data. Re-measure as the model trains (peakedness
+shifts), and always with REAL targets, never synthetic.
+
+## Standing L3 gates (what to run, since there's no oracle)
+
+Per the project's parity-as-correctness model (bit-exact where an ancestor
+overlaps; KL + quality + benchmarks where it doesn't):
+
+1. **Filter-on-real-data** — ✅ DONE (above): near-free at eps ~1e-5 on real
+   targets (0.16% dh, 1.35× backward). `flash-cce-filter-realdata.ts`.
+2. **Teacher-forced grad fidelity** — flash head vs full-logits head `dh`
+   (gradient cosine + relnorm) on real data, as the standing L3 regression gate
+   (replaces the synthetic-data parity number for *quality* claims).
+3. **End-to-end quality eval** — the completed training run's downstream task
+   quality (the real proof that L3 trains correctly).

package/docs/investigations/orpo-uf-testing-handoff.md

@@ -0,0 +1,126 @@
+# ORPO-on-UF — results + testing-phase handoff
+
+**Status:** training done (2026-06-20); **testing phase = next.** This doc is the
+"where we are + what to run" so anyone (or a fresh session) can execute the measurements.
+Companion: [trainer-validation-experiment.md](trainer-validation-experiment.md) (the
+validation that got us here).
+
+## The run that just finished
+
+From-**base** ORPO on UltraFeedback, MiniCPM5-1B.
+- **Base:** `mlx-community/MiniCPM5-1B-OptiQ-4bit` (the full post-trained model; optiq quant).
+- **Data:** `~/.cache/mlx-bun/eval-data/uf-cpm5` — 14,691 train / 148 val, `{prompt,chosen,rejected}`.
+- **Config:** ORPO, **1 epoch (14,691 iters)**, B=1, seq 2048, rank 16, **λ=0.1**, cosine LR
+  1e-5, **no-segment**. From base (no warm-start).
+- **Command:** `SEQ=2048 EVAL_EVERY=500 RANK=16 SCHED=cosine LAMBDA=0.1
+  ADAPTER_DIR=./adapters/cpm5-uf-frombase bun scripts/experiments/pref-control.ts orpo
+  ~/.cache/mlx-bun/eval-data/uf-cpm5 uf-frombase 14691 1e-5`
+
+## Key result — the two ORPO halves decoupled
+
+| half | metric | best | at step | final (14690) |
+|---|---|---|---|---|
+| **SFT** | loss (NLL+λOR) ↓ | **1.5057** | **14500** (the end) | 1.506 |
+| **DPO** | margin ↑ | **0.143** | **6000** | 0.127 |
+| — | accuracy | **81/148** | ~5500 | 80/148 |
+
+- **SFT improved across the *whole* run** (loss bottomed at the end) → **final adapter = best SFT.**
+- **DPO peaked at ~6000 and faded** → **step ~6000 = best DPO**, slightly worse SFT (loss 1.5077).
+- **No single checkpoint has both.** Honest summary: a **small, early-peaking preference
+  signal** (~+6 questions, chance 74 → peak 81/148) that **drifted down in the tail**, while
+  SFT kept inching to its ~1.506 floor. Whether even that small lift **converts** on an
+  independent test is the open question the testing phase answers.
+
+## Artifacts + baselines
+
+- **Final adapter:** `./adapters/cpm5-uf-frombase/` (mountable: `eval.ts --adapter`).
+- **Trajectory:** `runs/uf-frombase.{json,csv}`; **plot:** `runs/uf-frombase.png` via
+  `scripts/experiments/plot-runs.py` (panels: loss=val+train-EMA; questions-correct/148;
+  margin — each zoomed to its band).
+- **x1 baseline** (base CPM5, full capability suite): **28.69** (in the eval DB; needs the
+  MMLU-argmax + HumanEval-python fixes — see Gotchas). optiq published is 30.28.
+- **⚠ No intermediate checkpoints this run** — `pref-control.ts` lacked `saveCheckpoints`
+  (now **fixed**; future runs checkpoint every val eval + track best-margin). The ~6k peak
+  is recoverable by **deterministic replay** (same seed/order → same adapter).
+
+## Testing plan
+
+**PRIMARY — the headline (READY TO RUN):**
+- **A. final adapter vs base** on x2 (capability) + Exp 2 (judge win-rate). Answers "did a
+  full epoch of ORPO-on-UF move a 1B, on an *independent* test?"
+
+**SWEEP / SYNTHESIS:**
+- **B. recover the ~6k peak** — re-run from base to ~7k (checkpoints now on), grab the
+  step-6000 (and neighbors, pick best-val). The early-stop peak arm.
+- **C. four-way** — base / 6k-peak / final / final+6k on x2 + Exp 2 → traces the
+  early-stopping curve on an independent test.
+- **D. DPO-boost (the day's synthesis) — keep best SFT, re-pressure the DPO.** Warm-start
+  from the **final** adapter (best SFT, NLL at floor), re-apply the **DPO-helpful examples
+  with HIGH λ** (or DPO-only). = dynamic-λ (ramp preference at SFT-saturation) + data-curation
+  in one move. Tests the fork: DPO stalled because of a **capacity ceiling** (won't help) or
+  **dilution / λ-too-gentle** (margin recovers past 0.143 while loss holds ~1.506 → better on
+  *both* axes than any single checkpoint).
+- **E. data-valuation control** — `final + (0–6000)` vs `final + (6000–14691)`, same size.
+  First-half wins → those examples carry the DPO signal; tie → it's "more preference training,"
+  not *which* examples. (We control which pairs go in, so this generalizes to **curating** the
+  data with the model as the probe.)
+
+## Commands
+
+`ORC=/Users/joshrossi/Code/mlx-lm-example/.venv/bin/python` (HumanEval interpreter, this machine).
+
+```bash
+# A — x2 capability (final vs base; base x1=28.69 already recorded)
+MLX_BUN_ORACLE_PYTHON=$ORC bun scripts/eval.ts capability --candidate MiniCPM5 --adapter ./adapters/cpm5-uf-frombase
+# A — Exp 2 (gen base+trained responses, then judge head-to-head)
+bun scripts/experiments/prefbench.ts gen mtbench ./adapters/cpm5-uf-frombase mtbench-uf   # GPU
+bun scripts/experiments/prefbench.ts judge mtbench-uf                                      # codex
+# B — replay base→7k to recover the ~6k checkpoint (now checkpoints)
+SEQ=2048 EVAL_EVERY=500 RANK=16 SCHED=cosine LAMBDA=0.1 ADAPTER_DIR=./adapters/cpm5-uf-replay \
+  bun scripts/experiments/pref-control.ts orpo ~/.cache/mlx-bun/eval-data/uf-cpm5 uf-replay 7000 1e-5
+```
+
+## Harness gaps to close for D / E
+
+`pref-control.ts` does **not yet** support (needed for the DPO-boost + control arms):
+1. **Warm-start from an adapter** (`--resume` / load a non-zero adapter before training) —
+   the trainer has `warm_start_adapter` (see `trainer.ts`); just not wired into pref-control.
+2. **Data-subset selection** (feed a chosen subset / order of pairs, e.g. rows 0–6000) —
+   pref-control feeds the whole `train.jsonl`; needs a row-range / id-list option.
+3. λ is already env-controllable (`LAMBDA=`), so high-λ is ready once 1+2 exist.
+
+## Validated infrastructure (all green — see trainer-validation-experiment.md)
+
+- **Trainer:** positive controls pass — DPO & ORPO drive prefer-uppercase val acc → 1.0;
+  grads match autograd. `scripts/experiments/pref-control.ts`.
+- **Capability harness:** reproduces optiq across 6 tasks after two fixes (below).
+  `eval.ts capability --adapter` wired.
+- **Judge:** `gpt-5.4-mini` via `codex exec`, **keyless**, structured, order-swapped,
+  3/3 positive control. `scripts/experiments/llm-judge.ts`.
+- **Exp 2 harness:** `scripts/experiments/prefbench.ts` (gen/judge); prompts downloaded
+  (`mt_bench.jsonl` 80, `alpaca_eval.json` 805).
+- **MMLU frozen set:** optiq's exact 969 questions + dev exemplars captured
+  (`mmlu_optiq_frozen.jsonl` / `mmlu_optiq_dev.jsonl`) for parity work.
+
+## Gotchas / load-bearing facts
+
+- **`AdapterManager.mount` affects the *serving* forward** (eval generation, chat) **but NOT
+  the *training* `branchLogpMean`.** Measure training metrics via the trainLora emit, not mount.
+- **Eval env:** `MLX_BUN_ORACLE_PYTHON=$ORC` (HumanEval code-exec — wrong default path on this
+  machine); evals default to **non-thinking** (optiq parity; `MLX_BUN_EVAL_THINK=1` to re-enable).
+- **MMLU parity:** our number (47.7 old / ~50 argmax) vs optiq 52.4 is **sampling + forward
+  parity**, not noise (both deterministic argmax). Use the frozen set; remaining gap = our
+  forward vs `mlx_lm` (logit-diff test on the frozen prompts — **queued, GPU**).
+- **Determinism:** same seed/order → same adapter. Replay recovers any checkpoint.
+- **Server adapter API** (`POST /v1/adapters` + per-request `"adapter":"id"`) is the clean way
+  to compare *many* adapters via chat; **raw eval (`--adapter`) is required for capability**
+  (MMLU is logit-argmax, not chat generation).
+
+## Recommended order
+
+1. **A** (final vs base) — the headline, READY. Run x2, then Exp 2.
+2. **B** — replay to recover the 6k peak (cheap, ~45 min).
+3. **C** — eval base/6k/final on the suite (four-way needs D for the +6k arm).
+4. Close the **pref-control harness gaps** (warm-start + subset), then **D** (DPO-boost) and
+   **E** (data-valuation control) — the synthesis experiments.
+5. (Parallel/queued) **MMLU forward-parity** logit-diff to nail x1 ≈ 30.28.

package/docs/investigations/steel-flash-cce-handoff.md

@@ -6,6 +6,64 @@ coeff + dh) on top — to get **MLX-GEMM speed AND the `[M,V]`-free memory prope
 in one kernel. MLX ships no fused "quantized-matmul + softcap + response-logp" op, so
 we build the composition; the heavy GEMM is theirs, verbatim.
 
+## HANDOFF — current state (2026-06-19): SHIPPED in v0.0.5
+
+The whole ORPO training stack is **merged (PR #16) and released as v0.0.5** — `npm
+i mlx-bun` / `brew install joshuarossi/tap/mlx-bun` (tap formula at 0.0.5, sha verified).
+What shipped: the flash-CCE head (fwd+bwd), prefix-sharing, segmented composition, warm-start,
+the `scripts/train-orpo.ts` launcher, and the IFEval/UltraFeedback eval pieces. Release notes:
+GitHub release v0.0.5. Supported scope = **OptiQ-quantized MiniCPM5 + Gemma-4 (e4b/12B/26B)**.
+
+**Post-merge PR-review fixes (CodeRabbit, commits 890f380 / 63759fa) — read these, they
+caught real bugs:**
+- 🔴 **Prefix-shared head-primal UAF** (`branchLogpMeanGathered`/`orpoLossPrefixShared`): the
+  flash/fused-head CustomVjp PRIMAL (`hResp`, and `h`) was disposed *before* the lazy backward
+  recomputes from it → use-after-free (same class as the segmented `headSink` one). Fix: push
+  primals into the sink (freed post-eval) on the fused/flash path. **Tests passed before only
+  by lazy-eval timing luck** — this was latent. Watch for this pattern anywhere a CustomVjp's
+  primal is freed inline.
+- **Segmented "flash implies fused"** (`trainer.ts` `segFusedChunk`): `orpoFlashCe` alone left
+  the segmented non-prefix fallback on the full `[M,vocab]` head → e4b memory blowup. Now
+  `orpoFusedCe || orpoFlashCe` routes through the bounded chunked head.
+- Warm-start made atomic (validate all targets before disposing any), job payload exposes
+  `warm_start_adapter`/`mlp_split`, script GPU inputs use `fromBytesCopy` (page-align footgun),
+  env-precondition enforcement + flag validation in the probes.
+
+**Warm-start** (`warmStartFromAdapter` / launcher `RESUME=<adapter-dir>`): continue from a
+checkpoint's LoRA weights (optimizer + LR schedule restart). Insurance for interrupted runs.
+
+**The training dress-rehearsal (honest):** a CPM5 UltraFeedback ORPO run (full stack, lr 5e-5)
+reached ~step 4820, **val 1.66 → ~1.50, then plateaued** — expected: UF is open-ended, its loss
+floor is ~1.2-1.5, NOT the 0.3-0.4 you get on a narrow task. **IFEval moved 22.5% → 22.5%
+(flat)** — a real null result: general preference data + a partial run is the wrong lever for
+IFEval. The pipeline learns (val dropped); the metric just isn't the right one. Checkpoints
+every 200 under `./adapters/cpm5-uf-8h/checkpoints/` (best-val ~`step-04200-val1.5008`).
+
+**Gotchas the runs taught us:**
+- **Agent-spawned background runs get REAPED** (~47 min observed; not a crash/OOM — confirmed
+  via logs + `pmset` showed no sleep). **Run long training detached from your OWN shell**
+  (`nohup … &`), never as a session background task.
+- e4b training **requires** `MLX_BUN_PERF_KERNEL=0` + `MLX_BUN_FUSED_GELU=0` (the launcher
+  auto-sets them; the e4b parity probe now enforces them).
+- Adapters live in **`~/.cache/mlx-bun/mlx-bun-finetunes/`** (launcher default), never the repo.
+
+**THE next goal — the chunk segmenter (load-bearing).** UF was the rehearsal; the real run is
+**distilling Opus/GPT-5.5 conversation-segmentation into a local model** (data at
+`~/Code/lucien/benchmark/finetune/chunk-v3/dpo/orpo-curated-*.fixed.jsonl`, `{prompt:[sys,user],
+chosen, rejected}` — already supported by the loader). Score by **boundary/label accuracy vs
+the gold (chunk-eval), NOT val loss** (loss is dominated by low-entropy JSON/UUID tokens). This
+is what localizes the Lucien synthesis pipeline's `chunk-recent` stage (the whole "100% local
+memory" promise). ORPO is ideal here: chosen/rejected share the conversation + are both
+well-formed JSON, so the contrast isolates the segmentation decision, and the gold makes
+pair-ranking + eval objective. The stack is built for it (prefix-share the shared 8k
+conversation once, segment the long context, flash the big vocab). See memory
+[[training-tracks-are-appliance-components]]. After that: the e4b overnight, and the welcome
+assistant (a UF-tuned greeter — system prompt carries the mlx-bun knowledge, not training).
+
+**Run commands:** see [orpo-quickstart](../reference/orpo-quickstart.md) +
+[training.md](../reference/training.md). The full stack is `scripts/train-orpo.ts` (defaults on);
+resume with `RESUME=<ckpt-dir>`; run it detached (`nohup`).
+
 ## TL;DR status
 
 | Piece | State |