mlx-bun 0.0.7 → 0.0.9

package/src/train/segmented.ts

@@ -42,8 +42,12 @@ 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, orpoLossFromLogps, type ChunkCtx } from "./loss";
-import { PrefixSharedCache, prefixGatherIdx, branchLogpMeanGathered, blockSparsePrefixMaskGemma } from "./prefix-shared";
+import {
+  responseOnlyCe, responseOnlyLogpMean, fusedLogpMeanFromHidden, chunkedLogpMeanFromHidden,
+  orpoLossFromLogps, spanLogpMeanFromHidden, respSpanFromMask, combineFullNll,
+  type ChunkCtx, type SftScope, type HeadSpan,
+} from "./loss";
+import { PrefixSharedCache, 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);
@@ -54,11 +58,11 @@ import { PrefixSharedCache, prefixGatherIdx, branchLogpMeanGathered, blockSparse
 const SEG_HEAD = process.env.MLX_BUN_SEG_HEAD ?? "checkpoint";
 function boundedHeadFromHidden(
   model: MiniCPM5Model | Gemma4Model, h: MlxArray, ids: number[], mask: number[],
-  chunkSize: number, sink: Array<{ dispose(): void }>,
+  chunkSize: number, sink: Array<{ dispose(): void }>, span?: HeadSpan,
 ): MlxArray {
   return SEG_HEAD === "fused"
-    ? fusedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink)
-    : chunkedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink);
+    ? fusedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink, 0, false, span)
+    : chunkedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink, span);
 }
 import type { TrainableLora } from "./lora-params";
 import { type SftBatch, type DpoBatch } from "./dataset";
