mlx-bun 0.0.8 → 0.0.10

package/src/train/prefix-shared.ts

@@ -28,7 +28,7 @@ import * as ops from "../mlx/ops";
 import { createCausalMask, type Cache, type Mask } from "../model/gemma4-base";
 import { MiniCPM5Model, setMiniCpmPrefixPlan } from "../model/minicpm5";
 import { Gemma4Model, setGemmaPrefixPlan } from "../model/gemma4";
-import { orpoLossFromLogps, fusedRespLogpMean, type ChunkCtx } from "./loss";
+import { orpoLossFromLogps, fusedRespLogpMean, combineFullNll, logitsFromHiddenPadM, type ChunkCtx, type SftScope } from "./loss";
 import type { DpoBatch } from "./dataset";
 
 /** A model that can project final-norm hidden states to vocab logits — the only
@@ -63,17 +63,40 @@ export function blockSparseMask(P: number, Rc: number, Rr: number): MlxArray {
  *  MiniCPM5's runLayerRange calls makeMask once and disposes the returned arr. */
 export class PrefixSharedCache implements Cache {
   offset = 0;
-  constructor(private readonly P: number, private readonly Rc: number, private readonly Rr: number) {}
+  constructor(
+    private readonly P: number, private readonly Rc: number, private readonly Rr: number,
+    // Step-shared mask memo (kernel-review backlog #9): the [T,T] block-sparse
+    // mask is IDENTICAL for every layer, segment forward, and backward
+    // recompute of a step, but each runLayerRange call rebuilt it from its
+    // elementwise ops (2·nSeg+ rebuilds per step in the segmented prefix
+    // class). With a memo it is built once; makeMask hands out cheap slice
+    // VIEWS (consumers dispose their view — runLayerRange does), and the
+    // first cache disposed frees the memo. Use prefixSharedCaches() to build
+    // a memo-sharing batch.
+    private readonly memo?: { mask: MlxArray | null },
+  ) {}
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
     return [k.slice([0, 0, 0, 0], k.shape), v.slice([0, 0, 0, 0], v.shape)];
   }
   makeMask(_N: number, _windowSize: number | null): Mask {
-    return { mode: "array", arr: blockSparseMask(this.P, this.Rc, this.Rr) };
+    if (!this.memo) return { mode: "array", arr: blockSparseMask(this.P, this.Rc, this.Rr) };
+    if (!this.memo.mask) this.memo.mask = blockSparseMask(this.P, this.Rc, this.Rr);
+    const m = this.memo.mask;
+    return { mode: "array", arr: m.slice([0, 0], m.shape) };
   }
   state(): MlxArray[] { return []; }
   isTrimmable(): boolean { return true; }
   trim(_n: number): void { /* offset pinned at 0 */ }
