mlx-bun 0.0.9 → 0.0.10

package/src/train/segmented.ts

@@ -36,18 +36,31 @@
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
+import { peakMemory } from "../mlx/ffi";
 import { Vjp } from "../mlx/autograd";
+
+// MLX_BUN_SEG_MEM_LOG=1: log the within-step peak after each phase of the
+// segmented step (forward+boundaries / head vjp / each segment backward).
+// The trainer resets the peak at step start (MLX_BUN_MEM_LOG=1), so the phase
+// where the monotone peak JUMPS is where the step's memory actually lives.
+// (This is how the backlog-#2 refutation below was measured. An inter-segment
+// clearCache was also tried and had NO effect — the watermark is live memory,
+// not allocator cache retention.)
+const SEG_MEM = process.env.MLX_BUN_SEG_MEM_LOG === "1";
+const segMem = (label: string): void => {
+  if (SEG_MEM) console.log(`  [seg-mem] ${label}: peak=${(peakMemory() / 1e9).toFixed(2)} GB`);
+};
 import { MiniCPM5Model, setMiniCpmPrefixPlan } from "../model/minicpm5";
 import type { Gemma4Model } from "../model/gemma4";
 import { setGemmaPrefixPlan } from "../model/gemma4";
 import type { Cache, Mask, SharedKv } from "../model/gemma4-base";
 import { TrainingCache } from "./forward";
 import {
-  responseOnlyCe, responseOnlyLogpMean, fusedLogpMeanFromHidden, chunkedLogpMeanFromHidden,
+  responseOnlyLogpMean, fusedLogpMeanFromHidden, chunkedLogpMeanFromHidden,
   orpoLossFromLogps, spanLogpMeanFromHidden, respSpanFromMask, combineFullNll,
   type ChunkCtx, type SftScope, type HeadSpan,
 } from "./loss";
-import { PrefixSharedCache, prefixGatherIdx, branchLogpMeanGathered, blockSparsePrefixMaskGemma, prefixFullNll } from "./prefix-shared";
+import { prefixSharedCaches, prefixGatherIdx, branchLogpMeanGathered, blockSparsePrefixMaskGemma, prefixFullNll } from "./prefix-shared";
 
 // Experiment toggle: which bounded head to use inside the segmented mlx_vjp.
 // "checkpoint" = per-chunk Checkpoint recompute (bounds memory under nesting);
@@ -65,7 +78,36 @@ function boundedHeadFromHidden(
     : chunkedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink, span);
 }
 import type { TrainableLora } from "./lora-params";