@@ -582,6 +586,7 @@ export class SegmentedBackwardOrpo {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly fusedChunkSize = 0,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -609,9 +614,26 @@ export class SegmentedBackwardOrpo {
       const lr = cs > 0 // [1]
         ? boundedHeadFromHidden(model, hnR, this.curRejectedIds!, this.curRejectedMask!, cs, this.headSink)
         : responseOnlyLogpMean(model, hnR, this.curRejectedIds!, this.curRejectedMask!);
+      // sft_scope:"full": the chosen head additionally runs over the PROMPT span
+      // [0, startT) of the SAME hnC (same tier: bounded or whole-vocab), and the
+      // full-scope chosen NLL replaces mean(-ℓw). The prompt-position gradient
+      // rides into this vjp's dh_chosen automatically (the loss graph includes
+      // the prompt-span head). ℓw/ℓr stay response-only for the odds ratio.
+      let nllFull: MlxArray | null = null;
+      if (this.sftScope === "full") {
+        const { startT, M } = respSpanFromMask(this.curChosenMask!, this.curChosenIds!.length - 1);
+        let pm: MlxArray | null = null;
+        if (M > 0 && startT > 0)
+          pm = cs > 0
+            ? boundedHeadFromHidden(model, hnC, this.curChosenIds!, this.curChosenMask!, cs, this.headSink, { startT: 0, M: startT })
+            : spanLogpMeanFromHidden(model, hnC, this.curChosenIds!, 0, startT);
+        nllFull = combineFullNll(pm, lw, M > 0 ? startT : 0, M);
+        pm?.dispose();
+      }
       hnC.dispose(); hnR.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -801,6 +823,7 @@ export class SegmentedBackwardOrpoPrefix {
   private curP = 0;
   private curRc = 0;
   private curRr = 0;
+  private curPromptIds: number[] | null = null;
   private curChosenResp: number[] | null = null;
   private curRejectedResp: number[] | null = null;
   private curChunk: ChunkCtx | undefined = undefined;
@@ -816,6 +839,7 @@ export class SegmentedBackwardOrpoPrefix {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly chunkCtx?: ChunkCtx,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -838,9 +862,16 @@ export class SegmentedBackwardOrpoPrefix {
       const { chosenIdx, rejectedIdx } = prefixGatherIdx(this.curP, this.curRc, this.curRr);
       const lw = branchLogpMeanGathered(model, hn, chosenIdx, this.curChosenResp!, this.curChunk);
       const lr = branchLogpMeanGathered(model, hn, rejectedIdx, this.curRejectedResp!, this.curChunk);
+      // sft_scope:"full": prompt predictions (H[0..P-2] → prompt[1..P-1]) from the
+      // same concat hiddens; the gather's takeAxis vjp scatter-adds the prompt
+      // cotangent into dh alongside the responses'. ℓw/ℓr stay response-only.
+      const nllFull = this.sftScope === "full"
+        ? prefixFullNll(model, hn, this.curPromptIds!, lw, this.curRc, this.curChunk)
+        : null;
       hn.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -873,6 +904,7 @@ export class SegmentedBackwardOrpoPrefix {
     if (P < 1 || Rc < 1 || Rr < 1) throw new Error("SegmentedBackwardOrpoPrefix: need P,Rc,Rr >= 1");
     const T = P + Rc + Rr;
     this.curP = P; this.curRc = Rc; this.curRr = Rr;
+    this.curPromptIds = promptIds;
     this.curChosenResp = chosenResp; this.curRejectedResp = rejectedResp;
     // Per-step chunk ctx pointing the head's CustomVjp/slice disposables at this
     // step's headSink (the flash/fused head recompute runs during headVjp.apply;
@@ -942,7 +974,7 @@ export class SegmentedBackwardOrpoPrefix {
       inputIds.dispose();
       for (const c of this.curCaches) c.dispose();
       this.curCaches = [];
-      this.curChosenResp = this.curRejectedResp = null;
+      this.curPromptIds = this.curChosenResp = this.curRejectedResp = null;
       for (const d of this.headSink) d.dispose();
       this.headSink = [];
       dhOut?.dispose();
@@ -1024,6 +1056,7 @@ export class SegmentedBackwardOrpoGemma4 {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly fusedChunkSize = 0,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -1072,9 +1105,23 @@ export class SegmentedBackwardOrpoGemma4 {
       const lr = cs > 0
         ? boundedHeadFromHidden(this.model, hnR, this.curRejectedIds!, this.curRejectedMask!, cs, this.headSink)
         : responseOnlyLogpMean(this.model, hnR, this.curRejectedIds!, this.curRejectedMask!);
+      // sft_scope:"full": prompt-span head over the same hnC (see
+      // SegmentedBackwardOrpo); ℓw/ℓr stay response-only for the odds ratio.
+      let nllFull: MlxArray | null = null;
+      if (this.sftScope === "full") {
+        const { startT, M } = respSpanFromMask(this.curChosenMask!, this.curChosenIds!.length - 1);
+        let pm: MlxArray | null = null;
+        if (M > 0 && startT > 0)
+          pm = cs > 0
+            ? boundedHeadFromHidden(this.model, hnC, this.curChosenIds!, this.curChosenMask!, cs, this.headSink, { startT: 0, M: startT })
+            : spanLogpMeanFromHidden(this.model, hnC, this.curChosenIds!, 0, startT);
+        nllFull = combineFullNll(pm, lw, M > 0 ? startT : 0, M);
+        pm?.dispose();
+      }
       hnC.dispose(); hnR.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -1394,6 +1441,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
   private curP = 0;
   private curRc = 0;
   private curRr = 0;
+  private curPromptIds: number[] | null = null;
   private curChosenResp: number[] | null = null;
   private curRejectedResp: number[] | null = null;
   private curChunk: ChunkCtx | undefined = undefined;
@@ -1408,6 +1456,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly chunkCtx?: ChunkCtx,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -1455,9 +1504,15 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
       const { chosenIdx, rejectedIdx } = prefixGatherIdx(this.curP, this.curRc, this.curRr);
       const lw = branchLogpMeanGathered(model, hn, chosenIdx, this.curChosenResp!, this.curChunk);
       const lr = branchLogpMeanGathered(model, hn, rejectedIdx, this.curRejectedResp!, this.curChunk);
+      // sft_scope:"full": prompt predictions from the same concat hiddens (see
+      // SegmentedBackwardOrpoPrefix); ℓw/ℓr stay response-only.
+      const nllFull = this.sftScope === "full"
+        ? prefixFullNll(model, hn, this.curPromptIds!, lw, this.curRc, this.curChunk)
+        : null;
       hn.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -1525,6 +1580,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
     if (P < 1 || Rc < 1 || Rr < 1) throw new Error("SegmentedBackwardOrpoPrefixGemma4: need P,Rc,Rr >= 1");
     const T = P + Rc + Rr;
     this.curP = P; this.curRc = Rc; this.curRr = Rr;
+    this.curPromptIds = promptIds;
     this.curChosenResp = chosenResp; this.curRejectedResp = rejectedResp;
     this.headSink = [];
     this.curChunk = this.chunkCtx ? { ...this.chunkCtx, sink: this.headSink } : undefined;
@@ -1640,7 +1696,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
       this.curPerLayer?.dispose();
       this.curPerLayer = null;
       if (this.curMasks) { for (const m of this.curMasks.values()) m.arr?.dispose(); this.curMasks = null; }
-      this.curChosenResp = this.curRejectedResp = null;
+      this.curPromptIds = this.curChosenResp = this.curRejectedResp = null;
       for (const d of this.headSink) d.dispose();
       this.headSink = [];
       dhOut?.dispose();

package/src/train/trainer.ts

@@ -34,7 +34,7 @@ import {
   saveAdapter, disposeLora, warmStartFromAdapter, type TrainableLora, type SaveAdapterConfig,
 } from "./lora-params";
 import { resolveRanks, bitsMapFromModel, readPerLayerKl, DEFAULT_TARGET_MODULES, type RankScaling } from "./rank";
-import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics, orpoLoss, orpoMetrics, type ChunkCtx } from "./loss";
+import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics, orpoLoss, orpoMetrics, type ChunkCtx, type SftScope } from "./loss";
 import { orpoLossPrefixShared, orpoLossPrefixSharedGemma, splitPrefixBatch } from "./prefix-shared";
 import { AdamW, warmupCosineSchedule } from "./optimizer";
 import { writeFileSync, appendFileSync, mkdirSync } from "node:fs";
@@ -134,6 +134,15 @@ export interface TrainConfig {
    *  differ. Composes with orpoFlashCe/orpoFusedCe (each branch routes through the
    *  [M,V]-free head). See src/train/prefix-shared.ts. */
   orpoPrefixShared: boolean;
+  /** ORPO L_SFT scope (`sft_scope`). "full" (default; paper/TRL-faithful):
+   *  the chosen-NLL term is the token-mean cross-entropy over the FULL
+   *  prompt+response (only padding excluded), computed from the same chosen
+   *  forward — matching TRL's ORPOTrainer chosen_nll_loss. "response":
+   *  L_NLL = -ℓw (response-only), bit-exact to pre-sft_scope runs. The
+   *  odds-ratio ℓw/ℓr terms are response-only length-normalized means in BOTH
+   *  modes (that also matches TRL). Applies to every ORPO path (naive /
+   *  chunked / fused / flash / prefix-shared / segmented). orpo only. */
+  sftScope: SftScope;
   /** 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.
@@ -188,6 +197,7 @@ export const DEFAULT_TRAIN_CONFIG: TrainConfig = {
   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)
+  sftScope: "full", // paper/TRL-faithful chosen-NLL over prompt+response; "response" = pre-2026-07 behavior
   warmStartAdapter: "", // default off; path to an adapter dir to warm-start LoRA weights from
   adapterPath: "adapters",
   baseModel: "",
@@ -298,6 +308,7 @@ export async function trainLora(
     iters: cfg.iters, learning_rate: cfg.learningRate, max_seq_length: cfg.maxSeqLen,
     batch_size: cfg.batchSize, rank: cfg.rank, scale: cfg.scale,
     orpo_lambda: cfg.method === "orpo" ? cfg.orpoLambda : undefined,
+    sft_scope: cfg.method === "orpo" ? cfg.sftScope : undefined,
     adapter_path: cfg.adapterPath,
   });
 
@@ -349,6 +360,7 @@ export async function trainLora(
           segmentSize: cfg.segmentSize, iters: cfg.iters,
           numLayers: cfg.numLayers, baseModel: cfg.baseModel,
           orpoLambda: cfg.method === "orpo" ? cfg.orpoLambda : undefined,
+          sftScope: cfg.method === "orpo" ? cfg.sftScope : undefined,
         },
         wallSeconds: (Date.now() - startedAt) / 1000,
         peakGb: peakMemory() / 1e9,
@@ -466,14 +478,40 @@ async function sftLoop(
   // 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");
+  //
+  // GEMMA GUARD (same convention as the PERF_KERNEL/FUSED_GELU sanitization in
+  // cli.ts / job.ts, but enforced at the trainer): e4b on this path SIGTRAPed
+  // (uncatchable native crash) at multi-K sequence lengths (>=2K, reproduced in
+  // scripts/experiments/segmented-grad-test-e4b.ts; docs/reference/training.md)
+  // and has NOT been re-validated at that scale since the two kernel fixes —
+  // the regression tests stop at T<=256. cli.ts even defaults e4b seq to 8192,
+  // so a stale `export MLX_BUN_TRAIN_ATTN=flash` from a MiniCPM5 experiment
+  // would otherwise ride silently into a crash mid-run. Refuse it for Gemma
+  // until the >=2K re-validation lands; MiniCPM5 stays allowed.
+  if (process.env.MLX_BUN_TRAIN_ATTN === "flash") {
+    if (model instanceof Gemma4Model)
+      throw new Error(
+        "MLX_BUN_TRAIN_ATTN=flash is disabled for Gemma models: e4b SIGTRAPs on this " +
+          "path at seq >= 2048 (docs/reference/training.md) and it has not been " +
+          "re-validated at that scale since the kernel fixes. Unset MLX_BUN_TRAIN_ATTN " +
+          "(ops.sdpa, the default, is exact and ~30x faster) or train MiniCPM5.",
+      );
+    emit({ type: "stage", stage: "setup", progress: 0.045,
+      message: "training attention: hand-rolled flash kernel (MLX_BUN_TRAIN_ATTN=flash; O(L) memory, ~30x slower than ops.sdpa)" });
+    setTrainingAttn("flash");
+  }
 
   // Pad with the eos id when available (a real sentinel token); padded
   // positions are excluded from both the loss and attention masks, so the
   // exact pad value never affects the result — it just must be a valid id.
   const padId = tok.eosTokenId ?? 0;
   const batches = iterateSftBatches(train, cfg.batchSize, cfg.maxSeqLen, cfg.seed, true, padId);
-  const t0 = Date.now();
+  // tokens/sec follows mlx-lm's tuner convention: accumulate the response-token
+  // count across the REPORT WINDOW and divide by the window's own wall time
+  // (validation time excluded) — not one step's count extrapolated over the
+  // whole run, which jumps ~10x step-to-step on variable-length data.
+  let windowTokens = 0;
+  let windowStart = Date.now();
 
   try {
     for (let step = 1; step <= cfg.iters; step++) {
@@ -483,12 +521,11 @@ async function sftLoop(
       // 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);
+          windowTokens += countResponseTokens(currentBatch);
           return segmented ? segmented.step(currentBatch!) : vag!.apply(flatParams(lora));
         },
         ckptCtx
@@ -512,7 +549,7 @@ async function sftLoop(
       }
 
       if (step % cfg.stepsPerReport === 0 || step === 1) {
-        const elapsed = (Date.now() - t0) / 1000;
+        const elapsed = (Date.now() - windowStart) / 1000;
         emit({
           type: "metric",
           kind: "train",
@@ -520,19 +557,31 @@ async function sftLoop(
           loss: lossVal,
           grad_norm: gradNorm,
           learning_rate: opt.lr,
-          tokens_per_sec: elapsed > 0 ? (step * tokensThisStep) / elapsed : 0,
+          tokens_per_sec: elapsed > 0 ? windowTokens / elapsed : 0,
           progress: step / cfg.iters,
           message: `step ${step}/${cfg.iters} loss=${lossVal.toFixed(4)}`,
         });
+        windowTokens = 0;
+        windowStart = Date.now();
       }
 
       if (valid.length > 0 && step % cfg.stepsPerEval === 0) {
         // Eval is forward-only (no backward), so disable checkpointing for it —
         // it adds no benefit and a forward-only pass fits without it.
         if (ckptCtx) (model as Gemma4Model).gradCkpt = null;
-        const vLoss = evalSftLoss(model, valid, cfg);
+        const tVal = Date.now();
+        const { loss: vLoss, used, skipped } = evalSftLoss(model, valid, cfg);
         if (ckptCtx) (model as Gemma4Model).gradCkpt = ckptCtx;
-        emit({ type: "metric", kind: "val", step, loss: vLoss, progress: step / cfg.iters });
+        // A silently-shrinking val set must be visible: rows whose supervised
+        // span is empty after truncation are skipped BY DESIGN (mlx-lm drops
+        // them at dataset build time) — report how many, and over what n the
+        // mean was taken. Any other per-row failure now rethrows (see
+        // evalSftLoss).
+        emit({ type: "metric", kind: "val", step, loss: vLoss, progress: step / cfg.iters,
+          val_rows_used: used, val_rows_skipped: skipped,
+          ...(skipped > 0 ? { message: `val loss over ${used} rows (${skipped} skipped: no response tokens to supervise)` } : {}) });
+        // validation wall time is not training throughput — keep it out of the window
+        windowStart += Date.now() - tVal;
       }
     }
   } finally {
@@ -548,13 +597,16 @@ async function sftLoop(
   return { numIters: cfg.iters };
 }
 
-function evalSftLoss(model: RuntimeModel, valid: ReturnType<typeof Array.prototype.slice>, cfg: TrainConfig): number {
+function evalSftLoss(
+  model: RuntimeModel, valid: ReturnType<typeof Array.prototype.slice>, cfg: TrainConfig,
+): { loss: number; used: number; skipped: number } {
   let total = 0;
   let n = 0;
+  let skipped = 0;
   for (const ex of valid as { ids: number[]; promptLen: number }[]) {
     let ids = ex.ids;
     let promptLen = ex.promptLen;
-    if (ids.length < 2) continue;
+    if (ids.length < 2) { skipped++; continue; }
     if (ids.length > cfg.maxSeqLen) {
       ids = ids.slice(0, cfg.maxSeqLen);
       promptLen = Math.min(promptLen, Math.max(0, ids.length - 1));
@@ -565,12 +617,18 @@ function evalSftLoss(model: RuntimeModel, valid: ReturnType<typeof Array.prototy
       total += loss.toFloat32()[0]!;
       loss.dispose();
       n++;
-    } catch {
-      // skip examples with no response tokens
+    } catch (e) {
+      // ONLY the empty-supervision row is skippable (a truncated row whose
+      // response fell entirely past maxSeqLen — mlx-lm drops these at dataset
+      // build). Everything else (MLX OOM/shape/vjp errors surface as catchable
+      // JS Errors via the ffi handler) must rethrow: swallowing them silently
+      // reduces the val set to its easy rows and reports a plausible loss.
+      if (e instanceof Error && e.message === "sftLoss: no response tokens to supervise") skipped++;
+      else throw e;
     }
     clearCache();
   }
-  return n > 0 ? total / n : 0;
+  return { loss: n > 0 ? total / n : 0, used: n, skipped };
 }
 
 // ---------------------------------------------------------------------------
@@ -780,18 +838,18 @@ async function orpoLoop(
       // 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);
+        segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefix(model, lora, ranges, cfg.orpoLambda, segPrefixChunk, cfg.sftScope);
+        segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
       } else {
-        segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefixGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segPrefixChunk);
-        segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk);
+        segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefixGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segPrefixChunk, cfg.sftScope);
+        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)` });
     } else if (model instanceof MiniCPM5Model)
-      segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
+      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);
+      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");
     if (!useSegmentedPrefix) emit({ type: "stage", stage: "setup", progress: 0.04,
@@ -837,12 +895,12 @@ async function orpoLoop(
             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);
+                return orpoLossPrefixShared(model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk, cfg.sftScope);
+              return orpoLossPrefixSharedGemma(model as Gemma4Model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk, cfg.sftScope);
             }
             notePrefixFallback();
           }
-          return orpoLoss(model, currentBatch!, cfg.orpoLambda, chunk);
+          return orpoLoss(model, currentBatch!, cfg.orpoLambda, chunk, cfg.sftScope);
         } finally {
           restorePrimals(lora, saved);
         }
@@ -897,7 +955,7 @@ async function orpoLoop(
         // 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 m = orpoMetrics(model, currentBatch!, cfg.orpoLambda, cfg.sftScope);
         const elapsed = (Date.now() - t0) / 1000;
         emit({
           type: "metric",
@@ -929,7 +987,7 @@ async function orpoLoop(
             chosenIds: [ex.chosenIds], rejectedIds: [ex.rejectedIds],
             chosenMask: [ex.chosenMask], rejectedMask: [ex.rejectedMask],
           };
-          const m = orpoMetrics(model, vb, cfg.orpoLambda);
+          const m = orpoMetrics(model, vb, cfg.orpoLambda, cfg.sftScope);
           vl += m.loss; va += m.accuracy; vm += m.margin; vn++;
           clearCache();
         }
@@ -1105,10 +1163,16 @@ function clipGradsByNorm(grads: MlxArray[], normVal: number, maxNorm: number): v
   }
 }
 
+/** Supervised (response) tokens across ALL rows of the batch — the count
+ *  tokens/sec is reported over (mlx-lm's n_tokens counts every row too). */
 function countResponseTokens(batch: SftBatch): number {
-  const promptLen = batch.promptLens[0]!;
-  const len = batch.ids[0]!.length;
-  return Math.max(0, len - Math.max(promptLen, 1));
+  let total = 0;
+  for (let r = 0; r < batch.ids.length; r++) {
+    const promptLen = batch.promptLens[r]!;
+    const len = batch.ids[r]!.length;
+    total += Math.max(0, len - Math.max(promptLen, 1));
+  }
+  return total;
 }
 
 async function fileExists(path: string): Promise<boolean> {