mlx-bun 0.0.4 → 0.0.5

package/src/train/trainer.ts

@@ -20,7 +20,8 @@ import type { RuntimeModel } from "../model/factory";
 import { Gemma4Model, type GradCheckpointCtx } from "../model/gemma4";
 import { MiniCPM5Model } from "../model/minicpm5";
 import { setTrainingAttn } from "../model/flash-attention";
-import { SegmentedBackward, SegmentedBackwardGemma4, planSegmentsBySize } from "./segmented";
+import { setFusedGeluTraining } from "../model/fused-geglu-kernel";
+import { SegmentedBackward, SegmentedBackwardGemma4, SegmentedBackwardOrpo, SegmentedBackwardOrpoGemma4, SegmentedBackwardOrpoPrefix, SegmentedBackwardOrpoPrefixGemma4, planSegmentsBySize } from "./segmented";
 import type { LoadedTokenizer } from "../tokenizer";
 import type { ChatTemplate } from "../chat-template";
 import type { Emit } from "../jobs/types";
@@ -30,15 +31,16 @@ import {
 } from "./dataset";
 import {
   buildTrainableLora, attachForTraining, detachTraining, flatParams,
-  saveAdapter, disposeLora, type TrainableLora, type SaveAdapterConfig,
+  saveAdapter, disposeLora, warmStartFromAdapter, type TrainableLora, type SaveAdapterConfig,
 } from "./lora-params";
