mlx-bun 0.0.13 → 0.3.0

package/docs/design/orpo-dynamic-lambda.md

@@ -1,220 +0,0 @@
-# Dynamic λ controller for ORPO (adaptive preference pressure)
-
-**Status:** proposed (design only — not built). **Opened:** 2026-06-20.
-**Reaffirmed + reframed by Josh 2026-07-01** (see "The band framing" below).
-**Surface:** `src/train/loss.ts` (orpoMetrics/orpoLoss), `src/train/trainer.ts`
-(orpoLoop, the λ config), the segmented ORPO classes in `src/train/segmented.ts`.
-
-## The band framing (Josh, 2026-07-01 — the product statement)
-
-> "Vary the lambda term based on the loss or another signal — maintain the
-> interest inside a band: learn preference HARDER while you are still able to
-> stay coherent. A self-stabilizing AND self-discovering hyperparameter, which
-> optimizes the training to get as much preference as possible while not
-> resulting in degeneracy."
-
-That is the controller's contract in one sentence: **λ is not a hyperparameter
-to tune, it is a process variable to control** — push preference pressure to
-the maximum the model can currently bear (self-discovering), retreat on
-confirmed coherence damage (self-stabilizing), and hold the system inside the
-healthy band rather than at a fixed point. Everything below (phases, AIMD/PID,
-displacement decomposition) is the mechanism for that contract.
-
-**Interplay with `sft_scope` (landed 2026-07-01):** the default chosen-NLL is
-now full prompt+response. The controller's displacement signal must key on the
-**response-only `lw`** (which `orpoMetrics` computes for the odds ratio in both
-modes) — full-scope NLL mixes prompt modeling into the level and would mask a
-crushed `logp(chosen-response)`. Instrument both, control on response-only.
-
-## Paper lens (2026-07-01): is this citable?
-
-Closest prior: **β-DPO (arXiv 2407.08639)** — dynamic β for DPO, calibrated
-per-batch from the reward gap. Ours differs on three axes worth claiming:
-(1) **closed-loop trajectory control** (a controller tracking a non-stationary
-optimum across the run) vs per-batch heuristic calibration; (2) the control
-signal is **displacement-aware** — the (dLw, dLr) decomposition targets the
-known ORPO/DPO failure mode directly (cf. likelihood displacement,
-arXiv 2410.08847 "Unintentional Unalignment"); (3) **objective-agnostic**
-actuator (ORPO λ, SimPO β/γ). The incremental P→PD→PID build IS the ablation
-section. The claim to stake: *"the optimal preference pressure is
-non-stationary, so a controller strictly dominates any fixed λ a sweep can
-find — and a displacement-aware control law finds the coherence cliff without
-a priori knowledge of it."* Referees will demand: ≥2 datasets, ≥2 model sizes,
-the static-λ sweep as the baseline family (the controller must beat the BEST
-constant, not λ=0.1), seeds, and the positive-control sanity run this doc
-already specifies. Workshop/arXiv-preprint tier as a single-model study;
-main-track needs the multi-model sweep.
-
-## Motivation
-
-ORPO's loss is `L = NLL(chosen) + λ·OR`. The two terms learn on different
-timescales, and a **static λ** can't serve both:
-
-- **The SFT/NLL half learns easily** — the model quickly raises `logp(chosen)`
-  (fits the response *style*). Observed live: val loss falls steadily while
-  preference accuracy/margin barely move.
-- **The preference/OR half is hard** — on subtle data (UltraFeedback on a 1B) the
-  model doesn't generalize "which of two plausible answers is better."
-- **Over-weighting OR backfires** via **likelihood displacement**: the model widens
-  the *relative* margin by **crushing `logp(rejected)`** (and sometimes dropping
-  `logp(chosen)` too) rather than genuinely preferring chosen — degrading fluency.
-
-A fixed λ either under-pressures the preference (λ too low → never leans on OR) or
-over-pressures it (λ too high → displacement/degeneration). The hypothesis: an
-**adaptive, displacement-aware λ** — low while SFT learns, ramped once SFT
-saturates, backed off the moment displacement appears — can find a generalization
-regime a static λ sails past. (Static λ tuning was the blunt version we rejected;
-this is the targeted version.)
-
-## Control signals (what we measure)
-
-`orpoMetrics` (`loss.ts:1249`) already computes everything; we just don't surface it:
-
-- **`lw`** = mean `logp(chosen)` (currently folded into `nll = mean(−lw)`).
-- **`lr`** = mean `logp(rejected)` (currently only used inside the odds ratio).
-- **`margin`** = mean log-odds (chosen vs rejected) — already emitted.
-- **`accuracy`** = fraction `lw > lr` — already emitted.
-
-The controller keys on the **decomposition of margin growth**:
-- margin ↑ because **`lw` ↑** → *healthy* (genuinely preferring chosen).
-- margin ↑ because **`lr` ↓** while `lw` flat/falling → *displacement* (the cheat).
-
-So the load-bearing new instrumentation is **surfacing `lw` and `lr` separately**
-per eval (not just `nll`/`margin`).
-
-## The controller (state machine)
-
-Runs at **each val eval** (every `stepsPerEval` steps) — val (`lw`, `lr`) is the
-stable signal; per-step B=1 train metrics are too noisy to control on. Maintains a
-short window of recent val (`lw`, `lr`, margin, accuracy).
-
-```
-state: { phase, lambda, window[] }   # lambda starts at LAMBDA_MIN
-on each val eval (lw, lr, margin, acc):
-  dLw = lw - lw_prev ; dLr = lr - lr_prev
-
-  PHASE 1 — SFT warmup (lambda = LAMBDA_MIN, low):
-    stay while the NLL is still improving (dLw > EPS_SFT, i.e. chosen logp rising).
-    when dLw plateaus (|dLw| < EPS_SFT over K evals) → SFT saturated → PHASE 2.
-
-  PHASE 2 — preference ramp:
-    each eval, lambda += RAMP (cap LAMBDA_MAX), applying OR pressure —
-    AS LONG AS chosen holds: dLw >= -EPS_DISP.
-    if accuracy/margin respond (margin widening with dLw >= 0) → keep ramping.
-    if displacement detected (dLw < -EPS_DISP, i.e. chosen logp dropping, or
-       dLr << 0 with dLw <= 0) → PHASE 3.
-
-  PHASE 3 — displacement backoff:
-    lambda -= BACKOFF (floor LAMBDA_MIN) until chosen recovers (dLw >= 0),
-    then hold (hysteresis) before considering another ramp.
-```
-
-Knobs (all tunable): `LAMBDA_MIN≈0.05`, `LAMBDA_MAX≈0.5`, `RAMP/BACKOFF≈0.05`,
-window `K≈3` evals, `EPS_SFT`/`EPS_DISP` set from the observed val-`lw` noise scale.
-Hysteresis (require K consecutive evals before a phase flip) prevents oscillation.
-
-**Degeneration guard (hard stop):** if `lw` falls below its phase-1 starting value
-by more than a threshold, freeze λ at `LAMBDA_MIN` and log — the model is paying for
-margin with chosen's likelihood, which a backoff alone isn't recovering.
-
-## Control-theory framing (PID / AIMD) — added 2026-06-23
-
-The phase state-machine above is one instance of a more general and more honest
-frame: **this is closed-loop control.** The **process variable** is the SFT signal
-(NLL / `lw`), the **actuator** is the preference weight (ORPO's λ, *or* SimPO's
-β/γ — the controller is objective-agnostic), and the **control law** is what maps the
-measured degradation back to the actuator. Two laws worth naming:
-
-- **AIMD** (additive-increase / multiplicative-decrease — TCP congestion control):
-  ramp the preference up *gently* while the SFT is healthy, **slam it down hard** the
-  moment degradation ("congestion") appears. Famously robust and forgiving of its own
-  knobs. The PHASE-2 ramp / PHASE-3 backoff is essentially AIMD.
-- **PID** decomposes the response into three terms, each fixing a failure the others
-  can't:
-  - **P (proportional · the level):** how degraded is `lw` *right now* → back off in proportion. (Leaves a steady-state offset.)
-  - **D (derivative · the slope):** is degradation *accelerating* → a **leading indicator** that catches the helping→hurting **inflection before the level rises**, so you pull back *earlier* and leave less damage to heal. **Gotcha:** differentiating a noisy signal amplifies noise → use a *robust* slope (slope of the EMA, or a windowed regression), never raw step deltas. The slope is where most of the value is.
-  - **I (integral · accumulated error):** the "missing letter." P leaves a chronic offset and D ignores it (flat slope at steady-state) → the model can settle hovering with a *small persistent* NLL elevation = quiet permanent SFT damage. I accumulates that residual and drives `lw` **back to its floor**. **Gotcha — integral windup:** if the actuator is pinned (λ at floor) while error persists, the integral piles up and **overshoots** on recovery → needs **anti-windup** (clamp/stop accumulating when saturated).
-
-**Build it incrementally — P → PD → PID** — each term earning its place against an
-*observed* failure of the simpler controller (P oscillates/offsets → add D; chronic
-offset remains → add I with anti-windup). That's both good control practice and an
-**ablation that writes the experimental section**, and it stops us from tuning three
-gains blind.
-
-**Why a controller beats a sweep (the real argument):** the optimal preference
-pressure is **non-stationary** — fragile early (push less), robust late (push more) —
-so **no fixed λ is optimal.** A sweep finds the best *constant*; the controller tracks
-the moving target and finds the best *trajectory*, which strictly dominates. The pitch
-isn't just "one run instead of six," it's "**beats any fixed value** a sweep could
-find." (The displacement backoff is also why this is *more* than auto-tuning — it
-targets the degradation signal directly.)
-
-**Degenerate failure to guard (beyond the §degeneration hard-stop):** if the law backs
-off too eagerly it **collapses the preference term to ~0 → you've trained pure SFT.**
-The AIMD/PID balance must **favor pushing** (aggressive-ish increase, retreat only on
-*confirmed* degradation), and target a small **safety margin below the cliff** rather
-than the absolute max-bearable (the edge is fragile).
-
-**Synergy with the diffusion frontier:** all of this is noise-robustness machinery for
-controlling off a noisy process variable — so it transfers *especially* well to a
-diffusion LM, where the likelihood is a stochastic ELBO (even noisier). See
-[orpo-base-uf-experiment-and-directions.md](../investigations/orpo-base-uf-experiment-and-directions.md) §6.7.
-
-## Wiring
-
-1. **Instrument (small):** add `lw`, `lr` to `orpoMetrics`'s return + the trainer's
-   `metric` emit (train and val). `pref-control.ts` / `metrics.json` log them so the
-   controller (and our plots) can see the decomposition.
-2. **Make λ dynamic:** today `cfg.orpoLambda` is a constant threaded into
-   `orpoLoss(model, batch, cfg.orpoLambda, …)` and **baked into the segmented ORPO
-   classes' constructors** (`new SegmentedBackwardOrpo(model, lora, ranges,
-   cfg.orpoLambda, …)`). Change to a **mutable λ source**:
-   - simplest: a `{ value: number }` box (or a getter) passed instead of the scalar;
-     the loss / segmented step reads `lambdaRef.value` each step.
-   - the segmented classes expose a `setLambda(x)` (or read the ref) so the
-     constructor-baked constant becomes a live field.
-3. **The controller** lives in `orpoLoop` (`trainer.ts`): after each val eval, feed
-   the val (`lw`, `lr`, margin, acc) to `updateLambda(state)` and write the new λ
-   into the ref. Log every λ change (`· λ 0.10 → 0.15 (SFT saturated, ramping)`).
-4. **Off by default:** gate behind a config flag (`orpoLambdaSchedule: "static" |
-   "adaptive"`); static = today's behavior, exact.
-
-## Experiment design (how we validate it)
-
-Head-to-head, **same data / seed / everything else**, on the from-base UF setup:
-- **Arm A:** static λ=0.1 (the current run — the baseline).
-- **Arm B:** adaptive λ (this controller).
-
-Compare:
-- **val accuracy + margin** trajectories (does B clear the noise floor / widen
-  margin where A stalls?);
-- **`lw` trajectory** (does B avoid displacement — `lw` non-decreasing — where a
-  high static λ would crater it?);
-- the **downstream** numbers: Exp-1 capability x2 and Exp-2 win-rate (does B's model
-  actually win more head-to-head?).
-
-Success = B generalizes measurably better (val acc clears ~0.55 with margin lifting,
-and/or Exp-2 win-rate > A's) **without** degeneration (`lw` holds). Null result is
-also informative: "even adaptive λ can't beat the 1B capacity ceiling on UF."
-
-## Open questions / risks
-
-- **Cadence:** controlling on val every `stepsPerEval` may be too coarse (few control
-  steps per run). Option: a cheaper, more frequent "control eval" on a small fixed
-  probe set distinct from the reported val set.
-- **Confounds the replication claim:** an adaptive-λ result is no longer "standard
-  ORPO" — it's our method. Keep Arm A (static 0.1) as the honest paper-replication;
-  Arm B is the "can we do better" research arm. Report both.
-- **Tuning the controller's own knobs** risks the same trap (we'd be tuning a tuner).
-  Mitigate: set EPS thresholds from observed noise scale, not by hand-fitting to one
-  run; validate the controller on the *positive control* (prefer-uppercase) first —
-  it should ramp λ early (SFT saturates fast) and never trip displacement.
-
-## References
-
-- `src/train/loss.ts:1249` `orpoMetrics` (computes `lw`,`lr`,`nll`,`or`,`accuracy`,`margin`).
-- `src/train/trainer.ts` `orpoLoop` (λ config, val emit), `DEFAULT_TRAIN_CONFIG.orpoLambda`.
-- `src/train/segmented.ts` segmented ORPO classes (λ baked in constructor).
-- `scripts/experiments/pref-control.ts` (trajectory logging — extend with `lw`,`lr`).
-- Context: [[opssdpa-dk-vjp-bug]] (validate grads vs autograd), the dynamic-λ
-  discussion + the live from-base UF val trajectory (2026-06-20).

package/docs/design/orpo-future-enhancements.md

@@ -1,340 +0,0 @@
-# ORPO future enhancements
-
-Performance-first follow-on ideas for mlx-bun's ORPO trainer, plus a short
-appendix of objective variants that may be useful once the runtime path is
-comfortable.
-
-> **LANDED since this was written (2026-06-19):** the top items here are now in
-> production — **Priority 1 (Cut Cross Entropy / fused linear CE)** and **§1
-> (vocab-axis chunked CE)** shipped as the `[M,vocab]`-free **flash-CCE Metal head**
-> (verbatim MLX steel GEMM, fwd+bwd; `orpo_flash_ce`), and **§3 (prefix sharing in the
-> trainer)** shipped as `orpo_prefix_shared`, composed with the flash head AND the
-> segmented backward for MiniCPM5 **and** e4b. See [training.md](../reference/training.md),
-> [orpo-quickstart.md](../reference/orpo-quickstart.md), and the live record
-> [steel-flash-cce-handoff.md](../investigations/steel-flash-cce-handoff.md). The
-> remaining sections (§2, §4–§9, the objective variants) are still genuinely future.
-
-## Deferred levers — queued, do NOT start mid-run (2026-06-19)
-
-Two optimizations identified during the MiniCPM5 UltraFeedback ORPO run. **Both are
-deferred until the active run finishes** — the GPU is fully occupied, so these get
-*measured* later, not "quickly tested" now (a quick probe steals the training GPU).
-
-1. **Expose `--grad-accum` (gradient accumulation).** The machinery already exists —
-   `accumulateStep` (`src/train/trainer.ts:975`) loops `gradAccumSteps` B=1 micro-steps,
-   **averages** their grads (`scale = 1/accumSteps`, line 989), and evals+frees each
-   micro-graph so **peak memory stays flat** (line 1015). It's wired into the ORPO loop
-   (`trainer.ts:837`) but `gradAccumSteps` defaults to 1 and has **no CLI flag**.
-   - **Why it helps us:** ORPO here is hard **B=1** (segmented + prefix-share assume a
-     single stream — `segmented.ts:636` throws on B>1), so every optimizer update is
-     computed from exactly **one** preference pair — the noisiest possible gradient (the
-     per-step loss swinging 0.4↔3.3 is that noise). Grad-accum gives an effective batch of
-     G — averaging G pairs/update, variance ↓ ~1/G, steadier descent, tolerates a higher
-     LR — at **B=1 memory**. It's the *only* way to a bigger effective batch when B>1 is
-     impossible.
-   - **What it is NOT:** not a speedup (G× passes per optimizer update → same total
-     passes, same wall-clock, same data seen); and not a guaranteed win (the prior run
-     converged fine at B=1, val 1.66→1.50). It's a **try-and-measure** lever — expose the
-     flag, A/B G=4/8 judged on **val-margin**, don't assume.
-   - **Note on prefix-share:** the existing `orpo_prefix_shared` is *already* a hand-built
-     batch-of-2-with-shared-prefix (1F + 1B over `[P; Rc; Rr]`, prompt deduped in **both**
-     directions — strictly cheaper than a generic 2-batch's `2F + 2B`). Generic N-way
-     batching would re-duplicate per-pair prefixes and need a block-diagonal mask layer
-     *on top of* the prefix-share mask — high friction, and the payoff is dominated by
-     grad-accum (effective batch, free) + W4A16 (the real GPU bottleneck). Park batching
-     unless a profile ever shows real short-row GPU idle.
-
-2. **W4A16 compute-precision spike.** The head GEMM dequants 4-bit weights into **fp32**
-   and matmuls in fp32 (W4A32); dropping the operands to `half` (fp32 accumulate) attacks
-   the actual bottleneck. Confirmed **compute-bound live** (GPU 100% / BW ~⅓ of 400 on
-   M1 Max). Full plan + parity gating: [w4a16-compute-precision-spike.md](w4a16-compute-precision-spike.md).
-
-**Context — seg vs no-segment A/B (MiniCPM5, seq 4096, same-machine M1 Max):**
-`--no-segment` measured **~1.25 s/step steady @ peak 9.26 GB** vs `--seg 2`
-**~1.58 s/step @ 2.62 GB** — ~20% faster for 3.5× peak memory, **bit-identical output**
-(segmenting is just gradient checkpointing). Dropping the segs removed the recompute (real
-GPU work) → GPU busy but no longer pinned. Fine when memory is spare; the current 3-epoch
-run uses `--no-segment`.
-
-Product relevance: ORPO is not only a researcher feature. It is the path to
-small, downloadable task adapters that make the default local assistant better
-without forcing every user onto a larger base model. The performance work below
-matters when it turns adapter training from a fragile lab run into a repeatable
-part of the mlx-bun product loop.
-
-## Priority 1: Cut Cross Entropy / Liger Fused Linear CE
-
-The top ORPO performance priority is a Cut Cross Entropy / Liger-style Fused
-Linear Cross Entropy path for the LM head. This should be treated as the next
-major kernel project before lower-impact objective variants.
-
-Why it is first:
-
-- Gemma's vocab is huge, so materializing `[tokens, vocab]` logits dominates
-  peak memory in ORPO.
-- ORPO needs chosen and rejected response log-probs, so the head cost is paid
-  twice unless the trainer can share or fuse more work.
-- Token-axis chunking already helps fit the run, but vocab-axis chunking cuts
-  the largest remaining logits allocation at the source.
-- The same fused engine can support SFT, DPO, ORPO, SimPO, IPO, and hinge-style
-  preference losses once it exposes mean log-prob with a custom upstream
-  cotangent.
-- A Metal fused version can avoid extra MLX dispatches and intermediate memory
-  traffic in the hot training path.
-
-The implementation target is not just "chunk logits"; it is a fused log-prob
-head that computes the linear projection, online vocab logsumexp, target-logit
-gather, loss reduction, and backward replay without ever storing full logits.
-
-## Performance focus
-
-### 1. Vocab-axis chunked cross entropy
-
-The current ORPO head optimization is **token-axis chunking**: split the response
-positions into chunks, run the LM head on `[chunk, vocab]`, and rematerialize
-those chunks in the backward. That bounds peak logits by `chunk_tokens * vocab`.
-
-The complementary trick is **vocab-axis chunking**: split the vocabulary instead,
-run the head on `[tokens, vocab_chunk]`, and combine chunks with an
-online-logsumexp recurrence. This is the idea behind **Cut Cross Entropy** and
-Liger-style **Fused Linear Cross Entropy**. This is priority 1 for the ORPO
-performance roadmap.
-
-Exact forward sketch:
-
-1. For each vocab chunk, compute partial logits `h @ W_chunk`.
-2. Gather the target logit only if the target id falls inside this chunk.
-3. Maintain per-token online `logsumexp` across vocab chunks.
-4. At the end, `logp = target_logit - lse`.
-
-Backward reconstructs the softmax chunk-by-chunk from the saved per-token LSE
-and never materializes full `[tokens, vocab]` logits. For Gemma's 262k vocab,
-this is likely the next real moat after segmented backward.
-
-Two implementation levels:
-
-- Pure MLX chunked VJP first: exact, easier to validate, probably enough to fit
-  longer contexts.
-- Metal fused kernel later: fuse quantized head matmul + online LSE + target
-  gather + backward. This is the high-performance version.
-
-### 2. Fold ORPO gradients into the fused CE path
-
-ORPO needs two mean log-probs, `lw` and `lr`, then applies the odds-ratio loss.
-For a fused CE/log-prob kernel, the clean design is to feed the upstream scalar
-gradient for each branch back into the chunked head backward. In other words,
-the fused head should expose "mean logp with custom upstream cotangent", not
-only cross-entropy.
-
-This avoids materializing logits while still letting ORPO, SimPO, and future
-preference losses share the same head engine.
-
-### 3. Prefix sharing in the trainer
-
-`src/train/prefix-shared.ts` already proves the MiniCPM5 shared-prefix idea:
-one forward over `[prompt; chosen; rejected]` with block-sparse masking and
-block-wise RoPE. It is not yet wired into `orpoLoop`.
-
-Next useful path:
-
-- Add an opt-in flag, for example `orpo_prefix_sharing`.
-- Gate it by prompt/response ratio so short-prompt data stays on the simple
-  two-forward path.
-- Wire MiniCPM5 first.
-- Port e4b later, where sliding-window masks, donor-KV reuse, and per-layer
-  input make it trickier.
-
-### 4. Sequence-chunk the MLP sub-block
-
-The MLP is per-token independent. Inside segmented backward, run MLP blocks over
-sequence chunks so the GeGLU/SwiGLU intermediate is `chunk_tokens * d_ff` instead
-of `seq_len * d_ff`.
-
-This stacks with depth segmentation and vocab-axis CE:
-
-`layer segment -> layer -> MLP seq chunk -> vocab chunk`
-
-Attention cannot be chunked this naively, but the MLP can.
-
-### 5. Rotor / budget-aware segment planner
-
-Replace fixed `segment_size` with a planner that chooses segment boundaries from
-estimated activation bytes and a target memory budget. The planner should isolate
-expensive full-attention layers and larger MLP regions instead of treating every
-layer equally.
-
-This is the practical version of rotor/checkpoint scheduling for this codebase:
-less hand tuning, better fit at 8K+.
-
-### 6. Faster flash-attention backward
-
-The custom flash backward is now correctness-useful but too slow. A serious
-performance version should mimic MLX's tuned SDPA style: simdgroup tiling,
-larger threadgroups, and FA2-style backward that stores only per-row LSE and
-recomputes attention probabilities.
-
-This matters most when `ops.sdpa` backward's O(L^2) memory is the blocker.
-
-### 7. Fused LoRA epilogue
-
-Today LoRA adds multiple ops after the quantized base matmul:
-
-`base(x) + scale * ((x @ A) @ B)`
-
-A fused dequant-GEMM + LoRA epilogue would reduce dispatches and intermediate
-traffic. Backward can recompute the cheap rank projection:
-
-- `h = x @ A`
-- `dB = scale * h.T @ dy`
-- `dA = scale * x.T @ (dy @ B.T)`
-- `dx = dy @ W.T + scale * (dy @ B.T) @ A.T`
-
-This is more kernel work, but it benefits SFT, DPO, ORPO, and serving with
-adapters.
-
-### 8. Shape-specialized compile for training graphs
-
-Shapeless compile is mainly required when a dimension changes every replay, such
-as autoregressive decode with a growing KV cache:
-
-`active_kv_len = prompt_len + generated_tokens`
-
-ORPO training is different. Each dataset row already has concrete lengths:
-
-- `P`: prompt length
-- `Rc`: chosen response length
-- `Rr`: rejected response length
-- `T`: prefix-shared length, `P + Rc + Rr`
-
-That means ORPO does not need shapeless compile in the same way decode does. A
-training graph can be compiled for an exact `(P, Rc, Rr)` shape, or for a padded
-length bucket, and cached under a key that includes the model, flags, segment
-plan, token/vocab chunk sizes, and objective mode.
-
-This is especially relevant for e4b prefix-sharing: the proof path should stay
-eager first, but a later L3 speed tier can use shape-specialized compiled graphs
-for known training shapes instead of requiring every primitive to support
-shapeless output-shape inference.
-
-The deeper L3 framing is **model-checksum-specific specialization**. For a fixed
-model snapshot, we know the architecture, layer count, hidden sizes, quantization
-layout, RoPE geometry, vocab, layer types, KV-sharing pattern, and supported
-context budget in advance. It is acceptable to spend much more work up front if
-that buys a faster hot path later.
-
-Practical version:
-
-- Treat the model snapshot checksum as part of the compilation key.
-- Generate/compile/cache a table of exact graph variants for the shapes and
-  modes we care about.
-- Bucket or pad examples into those graph variants when helpful.
-- Fall back to eager only for shapes/modes not in the table.
-
-This is effectively a lookup table of specialized execution plans: more compile
-and analysis cost up front, less branching, shape handling, and dispatch overhead
-at runtime.
-
-### 9. Fixed-shape decode cache as a compile route
-
-Decode is harder because the active KV length grows by one token per step. The
-current shapeless-compile path exists to avoid recompiling for:
-
-`L = 17, 18, 19, 20, ...`
-
-A separate route is to make decode shapes fixed:
-
-1. Preallocate K/V buffers to a known maximum context length.
-2. Write the new token into a dynamic position.
-3. Carry `active_length` / `write_position` as scalar array inputs.
-4. Mask out inactive cache positions, or preferably use an
-   active-length-aware attention kernel that simply does not scan them.
-
-The correctness version is "fixed cache + mask". The performance version is
-"fixed cache + active-length-aware attention"; otherwise early decode steps may
-pay for the full max context even when only a short prefix is active.
-
-There is also a more aggressive specialization route: prebuild a graph table for
-specific active lengths or length buckets. Even if that means many compiled graph
-variants, the trade can be valid for a fixed model/checksum if compile cost is
-paid once and reused. In that regime, the important design question is not "can
-one shapeless graph infer every shape?" but "which finite set of exact shapes and
-runtime states should we specialize for?"
-
-If every op in a chosen route is compile-safe, the whole decode graph can become
-a stable compiled graph. Custom Metal kernels still need to stay outside compiled
-regions unless the route is intentionally eager/custom-kernel-only or MLX compile
-can accept the primitive's output contract.
-
-## Quality and robustness ideas
-
-These are lower priority if the current goal is performance, but they are useful
-once ORPO can run comfortably.
-
-### SimPO
-
-SimPO is a reference-free preference loss that uses the **average log-probability
-of each response** as the reward and adds a target reward margin. It reuses the
-same `lw` / `lr` values ORPO already computes.
-
-Rough shape:
-
-`loss = softplus(-beta * ((lw - lr) - gamma))`
-
-where `gamma` is the desired chosen-vs-rejected margin. It should be cheap to add
-as a sibling objective because the expensive part, response mean log-probs, is
-already implemented.
-
-### IPO
-
-IPO means **Identity Preference Optimization**. It was introduced as a more
-overfit-resistant alternative to the original DPO logistic objective. Instead of
-driving preference probabilities toward hard 0/1 labels indefinitely, IPO uses a
-smoother identity-style objective over preference gaps. The practical point:
-when preference labels are noisy, IPO is often less eager to push margins to
-extremes.
-
-For mlx-bun, IPO is interesting mainly as a loss-family option after the shared
-logp engine exists.
-
-### Hinge / SLiC-style loss
-
-The hinge loss is a margin-ranking objective:
-
-`loss = max(0, margin - reward_gap)`
-
-For preference training, `reward_gap` is usually the chosen score minus rejected
-score. Once the gap exceeds the margin, that pair contributes no gradient. This
-can be useful when you want to stop over-optimizing already-solved pairs and
-focus compute on pairs whose preference margin is still too small.
-
-### Label smoothing and robust weighting
-
-Preference pairs are noisy. Cheap robustness knobs:
-
-- `label_smoothing`: treat a pair as, for example, 95% chosen / 5% rejected
-  rather than a hard label.
-- Per-example `weight`: let curated data carry confidence or score gap.
-- Loss clipping / winsorization: cap extreme high-loss pairs so noisy examples
-  do not dominate gradients.
-- Ambiguity filtering: skip pairs with tiny initial margin or near-identical
-  chosen/rejected responses.
-
-### Better curation
-
-The current UltraFeedback curation mostly extracts rows and filters by length.
-Quality can likely improve by filtering for:
-
-- Large enough chosen-vs-rejected score gap when available.
-- Similar response length, or at least bounded length ratio.
-- High semantic contrast but not unrelated answers.
-- Instruction-following categories if IFEval is the target metric.
-
-Data quality may beat objective cleverness on small ORPO runs.
-
-## Short answer to "token vs vocab"
-
-The thing to look for is **vocab-axis chunked / fused cross entropy**, especially
-under the names **Cut Cross Entropy** or **Liger Fused Linear Cross Entropy**.
-Token chunking slices `[tokens, vocab]` by rows. Vocab chunking slices it by
-columns and combines chunks with online logsumexp so the full logits matrix is
-never live.