mlx-bun 0.5.0 → 1.0.0

package/docs/archive/investigations/optimization_plan.md

@@ -1,346 +0,0 @@
-# mlx-bun: gemma4.ts decompose & optimize — action plan
-
-> This is the canonical **design/rationale** for the optimization work.
-> Live **status/tracking** (checkboxes, findings) for Phases A–D lives in
-> PLAN-archive.md; Phase E remains in PLAN.md.
-
-**Status (verified 2026-06-17):** Phases A–D are **DONE** (`[x]` in PLAN-archive.md, all landed 2026-06-10/11). Phase E — per-config fused decode-SDPA kernel — is **in progress** (`[~]` in PLAN.md, started 2026-06-11; fused kernel exists and is quality-gated per PLAN.md Step 4). The body of this doc remains the durable design rationale; completed phase checkboxes and findings live in PLAN-archive.md; Phase E tracking lives in PLAN.md.
-
-**Premise (confirmed from the source read):** the current engine sits at
-parity-to-ahead vs mlx-lm on the 12B *while still interpreting the decode
-graph every token through hundreds of `bun:ffi` crossings, uncompiled.* The
-largest remaining cost is host-side graph rebuilding that is identical
-token-to-token — exactly what `mx.compile` exists to erase. Specialization and
-custom kernels are real but smaller, and both are only *measurable* once the
-host overhead is gone.
-
-**Therefore the ordering is non-negotiable:** compile first (the big lever),
-then decompose + specialize (enables the kernel and cleans the graph), then the
-per-config fused kernel (the 2–3× on top). Doing it in any other order means
-optimizing against a baseline that compile would have changed anyway.
-
-**Invariant for the entire plan:** the compat-mode bit-exact parity gate
-(`toBe(0)` logits, identical greedy trajectories vs the oracle venv) stays green
-at *every* checkpoint. It is the instrument that makes refactoring safe to do
-fast — any step that changes behavior turns it red the same hour. No step
-proceeds with a red gate. No tolerance is ever added to make a refactor pass;
-divergence is root-caused, never papered over.
-
----
-
-## Phase A — Compile the decode step (on the current monolith, before any split)
-
-**Why first:** this is the 14×-class lever for *this* runtime specifically,
-because mlx-bun's per-step cost is FFI-crossing-heavy in a way Python's
-interpreter loop is not. It is independent of decomposition. Landing it first
-means every later phase is measured against a compiled baseline, so you learn
-exactly how much headroom specialization actually has left to chase instead of
-chasing host overhead that compile already removed.
-
-**Steps:**
-
-1. **Characterize before touching anything.** Run `decode-split.ts (deleted; git history)`
-   (the per-step wall-time splitter that already exists) on all three models at
-   @600 and @8k. Record GPU-blocked time vs JS-graph-build time per step. This
-   is the control measurement; everything in Phase A is judged against it. The
-   prediction: e4b shows the largest JS-graph-build fraction (it has the
-   per-layer-input slicing loop no other model runs).
-
-2. **Wire `mx.compile` around the single-token decode graph.** The unit to
-   compile is one `forward(tokens=[next], cache)` step: embed → forwardLayers →
-   finalNorm → logitsFromHidden → argmax. At fixed L=1 every op is shape-stable
-   token-to-token; only input values and cache offsets change. Use
-   `shapeless=True` semantics so the same compiled graph serves any decode step.
-   The cache update is the one stateful part — confirm the compiled region
-   either includes the sliceUpdate writes or is structured so the cache mutation
-   happens at stable shapes (see step 4).
-
-3. **Keep prefill uncompiled (or separately compiled).** Prefill is
-   variable-length (L = prompt size), so it does not share a compiled graph with
-   decode. Leave prefill on the current path initially; a separate
-   fixed-tile-size compiled prefill is a later option, not part of this phase.
-
-4. **Resolve the cache-mutation/compile boundary.** The quantized caches do
-   `quantize → sliceUpdate(ring) → activeSlices` per step (≈670 FFI crossings
-   per token on the 12B across 48 layers). Decide whether these live inside the
-   compiled graph (preferred — collapses the crossings into one replay) or
-   immediately outside it. The ring-index bookkeeping is host-side integer math
-   and stays in JS; only the array ops compile.
-
-5. **Parity gate.** Compiled decode must produce bit-exact identical token
-   trajectories to the uncompiled path on the fixed prompt suite. Compile must
-   not change numerics — if it does, a fusion inside the compiled graph reordered
-   an op; root-cause it.
-
-6. **Measure.** Re-run `decode-split.ts`. Expected: JS-graph-build fraction
-   collapses, e4b ~5% residual shrinks or closes, 12B short-context moves
-   further ahead. Record the delta — this number tells you how much the rest of
-   the plan is even worth.
-
-7. **Re-run the fused-decode A/B *after* compile, not before.**
-   `MLX_BUN_FUSED_DECODE` measured ~3% *worse* uncompiled, almost certainly
-   because the tile loop's per-iteration host overhead dominates at L=1. Compiled,
-   that host cost vanishes and the flag may flip to a win with no kernel written.
-   This is a free re-measurement that may close part of the gap on its own.
-
-**Phase A exit:** compiled decode green on parity, measured speedup recorded per
-model, fused-decode A/B re-evaluated. Decision point: how much headroom remains?
-That answer scopes Phases D–E.
-
----
-
-## Phase B — Extract a shared, config-independent base
-
-**Why:** before generating per-model files, separate the machinery that is
-genuinely identical across all architectures from the parts that vary. This is
-mechanical, parity-preserving code movement — no behavior change, gate stays
-green throughout.
-
-**Moves into `gemma4-base.ts` (or `src/model/gemma4/base.ts`):**
-
-- Cache classes: `KVCache`, `QuantizedKVCache`, `RotatingKVCache`,
-  `RotatingQuantizedKVCache` — these are architecture-independent storage; which
-  *combination* a model uses varies, but the classes themselves don't.
-- SDPA functions: `quantizedSdpaUnfused`, `quantizedSdpaTiled`,
-  `fusedSdpaSupported`, `quantizedSdpa`, plus `ops.sdpa` plain dispatch.
-- Mask builders: `createCausalMask`, `bidirMask`, the `Mask`/`MaskMode` types.
-- Primitives: `RMSNorm`, `QuantizedLinear`, `QuantizedEmbedding`,
-  `QuantizedSwitchLinear` and the LoRA machinery (`LoraWeights`, `LoraState`).
-- Helpers: `disposing`, `mapTriple`, `disposeTriple`, `logitSoftcap`,
-  `argmaxLastPosition`, `lastPositionLogits`, `FINFO_MIN`.
-
-**Stays model-specific (does NOT move to base):** `Attention`, `MLP`, MoE
-(`Router`/`Experts`/`SwitchGLU`), `DecoderLayer`, `Gemma4Model.forwardLayers`,
-`computePerLayerInputs` — these branch on architecture and become the
-per-model generated code.
-
-**Steps:**
-
-1. Create the base module, move the listed classes verbatim, fix imports.
-2. Run the full parity suite + all behavioral tests. Green = the move was
-   behavior-preserving. This is the only acceptance criterion for Phase B.
-3. No optimization in this phase. Pure extraction. Resist the urge to "improve
-   while moving" — that conflates a refactor bug with an optimization bug.
-
-**Phase B exit:** base module extracted, every test green, monolith now contains
-only the model-specific assembly that Phase C will generate.
-
----
-
-## Phase C — Generate three per-architecture files (not author — generate)
-
-**Reframe from the original request:** the per-model files are *outputs of a
-generator*, not hand-written variants. Loop unrolling and branch elimination are
-mechanical transforms; a codegen script performs them from each model's config.
-Three files = three generator outputs. When the base algorithm changes, you
-regenerate, not hand-edit three copies.
-
-**The three targets (your actual runnable set):**
-
-- `gemma4-12b.ts` — dense. Full-attention `KVCache` (or quantized variant),
-  GQA, no per-layer-input, no MoE, no KV-sharing. The simplest graph.
-- `gemma4-e4b.ts` — the complex one. Per-layer-input gating, KV-sharing
-  (donor/sharer layers), sliding/rotating quantized cache. This is where the
-  per-step host overhead concentrated, so it's where specialization pays most.
-- `gemma4-26b.ts` — MoE. Router + experts (`SwitchGLU`/`gather_qmm`), parallel
-  dense+routed branches.
-
-**The generator (`scripts/gen-model.ts`):** reads `config.json` +
-`kv_config.json` for a model and emits the specialized file. What it eliminates,
-all resolvable at generation time from known constants:
-
-- The per-layer `instanceof` cache-type check (gemma4.ts:1327) — each layer's
-  cache type is known from `layer_types`, so the generated layer calls the right
-  fetch directly, no runtime `instanceof`.
-- The `if (this.router)` / `if (perLayerGate)` / KV-sharing forks in
-  `DecoderLayer.forward` and `forwardLayers` — each model takes exactly one
-  branch; the others are deleted from its file.
-- The `forwardLayers` layer loop — unrolled into straight-line per-layer calls
-  with each layer's constants (type, cache index, donor index, sliding/full)
-  baked in. (Note: post-compile this is largely cosmetic for *speed* — compile
-  already erased the per-step cost — but it produces a cleaner compiled graph
-  and removes the `intermediates`/`previousKvs`/`cacheIndex` indirection arrays
-  and their per-step lookups.)
-- The per-layer-input `slice`+`reshape` loop (gemma4.ts:1865) on non-e4b models:
-  deleted entirely (they don't have it). On e4b: the slice bounds are constant
-  per layer, so they're baked.
-
-**Steps:**
-
-1. Write `gen-model.ts`: input = model dir (config + kv_config), output = one
-   `.ts` file importing from `gemma4-base.ts`. Start with the 12B (simplest) to
-   prove the generator end-to-end.
-2. Generated 12B must pass parity bit-exact against the oracle. This proves the
-   generator preserves behavior — the gate is the generator's test.
-3. Generate e4b and 26B. Each passes parity. The KV-sharing donor/sharer wiring
-   in e4b is the trickiest to generate correctly; its parity test is the proof.
-4. Dispatch: model load picks the generated module by config fingerprint
-   (arch + kv_config hash), with the **monolithic `gemma4.ts` retained as the
-   general fallback** for any model without a generated specialization. Nothing
-   is deleted; specialized files are opportunistic overrides.
-
-**Phase C exit:** three generated, branchless, unrolled per-model files, each
-bit-exact, dispatched by fingerprint, with the monolith as fallback. Measure
-each vs its Phase-A compiled baseline — expect small wins (cleaner graph, fewer
-indirection lookups), not large ones. The large one was compile.
-
----
-
-## Phase D — Per-config kernel specialization (kv_config-driven)
-
-**Why:** now that each model has its own file and the cache/SDPA dispatch is
-already branch-free per model, bake the *quantization* constants. The
-kv_config.json's per-layer (bits, group_size) becomes compile-time information
-in the generated SDPA and cache calls.
-
-**What gets specialized, from the 12B's kv_config as the worked example**
-(42 layers @ 4-bit/gs64, 6 layers @ 8-bit/gs64, group_size uniformly 64):
-
-- **Per-layer bit-width resolved at generation, not per-step.** The generated
-  layer for index 6 calls the 8-bit SDPA path; index 0 calls the 4-bit path.
-  No runtime bit-width parameter, no `fusedSdpaSupported` bit check — the
-  generator already knows.
-- **group_size folded as a compile-time constant** (64 everywhere in this
-  config) — `dim / group_size`, `32 / bits` (elements-per-int) become literals.
-- **nRep (GQA ratio) baked** — the `nRep > 1` branch and the `expandDims`
-  broadcast become a constant-folded reshape with the literal ratio.
-- **head_dim baked** — sliding (256) vs full-attention (512) head dims are known
-  per layer.
-
-**Important honesty:** Phase D is specialization of the *existing* op-graph
-(constant-folding, branch deletion), still using `ops.quantizedMatmulQT` etc.
-It is NOT yet a custom kernel. Expected gain is modest — it removes parameter
-passing and branches the compiled graph would partly handle anyway. Its real job
-is to make Phase E's custom kernel *possible* by having a single known
-(bits, group_size, nRep, head_dim) per dispatch site.
-
-**Steps:**
-
-1. Extend `gen-model.ts` to read kv_config and emit per-layer-specialized SDPA
-   dispatch (right bit-width path per layer, constants folded).
-2. Parity gate per generated model. Bit-exact required — constant-folding must
-   not change numerics.
-3. Measure vs Phase C. Record the (likely small) delta. If negligible, that's a
-   *finding*, not a failure — it means compile already captured it, and it tells
-   you the kernel in Phase E is the only remaining lever.
-
-**Phase D exit:** each generated model has branch-free, constant-folded,
-per-layer-specialized quant dispatch. Every dispatch site now has a single known
-shape — the precondition for Phase E.
-
----
-
-## Phase E — Per-config fused decode-SDPA kernel (the 2–3×, last)
-
-**Why last:** this is the custom Metal kernel, the fun part, and it's the
-*smallest* lever (the search and the turboquant retro both say the fused kernel
-was 2–3× while the buffer/host path was the big win — which you've already
-captured in Phases A and D). Doing it last means you measure its true
-contribution against a compiled, specialized baseline instead of attributing
-host-overhead noise to it.
-
-**What the kernel does** (the mlx-qsdpa shape, specialized per bit-width):
-
-- Decode-time (L=1) fused attention: stream the quantized KV through QK^T and
-  ×V in one Metal dispatch, dequantizing inline, **never materializing** the
-  `[B,H,1,N]` scores row or the dequantized cache.
-- **nRep baked in** — no `expandDims` GQA broadcast (gemma4.ts:916); the kernel
-  loops the query-head group internally with the literal ratio.
-- **bits/group_size as compile-time template parameters** — generate the 4-bit
-  variant and the 8-bit variant; dispatch the right one per layer from the
-  config (Phase D already resolved which).
-- Online-softmax fused (score → max → exp → ×V accumulate) so the scores never
-  hit memory.
-
-**Mechanism:** `mx.fast.metal_kernel` (the intended extension point — no MLX
-fork). Lift the *dequant unpacking logic* from a known-correct reference
-(mlx-qsdpa is open, `_BaseCache`-compatible, ~200 lines) — the bit-unpacking and
-scale/zero-point application is the well-trodden, copyable part. The novel part
-is the *fusion boundary* (not writing the intermediate) and the *Bun dispatch*
-(every existing impl dispatches from Python; a Bun-dispatched fused quant-SDPA
-is genuinely unoccupied territory).
-
-**Correctness regime — this is where the gate model changes:**
-
-- A fused kernel that's faster will **not** be bit-exact with the unfused path
-  (online softmax ≠ one-shot softmax in bf16 — this is already documented in the
-  tiled path's "tier b" note). So the compat bit-exact gate does NOT apply to
-  the kernel output directly.
-- Instead: the kernel runs under the **perf-mode flag** and is gated against the
-  **frozen ground-truth snapshot** (capture today's compat-mode logits/
-  trajectories on the fixed prompt suite *before* starting Phase E — that's the
-  quality oracle once bit-exact-vs-optiq no longer applies).
-- Perf-mode gate = bounded logit tolerance + token-trajectory agreement above a
-  set threshold, *labeled with the measured delta* (e.g. "+X% tok/s, −Y% on the
-  eval set"), the way a quant ships with its perplexity hit. "No guarantee of
-  identical output" means **bounded and measured**, never unchecked.
-- Compat mode (bit-exact, unfused) stays in the codebase permanently as the
-  `-O0` reference and differential-testing oracle: when the kernel misbehaves,
-  flip to compat, same prompt, and the diff localizes kernel-vs-expectation.
-
-**Steps:**
-
-1. **Freeze the ground-truth snapshot** (compat-mode logits + trajectories,
-   fixed prompt suite, all three models) before writing any kernel. This is the
-   perf-mode oracle.
-2. **Metal capture one decode step** (`mlx_metal_start_capture`) on the
-   compiled+specialized 12B to confirm the prize size: is the dequantized-cache
-   round-trip still fat after compile, or did fusion upstream already thin it?
-   If thin, the kernel ceiling is low — record and decide. If fat, proceed.
-3. **Toolchain derisk first:** get a trivial `mx.fast.metal_kernel` compiling,
-   dispatched from Bun, running, verified — *before* the GQA numerics. (RMSNorm
-   or a no-op fusion proves the `.metal` → mlx-c → `bun:ffi` pipeline end-to-end
-   so the real kernel debugs numerics OR plumbing, not both at once.)
-4. **Generate the 4-bit fused decode kernel** for the 42 four-bit layers (where
-   42/48 of KV traffic lives and 4-bit's dequant round-trip is most lopsided
-   against compute). nRep, group_size baked.
-5. **A/B vs compiled+specialized baseline**, paired, clean machine, at @8k (where
-   KV traffic dominates and fusion shows). Isolate the kernel's contribution.
-6. **Decide on the 8-bit six.** 8-bit dequant is cheaper and the data fatter, so
-   the fusion win per-byte is smaller — possibly not worth a second kernel; stock
-   MLX path may stay for those layers. Measure before committing kernel #2.
-7. **Per-model:** repeat the generation for e4b and 26B configs only if the 12B
-   kernel shows a real win. The generator (Phase C/D) already produces the
-   dispatch; the kernel is one more generated, per-bit-width artifact keyed by
-   fingerprint, each with its own auto-generated parity/quality test.
-
-**Phase E exit:** per-config fused decode kernel(s), perf-mode flag, gated
-against the frozen quality oracle with the trade *labeled*, compat mode retained
-as permanent reference. Measured isolated contribution recorded.
-
----
-
-## Cross-cutting discipline (applies to every phase)
-
-- **Parity gate green at every checkpoint.** Compat bit-exact for Phases A–D;
-  frozen-snapshot quality tolerance for Phase E's perf mode. Never add a
-  tolerance to make a refactor pass — root-cause divergence.
-- **One flag, one measurement per lever.** Compile, fused-decode, each
-  specialization, the kernel — each is independently togglable and independently
-  measured ("run it both ways, compare"), so you always know what each bought.
-  Never measure two new levers tangled together.
-- **Clean machine or it didn't happen.** Reboot + preflight-gate every quotable
-  benchmark; `sudo purge` for cold-start rows.
-- **Losing experiments become documented default-off flags, never deletions.**
-- **Nothing deleted, fallback always present.** The monolith stays as the
-  general fallback; specialized/generated files are opportunistic overrides
-  dispatched by config fingerprint. An unhandled config runs slow, never broken.
-- **Generator over hand-edits.** Per-model files are regenerated from config,
-  not maintained by hand. A base-algorithm fix is one regeneration, not three
-  edits. Each generated artifact ships with its generated parity test.
-
-## Expected outcome, honestly stated
-
-- **Phase A (compile):** the big win. Likely closes the e4b ~5% residual and
-  pushes the 12B further ahead. This is the lever that matters most and it's
-  pure host-side — no kernel, no Metal.
-- **Phases C–D (decompose + specialize):** small wins (cleaner graph, deleted
-  branches/indirection), main value is enabling Phase E and improving
-  maintainability via codegen.
-- **Phase E (kernel):** the 2–3× on the attention hot path, biggest at long
-  context, smallest lever of the three but the one with no upstream Bun-dispatched
-  equivalent — the genuinely novel contribution.
-
-The sequence front-loads the largest, cheapest win (compile) and ends with the
-most fun, most novel, smallest one (kernel), with every step gated by the oracle
-that's kept the project honest from the start.

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

@@ -1,120 +0,0 @@
-# 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/archive/investigations/orpo-flash-cce-pin-leak.md

@@ -1,167 +0,0 @@
-# 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).