-  dispose(): void { /* owns no arrays */ }
+  dispose(): void {
+    if (this.memo?.mask) { this.memo.mask.dispose(); this.memo.mask = null; }
+  }
+}
+
+/** One PrefixSharedCache per layer, all sharing a single block-sparse-mask
+ *  memo for the step (see the constructor note). */
+export function prefixSharedCaches(nLayers: number, P: number, Rc: number, Rr: number): Cache[] {
+  const memo: { mask: MlxArray | null } = { mask: null };
+  return Array.from({ length: nLayers }, () => new PrefixSharedCache(P, Rc, Rr, memo));
 }
 
 /** Length-normalized mean log-prob over the predictions at hidden positions
@@ -87,7 +110,7 @@ function gatheredLogpMean(model: LogitProjector, h: MlxArray, gatherIdx: number[
   const hSel = ops.takeAxis(h, idxArr, 1); // [1, M, hidden]
   idxArr.dispose();
   if (hSel.shape[2] !== hidden) throw new Error("gatheredLogpMean: hidden mismatch");
-  const logits = model.logitsFromHidden(hSel); // [1, M, V]
+  const logits = logitsFromHiddenPadM(model, hSel); // [1, M, V] (M in {2,3} padded — mlx qmm bug)
   hSel.dispose();
   const V = logits.shape[2]!;
   const logits2d = ops.reshape(logits, [M, V]);
@@ -147,6 +170,26 @@ export function prefixGatherIdx(P: number, Rc: number, Rr: number): { chosenIdx:
   return { chosenIdx, rejectedIdx };
 }
 
+/** sft_scope:"full" over the prefix-shared concat: the full-scope chosen NLL
+ *  gathers the PROMPT predictions too — hidden position t (t = 0..P-2) predicts
+ *  prompt token t+1 — and combines them with the response mean ℓw:
+ *  NLL_full = -((P-1)·promptMean + Rc·ℓw) / (P-1+Rc). The prompt head routes
+ *  through the same branchLogpMeanGathered tier as the responses (whole-vocab /
+ *  fused / flash), sharing the one concat forward's hiddens. Returns [1]
+ *  (caller owns); lw is NOT disposed. P=1 (no prompt predictions) → -ℓw. */
+export function prefixFullNll(
+  model: LogitProjector, h: MlxArray, promptIds: number[], lw: MlxArray, Rc: number, chunk?: ChunkCtx,
+): MlxArray {
+  const P = promptIds.length;
+  if (P <= 1) return combineFullNll(null, lw, 0, Rc);
+  const promptIdx = Array.from({ length: P - 1 }, (_, k) => k); // H[k] predicts promptIds[k+1]
+  const promptTargets = promptIds.slice(1);
+  const pm = branchLogpMeanGathered(model, h, promptIdx, promptTargets, chunk); // [1]
+  const nllFull = combineFullNll(pm, lw, P - 1, Rc);
+  pm.dispose();
+  return nllFull;
+}
+
 /** ORPO loss via the shared prompt-prefix single forward (MiniCPM5, B=1).
  *  `promptIds` is the shared prefix; `chosenResp`/`rejectedResp` are the response
  *  continuations. Returns the scalar loss (caller owns). Differentiable through
@@ -155,7 +198,7 @@ export function prefixGatherIdx(P: number, Rc: number, Rr: number): { chosenIdx:
 export function orpoLossPrefixShared(
   model: MiniCPM5Model,
   promptIds: number[], chosenResp: number[], rejectedResp: number[],
-  lambda: number, chunk?: ChunkCtx,
+  lambda: number, chunk?: ChunkCtx, sftScope: SftScope = "response",
 ): MlxArray {
   const P = promptIds.length, Rc = chosenResp.length, Rr = rejectedResp.length;
   if (P < 1 || Rc < 1 || Rr < 1) throw new Error("orpoLossPrefixShared: need P,Rc,Rr >= 1");
@@ -165,7 +208,7 @@ export function orpoLossPrefixShared(
   concat.set(chosenResp, P);
   concat.set(rejectedResp, P + Rc);
   const ids = MlxArray.fromInt32(concat, [1, T]);
-  const caches: Cache[] = model.layers.map(() => new PrefixSharedCache(P, Rc, Rr));
+  const caches: Cache[] = prefixSharedCaches(model.layers.length, P, Rc, Rr);
 
   let h: MlxArray;
   setMiniCpmPrefixPlan({ P, Rc, Rr });
@@ -188,10 +231,14 @@ export function orpoLossPrefixShared(
   if (keepForBwd) chunk!.sink.push(h);
   const lw = branchLogpMeanGathered(model, h, chosenIdx, chosenResp, chunk);
   const lr = branchLogpMeanGathered(model, h, rejectedIdx, rejectedResp, chunk);
+  // sft_scope:"full": add the prompt predictions (H[0..P-2] → promptIds[1..P-1])
+  // from the SAME concat forward; ℓw/ℓr stay response-only for the odds ratio.
+  const nllFull = sftScope === "full" ? prefixFullNll(model, h, promptIds, lw, Rc, chunk) : null;
   if (!keepForBwd) h.dispose();
-  const loss = orpoLossFromLogps(lw, lr, lambda);
+  const loss = orpoLossFromLogps(lw, lr, lambda, nllFull ?? undefined);
   lw.dispose();
   lr.dispose();
+  nllFull?.dispose();
   return loss;
 }
 
@@ -206,7 +253,7 @@ export function prefixSharedLogps(
   const concat = new Int32Array(T);
   concat.set(promptIds, 0); concat.set(chosenResp, P); concat.set(rejectedResp, P + Rc);
   const ids = MlxArray.fromInt32(concat, [1, T]);
-  const caches: Cache[] = model.layers.map(() => new PrefixSharedCache(P, Rc, Rr));
+  const caches: Cache[] = prefixSharedCaches(model.layers.length, P, Rc, Rr);
   setMiniCpmPrefixPlan({ P, Rc, Rr });
   let h: MlxArray;
   try { h = model.forwardHidden(ids, caches); }
@@ -349,7 +396,7 @@ class Gemma4PrefixSharedCache implements Cache {
 export function orpoLossPrefixSharedGemma(
   model: Gemma4Model,
   promptIds: number[], chosenResp: number[], rejectedResp: number[],
-  lambda: number, chunk?: ChunkCtx,
+  lambda: number, chunk?: ChunkCtx, sftScope: SftScope = "response",
 ): MlxArray {
   const P = promptIds.length, Rc = chosenResp.length, Rr = rejectedResp.length;
   if (P < 1 || Rc < 1 || Rr < 1) throw new Error("orpoLossPrefixSharedGemma: need P,Rc,Rr >= 1");
@@ -384,10 +431,14 @@ export function orpoLossPrefixSharedGemma(
   if (keepForBwd) chunk!.sink.push(h);
   const lw = branchLogpMeanGathered(model, h, chosenIdx, chosenResp, chunk);
   const lr = branchLogpMeanGathered(model, h, rejectedIdx, rejectedResp, chunk);
+  // sft_scope:"full": prompt predictions from the same concat forward (see
+  // orpoLossPrefixShared); ℓw/ℓr stay response-only for the odds ratio.
+  const nllFull = sftScope === "full" ? prefixFullNll(model, h, promptIds, lw, Rc, chunk) : null;
   if (!keepForBwd) h.dispose();
-  const loss = orpoLossFromLogps(lw, lr, lambda);
+  const loss = orpoLossFromLogps(lw, lr, lambda, nllFull ?? undefined);
   lw.dispose();
   lr.dispose();
+  nllFull?.dispose();
   return loss;
 }