-import { resolveRanks, DEFAULT_TARGET_MODULES, type RankScaling } from "./rank";
-import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics } from "./loss";
+import { resolveRanks, bitsMapFromModel, readPerLayerKl, DEFAULT_TARGET_MODULES, type RankScaling } from "./rank";
+import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics, orpoLoss, orpoMetrics, type ChunkCtx } from "./loss";
+import { orpoLossPrefixShared, orpoLossPrefixSharedGemma, splitPrefixBatch } from "./prefix-shared";
 import { AdamW, warmupCosineSchedule } from "./optimizer";
 import { writeFileSync } from "node:fs";
 
 export interface TrainConfig {
-  method: "sft" | "dpo";
+  method: "sft" | "dpo" | "orpo";
   rank: number;
   scale: number;
   rankScaling: RankScaling;
@@ -54,10 +56,27 @@ export interface TrainConfig {
   stepsPerEval: number;
   betas: [number, number];
   weightDecay: number;
+  /** LoRA-input dropout (PEFT-style), 0 = off. Training-only regularizer for
+   *  small preference sets. Recompute-safe: the mask is keyed by (step, layer)
+   *  so segmented/checkpointed recompute reproduces it. */
+  loraDropout: number;
+  /** rsLoRA: scale the LoRA update by α/√rank instead of α, so per-layer rank
+   *  scaling (by_bits/by_kl) changes capacity, not step size. Saved in the
+   *  adapter config so inference applies the same per-layer scale. */
+  rsLora: boolean;
+  /** LoRA+: LR multiplier for the B leaves (A stays at the base LR). >1 speeds
+   *  the B-driven early learning (B is zero-init). 1 = off. */
+  loraPlusRatio: number;
   /** Gradient checkpointing: recompute layer activations in the backward pass
    *  to bound memory at long context (mlx-lm grad_checkpoint; default in optiq).
    *  Numerically identical to off — pure memory↔compute trade. */
   gradCheckpoint: boolean;
+  /** Intra-layer MLP split (lever 4): with gradCheckpoint on (Gemma4 only),
+   *  checkpoint the attention and MLP sub-blocks separately with the post-attn
+   *  residual as the boundary, so the backward recompute holds max(attn,MLP)+hMid
+   *  instead of the whole layer. Numerically identical to gradCheckpoint alone —
+   *  a finer memory↔compute split. No effect unless gradCheckpoint is on. */
+  mlpSplit: boolean;
   /** Segmented backward: layers per segment (0 = off). Streams the SFT backward
    *  segment-by-segment so only one segment's activations are live at a time —
    *  beats per-layer checkpointing at long context, where naive checkpointing
@@ -75,6 +94,43 @@ export interface TrainConfig {
   dpoBeta: number;
   dpoWarmupIters: number;
   dpoLrSchedule: "constant" | "cosine";
+  // ORPO (reference-free; method-prefixed knobs, mirroring the dpo trio).
+  /** λ — weights ONLY the odds-ratio term (the SFT-NLL term stays unweighted). */
+  orpoLambda: number;
+  orpoWarmupIters: number;
+  orpoLrSchedule: "constant" | "cosine";
+  /** Token-chunk size for the ORPO head (0 = off). >0 computes the B=1
+   *  response-only LM head in token-chunks, each rematerialized in the backward,
+   *  bounding the dominant `[M, vocab]` logits term to `[chunkSize, vocab]`.
+   *  Exact (numerically identical). Non-segmented path only (segmentSize takes
+   *  precedence). See docs/design/orpo-training.md → chunked fused log-prob. */
+  orpoChunkSize: number;
+  /** Use the FUSED linear cross-entropy head (Liger/CCE-style): one CustomVjp with
+   *  an analytic softmax−onehot backward — no autograd through the head, no
+   *  retained `[M,vocab]` logits in either direction. Composes with
+   *  `orpoChunkSize` as the token-chunk size (defaults to 512 when unset). B=1,
+   *  non-segmented path only; Gemma + MiniCPM5 heads. Value bit-exact vs the
+   *  full-logits head, grads in the bf16 class. See docs/design/orpo-training.md
+   *  → fused linear-CE head. */
+  orpoFusedCe: boolean;
+  /** Route the fused linear-CE head through the flash-CCE Metal kernel: in-kernel
+   *  quantized logits + online softmax (fwd) and dh accumulation (bwd) — neither
+   *  `[M,vocab]` nor a dequantized head touches HBM. Lowest memory + fastest on
+   *  large vocab (e4b 262k / CPM 130k). Implies the fused head; the Apple-CCE coeff
+   *  filter is on by default in the backward. B=1, non-segmented path. */
+  orpoFlashCe: boolean;
+  /** Shared prompt-prefix ORPO (lever 7): one forward over [prompt; chosen; rejected]
+   *  with a block-sparse mask + block-wise RoPE, so the shared prompt is encoded ONCE
+   *  (token cost 2(P+R) → P+2R). B=1, non-segmented path; MiniCPM5 + Gemma4 (e4b).
+   *  Falls back to the two-forward orpoLoss for rows whose chosen/rejected prompts
+   *  differ. Composes with orpoFlashCe/orpoFusedCe (each branch routes through the
+   *  [M,V]-free head). See src/train/prefix-shared.ts. */
+  orpoPrefixShared: boolean;
+  /** Warm-start: path to an existing adapter dir (with adapters.safetensors) whose
+   *  LoRA weights initialize this run instead of the random/zero init. Continues
+   *  training from a checkpoint's weights; the optimizer state and LR schedule restart.
+   *  Rank/targets must match the checkpoint. "" = off (fresh init). */
+  warmStartAdapter: string;
   // output
   adapterPath: string;
   baseModel: string;
@@ -103,6 +159,10 @@ export const DEFAULT_TRAIN_CONFIG: TrainConfig = {
   stepsPerEval: 50,
   betas: [0.9, 0.999],
   weightDecay: 0.01,
+  loraDropout: 0.0, // off by default
+  rsLora: false, // off by default (opt-in; recommended when rank_scaling != constant)
+  loraPlusRatio: 1.0, // 1 = off
+  mlpSplit: false, // default off; only effective with gradCheckpoint (Gemma4)
   gradCheckpoint: false, // default off until validated; flip to true (optiq default) after
   segmentSize: 0, // default off; >0 enables segmented backward (MiniCPM5 SFT B=1)
   saveCheckpoints: false, // off by default; scripts opt in (keep-all-checkpoints + metrics.json)
@@ -110,6 +170,15 @@ export const DEFAULT_TRAIN_CONFIG: TrainConfig = {
   dpoBeta: 0.1,
   dpoWarmupIters: 0,
   dpoLrSchedule: "cosine",
+
+  orpoLambda: 0.1, // paper's primary value + TRL default; weights only L_OR
+  orpoWarmupIters: 0,
+  orpoLrSchedule: "cosine",
+  orpoChunkSize: 0, // default off; >0 = token-chunked head (bounds the [M,vocab] term)
+  orpoFusedCe: false, // default off; true = fused linear-CE head (analytic backward, no [M,vocab])
+  orpoFlashCe: false, // default off; true = flash-CCE Metal-kernel head (implies fused)
+  orpoPrefixShared: false, // default off; true = shared prompt-prefix single forward (B=1)
+  warmStartAdapter: "", // default off; path to an adapter dir to warm-start LoRA weights from
   adapterPath: "adapters",
   baseModel: "",
 };
@@ -127,16 +196,51 @@ export async function trainLora(
 ): Promise<TrainResult> {
   emit({ type: "stage", stage: "setup", progress: 0.02, message: "resolving ranks" });
 
+  // Mixed-precision rank scaling: feed the per-layer bits (from the loaded
+  // model's quant specs) and, for by_kl, the recorded KL sensitivities
+  // (optiq_metadata.json). Without these maps, by_bits/by_kl silently fell back
+  // to uniform rank — so the "give the sensitive (8-bit) layers more adapter
+  // capacity" policy (optiq: one sensitivity signal, two optimizations — the
+  // bit assignment AND the LoRA rank) never actually fired.
+  const bitsMap = bitsMapFromModel(model);
+  const klMap = cfg.rankScaling === "by_kl" ? await readPerLayerKl(cfg.baseModel) : undefined;
   const ranks = resolveRanks(model, {
     rank: cfg.rank,
     rankScaling: cfg.rankScaling,
     targetModules: cfg.targetModules,
     numLayers: cfg.numLayers,
+    bitsMap,
+    klMap,
   });
   const appliedRanks: Record<string, number> = Object.fromEntries(ranks);
 
-  const lora = buildTrainableLora(model, ranks, cfg.scale, cfg.seed);
+  // Surface the rank spread so mixed-precision scaling is observable (e.g.
+  // "r8×140 r16×28" means 28 sensitive layers got the wider adapter).
+  const rankCounts = new Map<number, number>();
+  for (const r of ranks.values()) rankCounts.set(r, (rankCounts.get(r) ?? 0) + 1);
+  const rankSpread = [...rankCounts.entries()]
+    .sort((a, b) => a[0] - b[0])
+    .map(([r, n]) => `r${r}×${n}`)
+    .join(" ");
+  emit({ type: "stage", stage: "setup", progress: 0.03,
+    message: `ranks (${cfg.rankScaling}): ${rankSpread}` });
+
+  const lora = buildTrainableLora(model, ranks, cfg.scale, cfg.seed, cfg.rsLora);
+  // Warm-start: continue from a saved adapter/checkpoint's weights (optimizer + LR
+  // schedule restart fresh). Must run before the optimizer is built (below) so it
+  // tracks the loaded leaves. Rank/targets must match the checkpoint.
+  if (cfg.warmStartAdapter) {
+    const n = warmStartFromAdapter(lora, cfg.warmStartAdapter);
+    emit({ type: "stage", stage: "setup", progress: 0.035,
+      message: `warm-start: loaded ${n} LoRA targets from ${cfg.warmStartAdapter} (optimizer + LR schedule restart)` });
+  }
   attachForTraining(model, lora, "train");
+  model.loraState.dropoutRate = cfg.loraDropout; // per-step seed set inside the loop
+  // Training-mode fused GeGLU (Gemma e4b): kernel forward + hand-derived vjp, so
+  // autograd flows and the backward recomputes the gelu from the primal instead
+  // of retaining the spelled-out intermediates. No-op for non-Gemma models.
+  // Cleared in the finally below.
+  setFusedGeluTraining(true);
 
   const saveCfg: SaveAdapterConfig = {
     rank: cfg.rank,
@@ -146,6 +250,7 @@ export async function trainLora(
     numLayers: cfg.numLayers,
     method: cfg.method,
     baseModel: cfg.baseModel,
+    rsLora: cfg.rsLora,
   };
 
   // Keep-ALL-checkpoints + a durable metrics record. At each eval step we
@@ -158,7 +263,11 @@ export async function trainLora(
   // metrics, and the live LoRA leaves are the current step's weights at emit
   // time, so saving here captures exactly that checkpoint.
   const trainHistory: { step: number; loss: number }[] = [];
-  const valHistory: { step: number; loss: number; checkpoint: string | null }[] = [];
+  // For preference methods (dpo/orpo) the val MARGIN (mean chosen−rejected
+  // log-odds) and accuracy are the meaningful early-stop signals — val loss is a
+  // weak proxy. Capture them when the loop emits them so metrics.json lets you
+  // pick the best-MARGIN checkpoint (overfit shows as train↑ / val-margin flat).
+  const valHistory: { step: number; loss: number; margin?: number; accuracy?: number; checkpoint: string | null }[] = [];
   const startedAt = Date.now();
   const collect: Emit = (e) => {
     if (e.type === "metric" && e.kind === "train") {
@@ -172,7 +281,8 @@ export async function trainLora(
         emit({ type: "stage", stage: "checkpoint", progress: e.step / cfg.iters,
                message: `checkpoint ${tag}` });
       }
-      valHistory.push({ step: e.step, loss: e.loss, checkpoint });
+      const ev = e as { step: number; loss: number; margin?: number; accuracy?: number };
+      valHistory.push({ step: ev.step, loss: ev.loss, margin: ev.margin, accuracy: ev.accuracy, checkpoint });
     }
     emit(e);
   };
@@ -181,7 +291,9 @@ export async function trainLora(
     const result =
       cfg.method === "dpo"
         ? await dpoLoop(model, tok, tmpl, dataDir, cfg, lora, collect)
-        : await sftLoop(model, tok, tmpl, dataDir, cfg, lora, collect);
+        : cfg.method === "orpo"
+          ? await orpoLoop(model, tok, tmpl, dataDir, cfg, lora, collect)
+          : await sftLoop(model, tok, tmpl, dataDir, cfg, lora, collect);
 
     // Final save (last adapter).
     detachTraining(model, lora);
@@ -193,18 +305,24 @@ export async function trainLora(
     if (cfg.saveCheckpoints) {
       let bestVal: { step: number; loss: number; checkpoint: string | null } | null = null;
       for (const v of valHistory) if (bestVal === null || v.loss < bestVal.loss) bestVal = v;
+      // Best by MARGIN too (preference methods) — the better checkpoint signal.
+      let bestMargin: typeof valHistory[number] | null = null;
+      for (const v of valHistory)
+        if (v.margin !== undefined && (bestMargin === null || v.margin > (bestMargin.margin ?? -Infinity))) bestMargin = v;
       writeFileSync(`${cfg.adapterPath}/metrics.json`, JSON.stringify({
         config: {
           method: cfg.method, rank: cfg.rank, scale: cfg.scale,
           learningRate: cfg.learningRate, maxSeqLen: cfg.maxSeqLen,
           segmentSize: cfg.segmentSize, iters: cfg.iters,
           numLayers: cfg.numLayers, baseModel: cfg.baseModel,
+          orpoLambda: cfg.method === "orpo" ? cfg.orpoLambda : undefined,
         },
         wallSeconds: (Date.now() - startedAt) / 1000,
         peakGb: peakMemory() / 1e9,
         finalTrainLoss: trainHistory.at(-1)?.loss ?? null,
         finalValLoss: valHistory.at(-1)?.loss ?? null,
         bestVal,
+        bestMargin, // null for SFT (no margin); the recommended pick for dpo/orpo
         valTrajectory: valHistory,
         trainLosses: trainHistory,
       }, null, 2));
@@ -220,6 +338,7 @@ export async function trainLora(
     });
     return { adapterPath: cfg.adapterPath, appliedRanks, numIters: result.numIters };
   } finally {
+    setFusedGeluTraining(false);
     disposeLora(lora);
   }
 }
@@ -290,22 +409,29 @@ async function sftLoop(
   // Checkpoints are created per forward and disposed after each step's backward.
   let ckptCtx: GradCheckpointCtx | null = null;
   if (!useSegmented && cfg.gradCheckpoint && model instanceof Gemma4Model) {
+    // Partition each layer's LoRA into the attention sub-block (self_attn.*) and
+    // the MLP sub-block (everything else in the layer: gate/up/down_proj +
+    // per_layer_*) so the split-MLP checkpoint can wrap them independently.
     const byLayer: GradCheckpointCtx["byLayer"] = new Map();
     for (const t of lora.targets) {
       const m = t.modulePath.match(/\.layers\.(\d+)\./);
       if (!m) continue;
       const li = Number(m[1]);
-      (byLayer.get(li) ?? byLayer.set(li, []).get(li)!).push(t.lw);
+      const ll = byLayer.get(li) ?? byLayer.set(li, { attn: [], mlp: [] }).get(li)!;
+      (t.modulePath.includes(".self_attn.") ? ll.attn : ll.mlp).push(t.lw);
     }
-    ckptCtx = { byLayer, keepAlive: [] };
+    ckptCtx = { byLayer, splitMlp: cfg.mlpSplit, keepAlive: [] };
     model.gradCkpt = ckptCtx;
+    if (cfg.mlpSplit) emit({ type: "stage", stage: "setup", progress: 0.04,
+      message: "grad checkpoint: intra-layer attn/MLP split" });
   }
 
-  // Training attention. Default = mlx fused SDPA (ops.sdpa): verified correct
-  // dQ/dK/dV vs a hand-rolled reference (scripts/sdpa-vs-manual.ts: 0.00%) and
-  // fast (one fused kernel). flash is opt-in via MLX_BUN_TRAIN_ATTN=flash for an
-  // O(L) memory backward at very long context, BUT its dK kernel currently has
-  // a bug (~100% off) and is ~30x slower here — so it is NOT the default.
+  // Training attention. Default = mlx fused SDPA (ops.sdpa): exact dQ/dK/dV
+  // (finite-difference verified) and fast, but O(L²) backward memory. The
+  // hand-rolled flash kernel is an opt-in (MLX_BUN_TRAIN_ATTN=flash) O(L)-memory
+  // path for memory-bound long context — now correct (its dK-transpose and
+  // dQ-causal-barrier bugs were fixed; FD-validated in
+  // scripts/experiments/flash-fd-check.ts), though ~30× slower than ops.sdpa.
   // Cleared in finally.
   if (process.env.MLX_BUN_TRAIN_ATTN === "flash") setTrainingAttn("flash");
 
@@ -318,22 +444,27 @@ async function sftLoop(
 
   try {
     for (let step = 1; step <= cfg.iters; step++) {
-      currentBatch = batches.next().value as SftBatch;
-      const tokensThisStep = countResponseTokens(currentBatch);
-
       if (MEM_LOG) resetPeakMemory();
-      const { value, grads } = segmented
-        ? segmented.step(currentBatch!)
-        : vag!.apply(flatParams(lora));
-      // Evaluate the loss AND grads in ONE mlx_eval so the backward shares the
-      // forward pass. Evaluating `value` on its own first (e.g. via toFloat32)
-      // materializes then frees the forward, forcing the backward to recompute
-      // and hold ALL layer activations — which silently defeats gradient
-      // checkpointing (verified: separate-eval ~0% memory savings, combined-eval
-      // restores it). Keep this single combined eval before reading either.
-      evalAll([value, ...grads]);
-      const lossVal = value.toFloat32()[0]!;
-      value.dispose();
+
+      // One optimizer step over cfg.gradAccumSteps micro-batches (pass-through
+      // when 1). afterMicroEval frees each micro-batch's gradient-checkpoint
+      // closures right after its backward so they don't pile up across
+      // accumulation.
+      let tokensThisStep = 0;
+      const { loss: lossVal, grads } = accumulateStep(
+        cfg.gradAccumSteps,
+        () => {
+          currentBatch = batches.next().value as SftBatch;
+          tokensThisStep += countResponseTokens(currentBatch);
+          return segmented ? segmented.step(currentBatch!) : vag!.apply(flatParams(lora));
+        },
+        ckptCtx
+          ? () => {
+              for (const ck of ckptCtx!.keepAlive) ck.dispose();
+              ckptCtx!.keepAlive.length = 0;
+            }
+          : undefined,
+      );
 
       const gradNorm = globalNorm(grads); // grads already materialized above
       opt.step(grads);
@@ -345,12 +476,6 @@ async function sftLoop(
       } else {
         clearCache();
       }
-      if (ckptCtx) {
-        // Grads are materialized (globalNorm/evalState), so the recompute is
-        // done — free this step's checkpoint closures before the next forward.
-        for (const ck of ckptCtx.keepAlive) ck.dispose();
-        ckptCtx.keepAlive.length = 0;
-      }
 
       if (step % cfg.stepsPerReport === 0 || step === 1) {
         const elapsed = (Date.now() - t0) / 1000;
@@ -467,28 +592,35 @@ async function dpoLoop(
 
   try {
     for (let step = 1; step <= cfg.iters; step++) {
-      currentBatch = batches.next().value as DpoBatch;
       if (schedule) opt.lr = schedule(step);
 
-      const refs = dpoRefLogps(model, lora, currentBatch);
-      refChosen = refs.refChosen;
-      refRejected = refs.refRejected;
-
-      const { value, grads } = vag.apply(flatParams(lora));
-      const lossVal = value.toFloat32()[0]!;
-      value.dispose();
+      // One optimizer step over cfg.gradAccumSteps micro-batches (pass-through
+      // when 1). Reference log-probs are recomputed per micro-batch OUTSIDE the
+      // policy value_and_grad and read as detached constants inside it;
+      // afterMicroEval frees them once that micro-batch's grads are materialized.
+      const { loss: lossVal, grads } = accumulateStep(
+        cfg.gradAccumSteps,
+        () => {
+          currentBatch = batches.next().value as DpoBatch;
+          const refs = dpoRefLogps(model, lora, currentBatch);
+          refChosen = refs.refChosen;
+          refRejected = refs.refRejected;
+          return vag.apply(flatParams(lora));
+        },
+        () => {
+          refChosen?.dispose();
+          refRejected?.dispose();
+          refChosen = refRejected = null;
+        },
+      );
 
       const gradNorm = globalNorm(grads);
       opt.step(grads);
       opt.evalState();
-
-      refChosen.dispose();
-      refRejected.dispose();
-      refChosen = refRejected = null;
       clearCache();
 
       if (step % cfg.stepsPerReport === 0 || step === 1) {
-        const m = dpoMetrics(model, lora, currentBatch, cfg.dpoBeta);
+        const m = dpoMetrics(model, lora, currentBatch!, cfg.dpoBeta);
         const elapsed = (Date.now() - t0) / 1000;
         emit({
           type: "metric",
@@ -529,8 +661,250 @@ async function dpoLoop(
   } finally {
     vag.dispose();
     opt.dispose();
-    refChosen?.dispose();
-    refRejected?.dispose();
+    // Cast: refs are only ever assigned inside the accumulateStep closures, so
+    // TS's flow analysis narrows these to their `null` initializer here.
+    (refChosen as MlxArray | null)?.dispose();
+    (refRejected as MlxArray | null)?.dispose();
+  }
+  return { numIters: cfg.iters };
+}
+
+// ---------------------------------------------------------------------------
+// ORPO loop (reference-free: 2 forwards/step, no reference pass)
+// ---------------------------------------------------------------------------
+
+async function orpoLoop(
+  model: RuntimeModel,
+  tok: LoadedTokenizer,
+  tmpl: ChatTemplate,
+  dataDir: string,
+  cfg: TrainConfig,
+  lora: TrainableLora,
+  emit: Emit,
+): Promise<{ numIters: number }> {
+  emit({ type: "stage", stage: "data", progress: 0.05, message: "loading preference dataset" });
+  const train = await loadDpoDataset(`${dataDir}/train.jsonl`, tok, tmpl, cfg.maxSeqLen);
+  const valid = (await fileExists(`${dataDir}/valid.jsonl`))
+    ? await loadDpoDataset(`${dataDir}/valid.jsonl`, tok, tmpl, cfg.maxSeqLen)
+    : [];
+
+  const params = flatParams(lora);
+  // LoRA+: B leaves (the second half of flatParams) get a higher LR than A.
+  const nT = lora.targets.length;
+  const lrScale = cfg.loraPlusRatio !== 1
+    ? params.map((_, i) => (i < nT ? 1 : cfg.loraPlusRatio))
+    : undefined;
+  if (lrScale) emit({ type: "stage", stage: "setup", progress: 0.04,
+    message: `LoRA+ : B leaves at ${cfg.loraPlusRatio}× LR` });
+  const opt = new AdamW(
+    params,
+    { lr: cfg.learningRate, betas: cfg.betas, eps: 1e-8, weightDecay: cfg.weightDecay, lrScale },
+    (i, p) => writeParam(lora, i, p),
+  );
+
+  const schedule =
+    cfg.orpoLrSchedule === "cosine" || cfg.orpoWarmupIters > 0
+      ? warmupCosineSchedule(cfg.learningRate, cfg.orpoWarmupIters, cfg.iters)
+      : null;
+
+  // Segmented backward for ORPO: stream chosen+rejected backwards segment-by-segment.
+  // MiniCPM5 + Gemma4 (e4b) — same model coverage as SFT segmented. B=1 only.
+  const useSegmented = cfg.segmentSize > 0;
+  // Prefix-sharing + segmented backward COMPOSE (M3-composition): the segmented
+  // backward streams the ONE prefix-shared concat forward [prompt; chosen; rejected]
+  // (block-sparse mask via PrefixSharedCache + block-wise RoPE via the prefix plan)
+  // segment-by-segment — so the prompt-encode-once saving holds AT long seq. Composed
+  // for BOTH MiniCPM5 and Gemma4/e4b (the e4b variant threads donor-KV + the
+  // logical-position sliding-window prefix mask through the segment boundaries).
+  const useSegmentedPrefix = cfg.orpoPrefixShared && useSegmented;
+  if (cfg.orpoPrefixShared && !(model instanceof MiniCPM5Model || model instanceof Gemma4Model))
+    throw new Error("orpoPrefixShared is only wired for MiniCPM5 and Gemma4");
+  // Fused linear-CE head INSIDE the segmented backward (orpo_fused_ce): bounds the
+  // head term to [chunk,V] alongside the per-segment layer savings. 0 = the
+  // full-[M,V] responseOnlyLogpMean head (unchanged). Reuses orpoChunkSize as the
+  // token-chunk (default 512), same as the non-segmented path.
+  // "flash implies fused": orpoFlashCe alone must still bound the segmented head to
+  // [chunk,V]. The non-prefix segmented class has no flash path, so flash routes
+  // through its fused/checkpoint chunked head (boundedHeadFromHidden) — never the
+  // full-[M,V] responseOnlyLogpMean, which would reintroduce [M,vocab] on e4b.
+  const segUsesFusedHead = cfg.orpoFusedCe || cfg.orpoFlashCe;
+  const segFusedChunk = segUsesFusedHead ? (cfg.orpoChunkSize > 0 ? cfg.orpoChunkSize : 512) : 0;
+  // The prefix segmented head gathers response positions then routes per-branch
+  // through the [M,V]-free flash/fused head (branchLogpMeanGathered, which takes a
+  // ChunkCtx). Its own sink is managed per-step inside the class.
+  const segPrefixChunk: ChunkCtx | undefined =
+    cfg.orpoFusedCe || cfg.orpoFlashCe
+      ? { chunkSize: cfg.orpoChunkSize > 0 ? cfg.orpoChunkSize : 512, fused: true, flash: cfg.orpoFlashCe, sink: [] }
+      : undefined;
+  let segmentedOrpo: SegmentedBackwardOrpo | SegmentedBackwardOrpoGemma4 | null = null;
+  let segmentedOrpoPrefix: SegmentedBackwardOrpoPrefix | SegmentedBackwardOrpoPrefixGemma4 | null = null;
+  if (useSegmented) {
+    const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
+    if (useSegmentedPrefix) {
+      // The composed prefix-shared segmented backward (MiniCPM5 or Gemma4/e4b), plus a
+      // plain segmented-ORPO fallback for rows whose chosen/rejected prompts differ
+      // (splitPrefixBatch returns null -> two-forward segmented step).
+      if (model instanceof MiniCPM5Model) {
+        segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefix(model, lora, ranges, cfg.orpoLambda, segPrefixChunk);
+        segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
+      } else {
+        segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefixGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segPrefixChunk);
+        segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk);
+      }
+      emit({ type: "stage", stage: "setup", progress: 0.04,
+        message: `orpo segmented + prefix-share (${model instanceof Gemma4Model ? "e4b" : "MiniCPM5"}): single concat forward streamed over ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers (two-forward fallback on prompt mismatch)` });
+    } else if (model instanceof MiniCPM5Model)
+      segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
+    else if (model instanceof Gemma4Model)
+      segmentedOrpo = new SegmentedBackwardOrpoGemma4(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
+    else
+      throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
+    if (!useSegmentedPrefix) emit({ type: "stage", stage: "setup", progress: 0.04,
+      message: `orpo segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers${segFusedChunk > 0 ? ` + fused linear-CE head (${segFusedChunk}/chunk)` : ""}` });
+  }
+
+  // Token-chunked head (non-segmented path only): the per-chunk checkpoints are
+  // collected here and disposed after each micro-step's grads are eval'd (the
+  // afterMicroEval hook below) — that is when their backward recompute has run.
+  const chunkSink: Array<{ dispose(): void }> = [];
+  // Fused linear-CE head reuses orpoChunkSize as the token-chunk (default 512 when
+  // unset); otherwise the Checkpoint token-chunked head fires only when
+  // orpoChunkSize > 0. Both are non-segmented-path only.
+  const fusedChunkSize = cfg.orpoChunkSize > 0 ? cfg.orpoChunkSize : 512;
+  const chunk: ChunkCtx | undefined = useSegmented
+    ? undefined
+    : cfg.orpoFusedCe || cfg.orpoFlashCe
+      ? { chunkSize: fusedChunkSize, fused: true, flash: cfg.orpoFlashCe, sink: chunkSink }
+      : cfg.orpoChunkSize > 0
+        ? { chunkSize: cfg.orpoChunkSize, sink: chunkSink }
+        : undefined;
+  if (chunk) emit({ type: "stage", stage: "setup", progress: 0.04,
+    message: chunk.flash
+      ? `orpo flash-CCE (Metal kernel) head: ${chunk.chunkSize} tokens/chunk`
+      : chunk.fused
+        ? `orpo fused linear-CE head: ${chunk.chunkSize} tokens/chunk`
+        : `orpo chunked head: ${cfg.orpoChunkSize} tokens/chunk` });
+  if (cfg.orpoPrefixShared && !useSegmentedPrefix) emit({ type: "stage", stage: "setup", progress: 0.04,
+    message: `orpo prefix-sharing: single forward over [prompt; chosen; rejected] (two-forward fallback on prompt mismatch)` });
+  const disposeChunkSink = () => { for (const c of chunkSink) c.dispose(); chunkSink.length = 0; };
+
+  // No reference forward (the whole point of ORPO) — the closure just builds the
+  // monolithic L_NLL + λ·L_OR over the active batch against the swapped primals.
+  let currentBatch: DpoBatch | null = null;
+  const vag = useSegmented
+    ? null
+    : new ValueAndGrad((primals) => {
+        const saved = swapPrimals(lora, primals);
+        try {
+          // Prefix-sharing (B=1): one forward over [prompt; chosen; rejected] when the
+          // row's chosen/rejected share an identical prompt; else fall back two-forward.
+          if (cfg.orpoPrefixShared) {
+            const split = splitPrefixBatch(currentBatch!);
+            if (split) {
+              if (model instanceof MiniCPM5Model)
+                return orpoLossPrefixShared(model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk);
+              return orpoLossPrefixSharedGemma(model as Gemma4Model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk);
+            }
+            notePrefixFallback();
+          }
+          return orpoLoss(model, currentBatch!, cfg.orpoLambda, chunk);
+        } finally {
+          restorePrimals(lora, saved);
+        }
+      }, params.map((_, i) => i));
+
+  const padId = tok.eosTokenId ?? 0;
+  const batches = iterateDpoBatches(train, cfg.batchSize, cfg.seed, true, padId);
+  const t0 = Date.now();
+
+  let dropoutCounter = 0;
+  // Prefix-share falls back to the two-forward path for any row whose chosen/rejected
+  // prompts aren't byte-identical (splitPrefixBatch → null). Log the FIRST fallback +
+  // count them, so a silent loss of the prompt-encode-once saving is visible.
+  let prefixFallbacks = 0;
+  const notePrefixFallback = () => {
+    prefixFallbacks++;
+    if (prefixFallbacks === 1) emit({ type: "stage", stage: "train", progress: 0,
+      message: `prefix-share: row prompt mismatch → two-forward fallback for this row (saving lost only on mismatched rows; counting…)` });
+  };
+  try {
+    for (let step = 1; step <= cfg.iters; step++) {
+      if (schedule) opt.lr = schedule(step);
+
+      // One optimizer step over cfg.gradAccumSteps micro-batches (pass-through
+      // when 1) — the long-context lever for a larger effective batch without
+      // the B>1 activation memory (see docs/design/orpo-training.md, Batching).
+      const { loss: lossVal, grads } = accumulateStep(cfg.gradAccumSteps, () => {
+        currentBatch = batches.next().value as DpoBatch;
+        // Per-micro-step dropout seed: constant across this micro-step's forward
+        // AND its backward recompute (segmented/chunked), so the mask reproduces.
+        if (cfg.loraDropout > 0) model.loraState.dropoutSeed = ++dropoutCounter;
+        // Segmented + prefix-share: stream the single concat forward when the row's
+        // chosen/rejected share an identical prompt; else fall back to the plain
+        // two-forward segmented step (segmentedOrpo, built alongside).
+        if (segmentedOrpoPrefix) {
+          const split = splitPrefixBatch(currentBatch!);
+          if (split) return segmentedOrpoPrefix.stepPrefix(split.promptIds, split.chosenResp, split.rejectedResp);
+          notePrefixFallback();
+          return segmentedOrpo!.step(currentBatch!);
+        }
+        return segmentedOrpo ? segmentedOrpo.step(currentBatch!) : vag!.apply(flatParams(lora));
+      }, disposeChunkSink); // free this micro-step's head checkpoints once its grads are eval'd
+
+      model.loraState.dropoutSeed = null; // metrics/val below run dropout-free
+      const gradNorm = globalNorm(grads);
+      opt.step(grads);
+      opt.evalState();
+      clearCache();
+
+      if (step % cfg.stepsPerReport === 0 || step === 1) {
+        // Diagnostics recompute the two forwards (like dpoMetrics) — cheap at
+        // B=1 (response-only head, no full logits) and gated by stepsPerReport.
+        // Reusing the value_and_grad forward is a shared-with-DPO follow-on.
+        const m = orpoMetrics(model, currentBatch!, cfg.orpoLambda);
+        const elapsed = (Date.now() - t0) / 1000;
+        emit({
+          type: "metric",
+          kind: "train",
+          step,
+          loss: lossVal,
+          grad_norm: gradNorm,
+          learning_rate: opt.lr,
+          accuracy: m.accuracy,
+          margin: m.margin,
+          tokens_per_sec: 0,
+          progress: step / cfg.iters,
+          message: `step ${step}/${cfg.iters} loss=${lossVal.toFixed(4)} nll=${m.nll.toFixed(4)} or=${m.or.toFixed(4)} acc=${m.accuracy.toFixed(2)} margin=${m.margin.toFixed(3)}`,
+        });
+        clearCache();
+      }
+
+      if (valid.length > 0 && step % cfg.stepsPerEval === 0) {
+        let vl = 0, va = 0, vm = 0, vn = 0;
+        for (const ex of valid) {
+          const vb: DpoBatch = {
+            chosenIds: [ex.chosenIds], rejectedIds: [ex.rejectedIds],
+            chosenMask: [ex.chosenMask], rejectedMask: [ex.rejectedMask],
+          };
+          const m = orpoMetrics(model, vb, cfg.orpoLambda);
+          vl += m.loss; va += m.accuracy; vm += m.margin; vn++;
+          clearCache();
+        }
+        emit({
+          type: "metric", kind: "val", step,
+          loss: vn ? vl / vn : 0,
+          accuracy: vn ? va / vn : 0,
+          margin: vn ? vm / vn : 0,
+          progress: step / cfg.iters,
+        });
+      }
+    }
+  } finally {
+    disposeChunkSink(); // defensive: free any head checkpoints a throw left behind
+    segmentedOrpoPrefix?.dispose();
+    segmentedOrpo?.dispose();
+    vag?.dispose();
+    opt.dispose();
   }
   return { numIters: cfg.iters };
 }
@@ -570,6 +944,89 @@ function writeParam(lora: TrainableLora, i: number, p: MlxArray): void {
   else lora.targets[i - n]!.lw.b = p;
 }
 
+/** Run one optimizer step's worth of gradient over `accumSteps` micro-batches:
+ *  backward each micro-batch, mean-accumulate the per-micro grads (each scaled
+ *  by 1/accumSteps), and return the mean loss + summed grads — a drop-in for the
+ *  single value_and_grad the loops used to call inline. The caller passes the
+ *  returned grads to opt.step (which takes ownership and disposes them).
+ *
+ *  accumSteps <= 1 is a pass-through identical to the pre-accumulation path (one
+ *  micro-batch, no scaling) so the proven B=1 loops stay byte-for-byte unchanged.
+ *  For accumSteps > 1 the running accumulator (and that micro-batch's loss) is
+ *  eval'd before the next micro-batch's forward builds, so only ONE micro-batch's
+ *  activations are ever live — peak memory does NOT grow with accumSteps (the
+ *  whole point: a larger effective batch without the B>1 activation memory).
+ *
+ *  The combined `evalAll([value, ...grads])` per micro-batch is deliberate: a
+ *  separate eval of `value` first would materialize then free the forward,
+ *  forcing the backward to recompute and hold every layer's activations — which
+ *  silently defeats gradient checkpointing (verified). Keep loss + grads in one
+ *  eval.
+ *
+ *  `afterMicroEval` runs once per micro-batch right after its grads are
+ *  materialized (its backward is complete), in BOTH the pass-through and
+ *  accumulating paths — it frees per-micro-batch state the loss closure captured
+ *  (gradient-checkpoint closures for SFT, reference log-probs for DPO) so that
+ *  state does not pile up across accumulation.
+ *
+ *  `micro` returns the micro-batch loss scalar + grads in flatParams order
+ *  ([...A, ...B] over targets); it transfers ownership of both to this helper,
+ *  which disposes the per-micro value and grads. */
+export function accumulateStep(
+  accumSteps: number,
+  micro: () => { value: MlxArray; grads: MlxArray[] },
+  afterMicroEval?: () => void,
+): { loss: number; grads: MlxArray[] } {
+  if (accumSteps <= 1) {
+    const { value, grads } = micro();
+    evalAll([value, ...grads]);
+    const loss = value.toFloat32()[0]!;
+    value.dispose();
+    afterMicroEval?.();
+    return { loss, grads };
+  }
+
+  const scale = 1 / accumSteps;
+  let acc: MlxArray[] | null = null;
+  let lossSum = 0;
+  try {
+    for (let i = 0; i < accumSteps; i++) {
+      const { value, grads } = micro();
+      if (acc === null) {
+        acc = grads.map((g) => {
+          const s = ops.mulScalar(g, scale);
+          g.dispose();
+          return s;
+        });
+      } else {
+        const prev = acc;
+        acc = grads.map((g, j) => {
+          const s = ops.mulScalar(g, scale);
+          const sum = ops.add(prev[j]!, s);
+          s.dispose();
+          prev[j]!.dispose();
+          g.dispose();
+          return sum;
+        });
+      }
+      // Eval the running accumulator (+ this micro-batch's loss) now so the
+      // micro-batch's forward/backward graph frees before the next one builds —
+      // this is what keeps peak flat across accumSteps.
+      evalAll([value, ...acc]);
+      lossSum += value.toFloat32()[0]!;
+      value.dispose();
+      afterMicroEval?.();
+      clearCache();
+    }
+    const grads = acc!;
+    acc = null; // ownership handed to the caller
+    return { loss: lossSum / accumSteps, grads };
+  } finally {
+    // Non-null only if a micro-batch threw mid-accumulation (nulled on success).
+    if (acc) for (const a of acc) a.dispose();
+  }
+}
+
 /** L2 norm over all grads (also forces eval). */
 function globalNorm(grads: MlxArray[]): number {
   const partial: MlxArray[] = grads.map((g) => {