-import { type SftBatch, type DpoBatch } from "./dataset";
+import { rowLength, type SftBatch, type DpoBatch } from "./dataset";
+
+// Token-chunk size for the bounded SFT head (the ORPO paths take theirs from
+// cfg.orpoChunkSize; SFT has no per-run knob, 512 matches the ORPO default).
+const SFT_HEAD_CHUNK = Number(process.env.MLX_BUN_SEG_HEAD_CHUNK ?? "512");
+
+/** [M,V]-bounded response-only CE from post-finalNorm hidden — the SFT
+ *  segmented head (kernel-review backlog #8). Same span arithmetic and
+ *  token-mean reduction as loss.ts responseOnlyCe, but the logits only ever
+ *  exist per token-chunk ([chunk,V] Checkpoint recompute, or the fused
+ *  CustomVjp head via MLX_BUN_SEG_HEAD) instead of the full [1,M,V] + its
+ *  cotangent inside the head vjp. CE = −(mean response logp); per-position
+ *  grads are bit-exact vs the whole-vocab head (chunking is over TOKEN
+ *  positions; each position's full-V lse is computed unchanged — only the
+ *  scalar loss VALUE can differ in fp association across chunks). */
+export function boundedSftCe(
+  model: MiniCPM5Model | Gemma4Model, hn: MlxArray, batch: SftBatch,
+  sink: Array<{ dispose(): void }>, chunkSize = SFT_HEAD_CHUNK,
+): MlxArray {
+  const L = batch.ids[0]!.length;
+  const startT = Math.max(0, batch.promptLens[0]! - 1); // first supervised input position
+  const M = Math.min(rowLength(batch, 0), L) - 1 - startT;
+  if (M <= 0) throw new Error("sftLoss: no response tokens to supervise");
+  const mean = boundedHeadFromHidden(model, hn, batch.ids[0]!, [], chunkSize, sink, { startT, M });
+  const neg = ops.mulScalar(mean, -1); // CE = −(mean logp); exact sign flip
+  mean.dispose();
+  const loss = ops.sumAxis(neg, 0, false); // [1] -> scalar, matches the vjp cotangent `one`
+  neg.dispose();
+  return loss;
+}
 
 /** Split `nLayers` into contiguous `[lo, hi)` ranges of at most `segSize`
  *  layers each (the last range may be shorter). `segSize` is the memory knob:
@@ -77,24 +119,50 @@ export function planSegmentsBySize(nLayers: number, segSize: number): [number, n
   return ranges;
 }
 
+// NOTE (2026-07-02, kernel-review backlog #2 REFUTED BY MEASUREMENT): a
+// "full-attention isolation" planner (singleton segments for e4b's 7 full
+// layers) was built and A/B'd @8K — ZERO peak win (18.09 vs 18.02 GB). The
+// design §5 premise is false: mlx's sdpa BACKWARD materializes the O(L²)
+// scores for EVERY layer (~3.5 GB/layer @8K on e4b — sliding's window is just
+// an additive mask; measured: sliding pair +7.1 GB ≈ sliding+full +7.15). The
+// worst segment is set by LAYER COUNT alone, so segSize IS the whole knob:
+// segment_size 1 measures 14.59 GB @8K (+3% step time) vs 17.5 @seg2. Below
+// that needs the [M,V]-bounded SFT head (backlog #8, ~3 GB head-vjp spike)
+// and/or an O(L)-memory attention backward. Evidence:
+// scripts/experiments/seg-isolation-smoke.ts + MLX_BUN_SEG_MEM_LOG=1 probes.
+
 /** Detach `a` into a graph-free LEAF holding its exact (evaluated) bytes, then
  *  dispose `a`. mlx retains the upstream graph after `eval`, so a segment
  *  boundary (or a value_and_grad output) would otherwise keep its whole layer
  *  stack's activations alive. Copying the bytes into a fresh leaf frees that
  *  graph. The leaf is a valid differentiable input for the per-segment vag. */
 function detachLeaf(a: MlxArray): MlxArray {
-  // Force a row-major copy first: `rawBytes` reads the underlying buffer
-  // linearly, so a NON-contiguous input (e.g. a transposed donor K/V view) would
-  // be byte-scrambled. `ops.contiguous` is a no-op for already-contiguous arrays
+  // Force a row-major copy first: the buffer is read linearly, so a
+  // NON-contiguous input (e.g. a transposed donor K/V view) would be
+  // byte-scrambled. `ops.contiguous` is a no-op for already-contiguous arrays
   // (hidden boundaries) and materializes row-major for strided views.
+  // detachCopy = ONE mlx→mlx copy (backlog #9: the old
+  // fromBytesCopy(rawBytes()) route was two copies via the JS heap).
   const c = ops.contiguous(a);
-  c.eval();
-  const leaf = MlxArray.fromBytesCopy(c.rawBytes(), c.shape, c.dtype);
+  const leaf = c.detachCopy();
   c.dispose();
   a.dispose();
   return leaf;
 }
 
