mlx-bun 0.0.4 → 0.0.6

package/docs/design/w4a16-compute-precision-spike.md

@@ -0,0 +1,232 @@
+# Research spike — W4A16 head GEMM (compute precision, not storage)
+
+**Status:** proposed (not started) · **Opened:** 2026-06-19 · **Owner:** Josh
+**Surface:** `src/train/flash-cce.ts` (steel + SG kernels) · **Gate:** the existing
+1e-5 CCE parity gate + greedy-prefix logit parity.
+
+> One-line thesis (Josh): *we've maxed **storage** quant (4-bit weights), but we
+> dequant into fp32 and run the matmul in fp32 — so the head GEMM is **W4A32**.
+> The untouched lever is **compute** precision: dequant 4-bit → half, matmul in
+> `half` operands with an fp32 accumulator (**W4A16**), gate on parity. Same 4-bit
+> storage, untouched. The only change is the float width we dequant into and
+> multiply in.*
+
+---
+
+## TL;DR
+
+- The flash-CCE head (forward logit GEMM `E@Wᵀ` and backward phase-2 `coeffᵀ@W`)
+  currently runs **W4A32**: 4-bit weights are dequantized into **fp32** threadgroup
+  tiles and multiplied with **fp32** simdgroup matrices. Verified in source (below).
+- The 4-bit packing bought the cheap HBM read; the fp32 matmul throws away the
+  compute advantage. Apple GPUs run **fp16 FMA at ~2× fp32** on the regular SIMD
+  ALUs (no tensor cores). Both head GEMMs are **compute-bound** (see roofline note),
+  so half operands are the lever.
+- **Industry-standard for the forward** (MLX `quantized_matmul`, llama.cpp, every
+  on-device 4-bit stack do exactly W4A16). Low risk; do it first.
+- **Backward needs the mixed-precision caveat** — fp32 accumulate is non-negotiable
+  (already planned), and the **operands** (`coeffᵀ`, `h`) need a dynamic-range probe
+  before trusting fp16, specifically the near-zero coeff tails after the skips fire.
+- **M1 Max constraint:** no `bf16` simdgroup_matrix (M2+ only), so **fp16 is the
+  only half option** — which makes the range-probe matter more here than on M3.
+- Gate exactly like the existing skip epsilons: **fp32 stays the default-exact
+  parity mode**; half is the training default; the 1e-5 gate test runs against fp32.
+
+---
+
+## Current state — verified W4A32
+
+The default production path for **both** forward and backward is the **steel**
+kernel (MLX's verbatim `BlockMMA` GEMM, `steel-qmm-header.ts`), selected when
+`H,V,blockV0 % 32 == 0` (`flash-cce.ts:1034`, `:1130`). The `simdgroup_float8x8`
+SG kernels and the scalar/lane kernels are fallbacks for non-32-divisible shapes or
+explicit `MLX_BUN_CCE_*` overrides.
+
+Every one of these instantiations is fp32-operand:
+
+| path | site | instantiation | operand width |
+|---|---|---|---|
+| steel fwd (default) | `flash-cce.ts:950` | `BlockMMA<float,float,…>` | fp32 |
+| steel bwd phase-1 (default) | `flash-cce.ts:799` | `BlockMMA<float,float,…>` (lmma_t) | fp32 |
+| steel bwd phase-2 (default) | `flash-cce.ts:800` | `BlockMMA<float,float,…>` (dmma_t) | fp32 |
+| steel staging tiles | `:596–598`, loaders `:801` | `threadgroup float` / `BlockLoader<float,…>` | fp32 |
+| SG fwd (fallback) | `:251–254` | `simdgroup_float8x8`, `Csh`/`Esh` = `threadgroup float` | fp32 |
+| SG bwd (fallback) | `:617`, `:684–704` | `simdgroup_float8x8`, `Wsh` = `threadgroup float` | fp32 |
+
+The dequant itself (`Csh[…] = sc * q + bi`, `flash-cce.ts:244`) computes into fp32
+and stores fp32. The dh accumulator is already fp32 (`MMATile<float,…>`,
+`flash-cce.ts:828`) — i.e. **the accumulator is already where it should be**; only
+the matmul *operands* are wider than they need to be.
+
+---
+
+## The lever — what actually changes
+
+W4A16 = **same 4-bit storage, half-width operands, fp32 accumulate.** Concretely:
+
+1. **Steel operand type** — the `BlockMMA<float,float,…>` first template param is the
+   GEMM input/operand type; flip operand to `half` while keeping the
+   accumulator/fragment in `float`. (**Spike task 0: confirm the exact MLX steel
+   template param that selects operand vs accumulator width** in
+   `steel-qmm-header.ts` before editing — don't assume the param order.)
+2. **Staging tiles** — `Esh`, `Cstg`, `Wsh`, `Csh` go `threadgroup half`; dequant
+   arithmetic can stay fp32 then narrow on store, or compute in half.
+3. **Loaders** — `BlockLoader<float,…>` → `BlockLoader<half,…>` to stage half tiles.
+4. **SG fallback** — `simdgroup_float8x8` → `simdgroup_half8x8` for the operands,
+   keep an `simdgroup_float8x8` accumulator (Apple supports the mixed `half×half→float`
+   `simdgroup_multiply_accumulate`).
+
+Everything reduction-shaped (the online-softmax `m`/`d` recurrence, the dh accum,
+logsumexp) **stays fp32**.
+
+---
+
+## First thoughts (mine) — what to add to Josh's analysis
+
+The analysis is right and the change is well-targeted. Five additions, each grounded
+in the current source:
+
+**1. Aim at the steel kernel, not the SG snippet.** Josh's note led with the SG path
+(`simdgroup_float8x8`, `Csh`). That path is a *fallback*. The lever that moves the
+measured 180 ms fwd / 754 ms bwd (e4b) numbers is the **steel `BlockMMA<float,float>`**
+instantiations (`:950`, `:799`, `:800`) and their `BlockLoader<float>`/staging tiles.
+Do the SG path too for completeness, but the steel path is where the production time is.
+
+**2. The 2× is an upper bound gated by the FMA fraction — get the roofline split
+first.** The whole flash-CCE design is `[M,V]`-free and amortizes dequant across
+`BLOCK_B` tokens specifically because the **dequant ALU** was "the dominant cost in
+v1" (`flash-cce.ts:59`). The runtime splits into (a) nibble-unpack + `sc·q+bi`
+dequant and (b) the FMA matmul. **W4A16 only speeds up (b).** The kernel is
+unambiguously compute-bound, not BW-bound — **confirmed live** (see the measured box
+below), and Josh flagged it from `mactop` before this doc existed. *But* the win
+ceiling is the FMA share of that ALU time. The dequant-amortization work was done
+precisely to shrink (a)'s share, which *should* leave (b) dominant and favorable — but
+**measure the (a)/(b) split before predicting 2×.** The live numbers prove the
+*whole kernel* is ALU-bound; they do **not** resolve the intra-kernel (a)/(b) split,
+which is what actually bounds the reclaim. If the kernel is still dequant-ALU-bound,
+half operands move it less.
+
+> **Live roofline — measured (M1 Max 32 GB, MiniCPM5-1B ORPO, `--no-segment`, seq 4096):**
+> GPU **100% @ 1296 MHz, 81 °C, Nominal** (sustained, not throttling) · memory
+> **BW 129.2 GB/s ≈ ⅓ of the M1 Max's ~400 GB/s ceiling** (bus two-thirds idle) ·
+> mem **17.37 / 32 GB (54%)**, swap 0 · CPU ~5%, ANE 0%. The training `bun` proc is
+> 92.6% GPU. **GPU pegged + memory bus mostly idle = ALU-bound.** This is the regime
+> where W4A16 pays; if it were BW-bound (BW near 400, GPU spiky) the spike would be
+> moot. The 4-bit storage already did its job — *that's why BW is low* — and the GPU
+> is now spending its cycles on the fp32 FMAs.
+
+**3. Threadgroup memory is a free secondary win — possibly a compounding one.** The
+staging tiles (`Esh`/`Cstg`/`Wsh`/`Csh`) are the kernel's threadgroup-memory pressure,
+and tile sizing is *already* occupancy-constrained — see the note at `flash-cce.ts:71`
+("the 8KB `Csh` tile measured SLOWER"). Half-width staging **halves that footprint**,
+which can either raise occupancy or admit larger tiles. So W4A16 may buy more than the
+raw 2× FMA: it could relieve the exact constraint that capped the current tiling.
+Worth re-sweeping `BLOCK_B`/tile dims *after* the precision change, not just measuring
+the drop-in.
+
+**4. SDPA in the segmented forward is probably already half — don't chase it.** Josh's
+aside ("if your SDPA tiles are also fp32, same logic"). The segmented forward uses
+`ops.sdpa` — MLX's own fused flash kernel ([[opssdpa-dk-vjp-bug]]) — which runs in the
+model's dtype (bf16), **not** our fp32 tiles. So that lever is likely already pulled by
+delegating to MLX. One-line confirm, then drop it; the head GEMM is the real target.
+
+**5. Free instrumentation hooks already exist for the backward range-probe.** The
+backward already computes `cmax = simd_max(coeff)` for the FILTER_EPS skip
+(`flash-cce.ts:682`). That same reduction is the natural hook to log the coeff dynamic
+range (min nonzero |coeff|) across a real run *before* committing the operand to fp16 —
+the underflow-to-zero case Josh flagged. Pair it with a max-abs probe on `h`: at e4b's
+`scale=1.0` peaked attention the residual stream can carry outliers
+([[e4b-forward-length-sensitivity]], [[siglip-vision-parity-cross-build]]), and with no
+bf16 on M1 Max we can't lean on a wider exponent to eat them.
+
+---
+
+## Risk analysis (preserving the forward/backward split)
+
+**Forward — low risk, boring industry default.** Logits are well-scaled; the
+online-softmax max-subtraction bounds the `exp` args. W4A16 here is what the dense
+path already assumes. Validate against the 1e-5 gate (fp16 operands won't hold 1e-5 —
+that's expected; the **fp32 path stays the exact-parity gate**, half is graded by
+greedy-prefix logit agreement, same as the vision tier-a bar).
+
+**Backward — fp16 matmul is mainstream, but as *mixed precision*, with eyes on
+operand range.** What makes the forward safe (well-scaled, max-subtracted) is **not**
+automatically true here:
+- `coeff_v = g·(onehot − softmax)·sech²` — after the skips fire, the survivors are the
+  **near-zero tails**. Small magnitudes are exactly where fp16's ~3-digit mantissa and
+  narrow range (subnormals < ~6e-5) bite. **`coeffᵀ` as an fp16 operand is the highest
+  risk** — tail coeffs underflowing to zero silently drop gradient contribution. Probe
+  the range first (task 5 above).
+- `h` (activations) as fp16 operand — residual outliers; the classic bf16-beats-fp16
+  case, but bf16 isn't available on M1 Max. Range-check before trusting.
+- **fp32 accumulate is non-negotiable** (already planned; `MMATile<float>` already is).
+- **Validate grads against autograd, not a hand-derived backward** ([[opssdpa-dk-vjp-bug]]).
+
+---
+
+## Spike plan
+
+Sequenced so each step is independently measurable and revertible (Josh: *"Don't flip
+both at once — you'll learn more, and if parity drifts you'll know which GEMM did it."*).
+
+- [ ] **0. Confirm the steel template semantics.** In `steel-qmm-header.ts`, identify
+  which `BlockMMA`/`BlockLoader` template params are operand vs accumulator width.
+  Exit: a one-line mapping, no code change.
+- [x] **1a. Whole-kernel roofline — DONE (live).** GPU 100% / BW 129 GB/s (~⅓ of 400)
+  on M1 Max confirms the kernel is ALU-bound, not BW-bound (see the measured box above).
+  The spike is worth chasing.
+- [ ] **1b. Intra-kernel split (still open).** Profile the steel fwd (180 ms) and bwd
+  (754 ms, e4b) into dequant-ALU vs FMA fractions — 1a says the GPU is the wall, this
+  says *which part* of the GPU work. Exit: a number that bounds the expected win. *If
+  FMA < ~40% of either kernel, reset expectations before coding.*
+- [ ] **2. Add the precision gate.** New codegen `#define` (e.g. `CCE_A_HALF`) +
+  env flag (e.g. `MLX_BUN_CCE_A16`), mirroring `CCE_BWD_FILTER_EPS`/`CCE_BWD_BLOCK_EPS`
+  (`flash-cce.ts:732`,`:739`). Default: fp32 (exact) for the parity suite; half is the
+  training default. Per [[three-level-fidelity-tree-model]] the exact path must stay
+  selectable.
+- [ ] **3. Forward → half operands.** Steel fwd `BlockMMA`/loader/staging to half.
+  Validate: 1e-5 gate (fp32 path) + greedy-prefix logit parity (half path) + measure
+  the fwd GEMM delta. Land as its own change.
+- [ ] **4. Backward range probe.** Instrument coeff (via the `simd_max` hook) and `h`
+  dynamic range across a real MiniCPM5/e4b ORPO step. Exit: min-nonzero-coeff and
+  max-|h| numbers + a go/no-go on fp16 operands.
+- [ ] **5. Backward → half operands** (only if 4 says go). `coeffᵀ`/`h`/`W` operands
+  half, dh accumulator stays fp32. Validate grads **vs autograd** + measure bwd delta.
+- [ ] **6. Re-sweep tiling.** With half-width staging, re-check `BLOCK_B`/tile dims and
+  occupancy (first-thought #3) — the precision win may unlock a tiling win.
+
+**Exit criteria for the spike:** a measured fwd/bwd speedup on e4b (V=262144) and
+MiniCPM5 (V=130560) on the bench machine (M4 Pro 24 GB; numbers verified against
+`~/.cache/mlx-bun/evals.sqlite`, not docs — [[verify-benchmarks-against-evaldb-not-docs]]),
+with the fp32 exact path still passing the 1e-5 gate and the half training path inside
+the greedy-prefix parity bar. Either it lands as the training default behind the gate,
+or the roofline/range data explains why it didn't pay and the spike closes with that.
+
+---
+
+## Open questions
+
+- Steel `BlockMMA` mixed operand/accum: does MLX's steel already support
+  `half`-operand / `float`-accum, or does it need a local template tweak? (task 0)
+- Is the dequant cheaper to do **in half** (fewer ALU bits) or **in fp32 then narrow**
+  (preserves the scale/bias precision)? A/B inside task 3.
+- Does the half staging change the `+4-float` padding trick (`flash-cce.ts:197`, the
+  BK_padded bank-conflict avoidance)? Half-width may re-introduce or shift bank
+  conflicts — re-derive the pad.
+
+---
+
+## References
+
+**Code:** `src/train/flash-cce.ts` — steel fwd `:937–1028`, steel bwd `:763–935`, SG
+fwd `:174–270`, SG bwd `:561–726`, gates `:732`/`:739`/`:1034`/`:1130`; staging tiles
+`:596–598`; dequant `:244`; coeff `simd_max` hook `:682`; dh fp32 accum `:828`.
+`src/train/steel-qmm-header.ts` (the verbatim MLX BlockMMA). `src/mlx/metal-kernel.ts`
+(kernel infra).
+
+**Cross-links:** [[flash-cce-metal-kernel-prototype]] · [[three-level-fidelity-tree-model]]
+· [[opssdpa-dk-vjp-bug]] · [[e4b-forward-length-sensitivity]] ·
+[[verify-benchmarks-against-evaldb-not-docs]] · `docs/investigations/steel-flash-cce-handoff.md`.
+
+**Sibling deferred lever** (queued together, post-run): `docs/design/orpo-future-enhancements.md`
+→ "Deferred levers" (expose `--grad-accum`).

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).