+/** Batch detach: ONE evalAll barrier for the whole set (backlog #9 — each
+ *  per-array eval was a separate blocking GPU drain; a segment's vjp returns
+ *  dh + 2·consumed-donors + 2·targets arrays, so the per-segment sync count
+ *  collapses from ~2+2n to 1), then one mlx→mlx copy each. Disposes inputs. */
+function detachLeaves(arrs: MlxArray[]): MlxArray[] {
+  const cs = arrs.map((a) => ops.contiguous(a));
+  ops.evalAll(cs);
+  const leaves = cs.map((c) => c.detachCopy());
+  for (const c of cs) c.dispose();
+  for (const a of arrs) a.dispose();
+  return leaves;
+}
+
 /** Layer index a target's modulePath belongs to, or -1 (e.g. embedding/head). */
 function targetLayer(modulePath: string): number {
   const m = modulePath.match(/\.layers\.(\d+)\./);
@@ -120,6 +188,10 @@ export class SegmentedBackward {
   private readonly one: MlxArray; // scalar cotangent 1.0 for the (scalar) head loss
   // Per-step context the (reused) closures read.
   private curBatch: SftBatch | null = null;
+  // Bounded-head per-step transients (chunk slices + Checkpoint/CustomVjp
+  // objects): the head backward recomputes through them, so they are disposed
+  // only AFTER the head vjp's outputs are detached (which forces the eval).
+  private headSink: Array<{ dispose(): void }> = [];
   private disposed = false;
 
   constructor(
@@ -140,11 +212,17 @@ export class SegmentedBackward {
     this.caches = Array.from({ length: this.nLayers }, () => new TrainingCache());
     this.one = MlxArray.fromFloat32(new Float32Array([1]), []); // d(loss)/d(loss) = 1
 
-    // Loss head: vjp of finalNorm + responseOnlyCe over the last boundary leaf.
+    // Loss head: vjp of finalNorm + the [M,V]-BOUNDED response CE over the last
+    // boundary leaf (kernel-review backlog #8: the old responseOnlyCe here
+    // materialized [1,M,V] logits + cotangents inside the head vjp — measured
+    // ~3 GB @8K on e4b, the exact transient the segmentation exists to kill).
+    // CE = −(mean response logp); span and reduction identical to
+    // responseOnlyCe, per-position grads bit-exact (chunking is over token
+    // positions — each position's full-V lse/softmax is computed unchanged).
     // outputs[0] = loss scalar; vjps[0] = dLoss/d(last residual stream) = dh.
     this.headVjp = new Vjp((p) => {
       const hn = model.finalNorm.forward(p[0]!);
-      const loss = responseOnlyCe(model, hn, this.curBatch!);
+      const loss = boundedSftCe(model, hn, this.curBatch!, this.headSink);
       hn.dispose();
       return [loss];
     }, 1);
@@ -204,9 +282,17 @@ export class SegmentedBackward {
       }
 
       // --- 2. Loss head over the last boundary (vjp with cotangent 1.0). ---
+      // The bounded head pushes per-step chunk slices + Checkpoints into
+      // headSink; the head BACKWARD recomputes through them, so materialize
+      // the full head-vjp graph BEFORE freeing the sink (lazy-eval UAF
+      // discipline, same as the ORPO classes).
+      this.headSink = [];
       const head = this.headVjp.apply([boundaries[nSeg]!], [this.one]);
+      ops.evalAll([head.outputs[0]!, head.vjps[0]!]);
       value = detachLeaf(head.outputs[0]!); // scalar loss, detached from the head graph (frees logits)
       dhOut = detachLeaf(head.vjps[0]!); // dLoss / d(last residual stream), graph-free
+      for (const d of this.headSink) d.dispose();
+      this.headSink = [];
       boundaries[nSeg]!.dispose();
       boundaries[nSeg] = null;
 
@@ -220,13 +306,15 @@ export class SegmentedBackward {
         res.outputs[0]!.dispose(); // segment output — unused (only its vjp matters)
         // vjps[0] = dh_in (grad w.r.t. this segment's input boundary) -> dh_out
         // for segment k-1; the rest scatter into the global flat vector. Detach
-        // each to a graph-free leaf (eval doesn't detach; keeps peak bounded).
+        // to graph-free leaves (eval doesn't detach; keeps peak bounded) — ONE
+        // barrier for the segment's whole vjp set.
+        const leaves = detachLeaves(res.vjps);
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          grads[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          grads[j] = leaves[1 + 2 * s]!; // dA
+          grads[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhOut!.dispose(); // consumed as this segment's cotangent
-        dhOut = detachLeaf(res.vjps[0]!);
+        dhOut = leaves[0]!;
         boundaries[k]!.dispose();
         boundaries[k] = null;
       }
@@ -238,6 +326,8 @@ export class SegmentedBackward {
     } finally {
       inputIds.dispose();
       this.curBatch = null;
+      for (const d of this.headSink) d.dispose(); // no-op unless the head threw mid-step
+      this.headSink = [];
       dhOut?.dispose();
       for (const b of boundaries) b?.dispose();
       if (!transferred) {
@@ -295,6 +385,8 @@ export class SegmentedBackwardGemma4 {
   private curBatch: SftBatch | null = null;
   private curMasks: Map<string, Mask> | null = null;
   private curPerLayer: MlxArray | null = null;
+  // Bounded-head per-step transients — see SegmentedBackward.headSink.
+  private headSink: Array<{ dispose(): void }> = [];
   private disposed = false;
 
   constructor(
@@ -337,11 +429,11 @@ export class SegmentedBackwardGemma4 {
     this.caches = Array.from({ length: model.numDonors }, () => new TrainingCache());
     this.one = MlxArray.fromFloat32(new Float32Array([1]), []);
 
-    // Loss head (identical to MiniCPM5): finalNorm + responseOnlyCe over the last
+    // Loss head (identical to MiniCPM5): finalNorm + bounded response CE over the last
     // boundary, vjp with cotangent 1.0 -> [loss], [dh].
     this.headVjp = new Vjp((p) => {
       const hn = model.finalNorm.forward(p[0]!);
-      const loss = responseOnlyCe(model, hn, this.curBatch!);
+      const loss = boundedSftCe(model, hn, this.curBatch!, this.headSink);
       hn.dispose();
       return [loss];
     }, 1);
@@ -437,14 +529,22 @@ export class SegmentedBackwardGemma4 {
         for (const [d, kv] of donorKvOut)
           if (!produces[k]!.includes(d) && kv.kind === "plain") { kv.keys.dispose(); kv.values.dispose(); }
         boundaries.push(detachLeaf(h));
+        segMem(`fwd seg ${k} [${lo},${hi})`);
       }
 
-      // --- 2. Loss head (vjp, cotangent 1.0). ---
+      // --- 2. Loss head (vjp, cotangent 1.0). Bounded-head sink discipline:
+      // materialize the head-vjp graph BEFORE freeing the sink (see the
+      // MiniCPM5 class / ORPO classes — lazy-eval UAF). ---
+      this.headSink = [];
       const head = this.headVjp.apply([boundaries[nSeg]!], [this.one]);
+      ops.evalAll([head.outputs[0]!, head.vjps[0]!]);
       value = detachLeaf(head.outputs[0]!);
       dhOut = detachLeaf(head.vjps[0]!);
+      for (const d of this.headSink) d.dispose();
+      this.headSink = [];
       boundaries[nSeg]!.dispose();
       boundaries[nSeg] = null;
+      segMem("head vjp");
 
       // --- 3. Backward, reverse over segments (vjp, cotangents dh + dKV). ---
       for (let k = nSeg - 1; k >= 0; k--) {
@@ -466,17 +566,17 @@ export class SegmentedBackwardGemma4 {
         const res = this.segVjps[k]!.apply(primals, cots);
         for (const o of res.outputs) o.dispose();
 
-        // vjps[0] = dh_in; then (dk,dv) per consumed donor; then (da,db) per LoRA.
-        const newDh = detachLeaf(res.vjps[0]!);
+        // vjps[0] = dh_in; then (dk,dv) per consumed donor; then (da,db) per
+        // LoRA — detached in ONE barrier (backlog #9).
+        const leaves = detachLeaves(res.vjps);
+        const newDh = leaves[0]!;
         cons.forEach((d, ci) => {
-          const dk = detachLeaf(res.vjps[1 + 2 * ci]!);
-          const dv = detachLeaf(res.vjps[1 + 2 * ci + 1]!);
-          this.accumulateDKV(dKV, d, dk, dv);
+          this.accumulateDKV(dKV, d, leaves[1 + 2 * ci]!, leaves[1 + 2 * ci + 1]!);
         });
         const base = 1 + 2 * cons.length;
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[base + 2 * s]!);
-          grads[n + j] = detachLeaf(res.vjps[base + 2 * s + 1]!);
+          grads[j] = leaves[base + 2 * s]!;
+          grads[n + j] = leaves[base + 2 * s + 1]!;
         });
 
         dh.dispose();
@@ -488,6 +588,7 @@ export class SegmentedBackwardGemma4 {
         }
         boundaries[k]!.dispose();
         boundaries[k] = null;
+        segMem(`bwd seg ${k} [${ranges[k]![0]},${ranges[k]![1]})`);
       }
 
       transferred = true;
@@ -498,6 +599,8 @@ export class SegmentedBackwardGemma4 {
       this.curPerLayer?.dispose();
       this.curPerLayer = null;
       if (this.curMasks) { for (const m of this.curMasks.values()) m.arr?.dispose(); this.curMasks = null; }
+      for (const d of this.headSink) d.dispose(); // no-op unless the head threw mid-step
+      this.headSink = [];
       dhOut?.dispose();
       for (const b of boundaries) b?.dispose();
       for (const kv of donorBnd.values()) { kv.keys.dispose(); kv.values.dispose(); }
@@ -725,12 +828,13 @@ export class SegmentedBackwardOrpo {
         for (const j of idxs) params.push(lora.targets[j]!.lw.a, lora.targets[j]!.lw.b);
         const res = this.segVjps[k]!.apply(params, [dhC!]);
         res.outputs[0]!.dispose();
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
         idxs.forEach((j, s) => {
-          gradsC[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          gradsC[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          gradsC[j] = leaves[1 + 2 * s]!; // dA
+          gradsC[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhC!.dispose();
-        dhC = detachLeaf(res.vjps[0]!);
+        dhC = leaves[0]!;
         bndC[k]!.dispose(); bndC[k] = null;
       }
       dhC!.dispose(); dhC = null; // embedding grad, discarded (no LoRA on embed)
@@ -742,12 +846,13 @@ export class SegmentedBackwardOrpo {
         for (const j of idxs) params.push(lora.targets[j]!.lw.a, lora.targets[j]!.lw.b);
         const res = this.segVjps[k]!.apply(params, [dhR!]);
         res.outputs[0]!.dispose();
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
         idxs.forEach((j, s) => {
-          gradsR[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          gradsR[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          gradsR[j] = leaves[1 + 2 * s]!; // dA
+          gradsR[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhR!.dispose();
-        dhR = detachLeaf(res.vjps[0]!);
+        dhR = leaves[0]!;
         bndR[k]!.dispose(); bndR[k] = null;
       }
       dhR!.dispose(); dhR = null; // embedding grad, discarded
@@ -912,7 +1017,7 @@ export class SegmentedBackwardOrpoPrefix {
     this.headSink = [];
     this.curChunk = this.chunkCtx ? { ...this.chunkCtx, sink: this.headSink } : undefined;
     // One prefix cache per layer (captures P/Rc/Rr; makeMask -> block-sparse mask).
-    this.curCaches = model.layers.map(() => new PrefixSharedCache(P, Rc, Rr));
+    this.curCaches = prefixSharedCaches(model.layers.length, P, Rc, Rr);
 
     const nSeg = ranges.length;
     const concat = new Int32Array(T);
@@ -957,12 +1062,13 @@ export class SegmentedBackwardOrpoPrefix {
         for (const j of idxs) params.push(lora.targets[j]!.lw.a, lora.targets[j]!.lw.b);
         const res = this.segVjps[k]!.apply(params, [dhOut!]);
         res.outputs[0]!.dispose();
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          grads[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          grads[j] = leaves[1 + 2 * s]!; // dA
+          grads[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhOut!.dispose();
-        dhOut = detachLeaf(res.vjps[0]!);
+        dhOut = leaves[0]!;
         boundaries[k]!.dispose(); boundaries[k] = null;
       }
       // Segment-0 dh_in = grad w.r.t. the embedding output (no LoRA on embed) — drop.
@@ -1232,16 +1338,15 @@ export class SegmentedBackwardOrpoGemma4 {
         const res = this.segVjps[k]!.apply(primals, cots);
         for (const o of res.outputs) o.dispose();
 
-        const newDh = detachLeaf(res.vjps[0]!);
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
+        const newDh = leaves[0]!;
         cons.forEach((d, ci) => {
-          const dk = detachLeaf(res.vjps[1 + 2 * ci]!);
-          const dv = detachLeaf(res.vjps[1 + 2 * ci + 1]!);
-          this.accumulateDKV(dKV, d, dk, dv);
+          this.accumulateDKV(dKV, d, leaves[1 + 2 * ci]!, leaves[1 + 2 * ci + 1]!);
         });
         const base = 1 + 2 * cons.length;
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[base + 2 * s]!);
-          grads[n + j] = detachLeaf(res.vjps[base + 2 * s + 1]!);
+          grads[j] = leaves[base + 2 * s]!;
+          grads[n + j] = leaves[base + 2 * s + 1]!;
         });
 
         dh.dispose();
@@ -1663,16 +1768,15 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
         const res = this.segVjps[k]!.apply(primals, cots);
         for (const o of res.outputs) o.dispose();
 
-        const newDh = detachLeaf(res.vjps[0]!);
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
+        const newDh = leaves[0]!;
         cons.forEach((d, ci) => {
-          const dk = detachLeaf(res.vjps[1 + 2 * ci]!);
-          const dv = detachLeaf(res.vjps[1 + 2 * ci + 1]!);
-          this.accumulateDKV(dKV, d, dk, dv);
+          this.accumulateDKV(dKV, d, leaves[1 + 2 * ci]!, leaves[1 + 2 * ci + 1]!);
         });
         const base = 1 + 2 * cons.length;
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[base + 2 * s]!);
-          grads[n + j] = detachLeaf(res.vjps[base + 2 * s + 1]!);
+          grads[j] = leaves[base + 2 * s]!;
+          grads[n + j] = leaves[base + 2 * s + 1]!;
         });
 
         dh.dispose();

package/src/train/trainer.ts

@@ -388,6 +388,16 @@ export async function trainLora(
   }
 }
 
+// ---------------------------------------------------------------------------
+// Segment planning: uniform planSegmentsBySize only. A full-attention
+// ISOLATION planner (kernel-review backlog #2 / segmented-backward-training.md
+// §5) was built and A/B-measured 2026-07-02: ZERO peak win — mlx's sdpa
+// backward is O(L²) for EVERY layer (~3.5 GB/layer @8K e4b), sliding included,
+// so the worst segment is set by layer count alone and segment_size is the
+// whole knob (seg1 = 14.59 GB @8K, +3% step time — fits the 24 GB M4 Pro).
+// Evidence: scripts/experiments/seg-isolation-smoke.ts; note in segmented.ts.
+// ---------------------------------------------------------------------------
+
 // ---------------------------------------------------------------------------
 // SFT loop
 // ---------------------------------------------------------------------------
@@ -421,14 +431,14 @@ async function sftLoop(
   const useSegmented = cfg.segmentSize > 0;
   let segmented: SegmentedBackward | SegmentedBackwardGemma4 | null = null;
   if (useSegmented) {
-    if (model instanceof MiniCPM5Model)
-      segmented = new SegmentedBackward(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
-    else if (model instanceof Gemma4Model)
-      segmented = new SegmentedBackwardGemma4(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
-    else
+    if (!(model instanceof MiniCPM5Model) && !(model instanceof Gemma4Model))
       throw new Error("segmented backward (segmentSize > 0) is only wired for MiniCPM5 and Gemma4");
+    const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
+    segmented = model instanceof MiniCPM5Model
+      ? new SegmentedBackward(model, lora, ranges)
+      : new SegmentedBackwardGemma4(model, lora, ranges);
     emit({ type: "stage", stage: "setup", progress: 0.04,
-      message: `segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers` });
+      message: `segmented backward: ${ranges.length} segments of <=${cfg.segmentSize} layers` });
   }
 
   // The value_and_grad loss closure: temporarily swap the differentiated
@@ -832,6 +842,8 @@ async function orpoLoop(
   let segmentedOrpo: SegmentedBackwardOrpo | SegmentedBackwardOrpoGemma4 | null = null;
   let segmentedOrpoPrefix: SegmentedBackwardOrpoPrefix | SegmentedBackwardOrpoPrefixGemma4 | null = null;
   if (useSegmented) {
+    if (!(model instanceof MiniCPM5Model) && !(model instanceof Gemma4Model))
+      throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
     const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
     if (useSegmentedPrefix) {
       // The composed prefix-shared segmented backward (MiniCPM5 or Gemma4/e4b), plus a
@@ -845,15 +857,13 @@ async function orpoLoop(
         segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
       }
       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)` });
+        message: `orpo segmented + prefix-share (${model instanceof Gemma4Model ? "e4b" : "MiniCPM5"}): single concat forward streamed over ${ranges.length} 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, cfg.sftScope);
-    else if (model instanceof Gemma4Model)
-      segmentedOrpo = new SegmentedBackwardOrpoGemma4(model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
     else
-      throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
+      segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
     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)` : ""}` });
+      message: `orpo segmented backward: ${ranges.length} 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