mlx-bun 0.0.5 → 0.0.6

package/src/quantize/quantizer.ts

@@ -392,8 +392,28 @@ async function computeMixedAllocation(
 
   progress("sensitivity", "Loading model for sensitivity analysis", 0.0);
 
-  const config = await loadModelConfig(srcDir);
-  const weights = await Weights.open(srcDir);
+  // The sensitivity sweep runs FORWARD passes, which need a model the runtime
+  // can load — and mlx-bun's model graph can only load QUANTIZED weights. If the
+  // source is bf16 (unquantized), build a temporary uniform 8-bit baseline
+  // (near-lossless) to probe against. The FINAL weights are still quantized from
+  // the original bf16 source per the allocation (quantizeModelDir reads srcDir),
+  // so fidelity is not capped by the baseline — the baseline only anchors the
+  // per-layer sensitivity (8→4-bit) measurement.
+  const srcConfig = await loadModelConfig(srcDir);
+  let probeDir = srcDir;
+  let tempBaseline: string | null = null;
+  if (!srcConfig.quantization) {
+    const { mkdtempSync } = await import("node:fs");
+    const { tmpdir } = await import("node:os");
+    const { join } = await import("node:path");
+    tempBaseline = mkdtempSync(join(tmpdir(), "mlx-bun-qbaseline-"));
+    progress("sensitivity", "Unquantized source — building an 8-bit baseline for the sweep", 0.0);
+    await quantizeModelDir(srcDir, tempBaseline, { bits: 8, groupSize, mode: "affine" });
+    probeDir = tempBaseline;
+  }
+
+  const config = await loadModelConfig(probeDir);
+  const weights = await Weights.open(probeDir);
   try {
     const model = createModel(weights, config);
     const tokenizer = await loadTokenizer(srcDir);
@@ -458,5 +478,9 @@ async function computeMixedAllocation(
   } finally {
     weights.dispose();
     clearCache();
+    if (tempBaseline) {
+      const { rmSync } = await import("node:fs");
+      rmSync(tempBaseline, { recursive: true, force: true });
+    }
   }
 }

package/src/server.ts

@@ -667,18 +667,49 @@ export function promptEndsInOpenThink(rendered: string): boolean {
   return open !== -1 && open > rendered.lastIndexOf("</think>");
 }
 
+/** Concatenate the text of an OpenAI content-part array, ignoring non-text
+ *  parts. Tolerant of the part shapes clients actually send: `{type:"text",
+ *  text}` (OpenAI/pi), `{type:"input_text", text}` (Responses-style), or any
+ *  part carrying a string `text`. */
+function contentPartsToText(parts: Array<Record<string, unknown>>): string {
+  return parts
+    .map((p) => (p && typeof p.text === "string" ? p.text : ""))
+    .join("");
+}
+
+/** True if any content part is an image (so the vision path must keep the
+ *  array form for extractImages). */
+function hasImagePart(parts: Array<Record<string, unknown>>): boolean {
+  return parts.some((p) => p && (p.type === "image" || p.type === "image_url"));
+}
+
 /** OpenAI sends assistant tool_call arguments as JSON strings; the
  *  template renders the object form natively — normalize before render.
- *  Also maps the OpenAI reasoning-model "developer" role to "system": pi-ai
- *  (and OpenAI's own SDKs) rename the system prompt to `developer` whenever a
- *  model advertises reasoning, but our chat templates only know
- *  system/user/assistant/tool and raise "Unexpected message role." otherwise.
- *  This is exactly why Qwen3.5/MiniCPM5 chat got no messages while Gemma (a
- *  non-reasoning model, so pi keeps `system`) worked. `developer` IS the
- *  instruction/system prompt, so the remap is semantics-preserving. */
+ *
+ *  Two more wire-format → template-format fixes, both "match the format the
+ *  model's chat template expects":
+ *
+ *  1. Map the OpenAI reasoning-model "developer" role to "system": pi-ai (and
+ *     OpenAI's own SDKs) rename the system prompt to `developer` whenever a
+ *     model advertises reasoning, but our chat templates only know
+ *     system/user/assistant/tool and raise "Unexpected message role." This is
+ *     why Qwen3.5/MiniCPM5 chat got no messages while Gemma (non-reasoning, so
+ *     pi keeps `system`) worked. `developer` IS the system prompt, so the remap
+ *     is semantics-preserving.
+ *
+ *  2. Flatten text-only content-part arrays to a plain string. pi (and any
+ *     OpenAI multimodal client) sends user content as `[{type:"text",text}]`,
+ *     but non-vision chat templates expect `content` to be a STRING and render
+ *     nothing for an array — so the user's turn silently vanishes and the model
+ *     replies "I don't see any message." Arrays that carry an image part are
+ *     left intact for the vision path (extractImages); only text-only arrays
+ *     are collapsed here, which is a no-op for the vision templates too. */
 export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
   return messages.map((raw) => {
-    const m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
+    let m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
+    if (Array.isArray(m.content) && !hasImagePart(m.content)) {
+      m = { ...m, content: contentPartsToText(m.content) };
+    }
     if (!m.tool_calls) return m;
     return {
       ...m,
@@ -880,43 +911,61 @@ export function createServer(
         `Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
     );
 
-  const toOptions = (req: ChatRequest): GenerateOptions & { stopSequences: string[] } => ({
-    maxTokens: req.max_completion_tokens ?? req.max_tokens ?? 1024,
-    temperature: req.temperature ?? serverOptions.defaultTemperature ?? ctx.genDefaults.temperature ?? 0.7,
-    topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
-    topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
-    seed: req.seed ?? (Date.now() & 0xffffffff),
-    repetitionPenalty: req.repetition_penalty ?? ctx.genDefaults.repetitionPenalty,
-    hlg: resolveHlg(req.hlg, serverOptions.hlg),
-    // generate() yields token ids; stop sequences match on decoded text
-    // (they can span token boundaries), so the StopMatcher sits at the
-    // decode layer below and halts the loop via onToken → false.
-    stopSequences: (typeof req.stop === "string" ? [req.stop] : req.stop ?? [])
-      .filter((s) => typeof s === "string" && s.length > 0),
-    ...kvScheme,
-  });
-
-  const templateOptionsFor = (req: ChatRequest, tools: ToolDefinition[] | null) => {
-    // Precedence: explicit per-request chat_template_kwargs wins; else the
-    // standard reasoning_effort (the OpenAI field Pi sends — "none" → no
-    // thinking, any level → thinking on); else the server-wide --thinking
-    // default; else the model's own default (false for MiniCPM5, otherwise
-    // leave the template's default). The Anthropic/Responses shims map their
-    // reasoning fields onto reasoning_effort before reaching here.
+  // Effective enable_thinking for a request, with the same precedence the chat
+  // template uses (extracted so sampling and template rendering can't disagree):
+  // explicit chat_template_kwargs.enable_thinking → reasoning_effort ("none" =
+  // off) → server --thinking default → model default (MiniCPM5 → off). undefined
+  // means "not a switchable-thinking model / leave the template default".
+  const resolveEnableThinking = (req: ChatRequest): boolean | undefined => {
     const explicit = req.chat_template_kwargs?.enable_thinking;
+    if (typeof explicit === "boolean") return explicit;
     const effort = req.reasoning_effort;
+    if (effort !== undefined) return effort !== "none";
+    if (serverOptions.defaultThinking !== undefined) return serverOptions.defaultThinking;
+    return isMiniCPM5Config(ctx.model.config) ? false : undefined;
+  };
+
+  const toOptions = (req: ChatRequest): GenerateOptions & { stopSequences: string[] } => {
+    // Sampling follows the thinking state — which the web UI's thinking button
+    // drives via enable_thinking. Model authors publish a SINGLE
+    // generation_config temperature (the think-mode value) but recommend a
+    // cooler one for direct, no-think replies (MiniCPM5 card: 0.9 think / 0.7
+    // no-think, top_p 0.95 for both). So with no explicit temperature set, a
+    // no-think turn runs at most NO_THINK_TEMPERATURE while a think turn keeps
+    // the model's hotter configured default. An explicit request/CLI
+    // temperature always wins. top_p is unchanged by mode.
+    const NO_THINK_TEMPERATURE = 0.7;
+    const genTemp = ctx.genDefaults.temperature ?? 0.7;
+    const defaultTemp = resolveEnableThinking(req) === false
+      ? Math.min(genTemp, NO_THINK_TEMPERATURE)
+      : genTemp;
     return {
-      tools,
-      enableThinking: typeof explicit === "boolean"
-        ? explicit
-        : effort !== undefined
-          ? effort !== "none"
-          : serverOptions.defaultThinking !== undefined
-            ? serverOptions.defaultThinking
-            : isMiniCPM5Config(ctx.model.config) ? false : undefined,
+      // A thinking turn emits <think>…</think> AND the answer, so a tight cap
+      // truncates the visible reply. Default very generously (the model's
+      // context is far larger); only an explicit max_tokens narrows it. The
+      // model still stops at its eos_token well before this in normal replies.
+      maxTokens: req.max_completion_tokens ?? req.max_tokens ?? 65_536,
+      temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
+      topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
+      topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
+      seed: req.seed ?? (Date.now() & 0xffffffff),
+      repetitionPenalty: req.repetition_penalty ?? ctx.genDefaults.repetitionPenalty,
+      hlg: resolveHlg(req.hlg, serverOptions.hlg),
+      // generate() yields token ids; stop sequences match on decoded text
+      // (they can span token boundaries), so the StopMatcher sits at the
+      // decode layer below and halts the loop via onToken → false.
+      stopSequences: (typeof req.stop === "string" ? [req.stop] : req.stop ?? [])
+        .filter((s) => typeof s === "string" && s.length > 0),
+      ...kvScheme,
     };
   };
 
+  const templateOptionsFor = (req: ChatRequest, tools: ToolDefinition[] | null) => {
+    // enableThinking resolution (and its precedence) lives in
+    // resolveEnableThinking so template rendering and sampling stay in sync.
+    return { tools, enableThinking: resolveEnableThinking(req) };
+  };
+
   const promptIdsFor = (
     req: ChatRequest,
     tools: ToolDefinition[] | null,
@@ -980,6 +1029,16 @@ export function createServer(
         try { html = readFileSync(new URL("./assets/curve-designer.html", import.meta.url), "utf8"); } catch { /* binary: use embedded */ }
         return new Response(html, { headers: { "content-type": "text/html; charset=utf-8", "cache-control": "no-store" } });
       }
+      // Training/inference DAG map — self-contained cytoscape artifact, served
+      // same-origin so the Routes tab can embed it in an <iframe src="/dag">.
+      if (url.pathname === "/dag" && request.method === "GET") {
+        try {
+          const html = readFileSync(new URL("../docs/dag/training-inference-map.html", import.meta.url), "utf8");
+          return new Response(html, { headers: { "content-type": "text/html; charset=utf-8", "cache-control": "no-store" } });
+        } catch {
+          return new Response("DAG map artifact not found at docs/dag/training-inference-map.html", { status: 404 });
+        }
+      }
       if (url.pathname === "/curve-terrain" && request.method === "GET") {
         try {
           const html = readFileSync(new URL("../docs/investigations/curve-terrain.html", import.meta.url), "utf8");
@@ -1000,7 +1059,11 @@ export function createServer(
           const { Registry } = await import("./registry");
           const { loadModelConfig } = await import("./config");
           const reg = new Registry();
-          if (reg.list().length === 0) await reg.scan();
+          // Always rescan when (re)building the 30s cache so models that appeared
+          // after boot — fresh downloads AND quants written into the HF cache —
+          // surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
+          // so it's idempotent and cheap for a local cache.)
+          await reg.scan();
           const rows = [];
           for (const m of reg.list()) {
             const supported = isSupportedModelRecord(m.modelType, m.repoId);
@@ -1747,6 +1810,78 @@ export function createServer(
         const { inspectModel } = await import("./quantize");
         return Response.json(await inspectModel(body.model_id ?? ""));
       }
+      // Turn an OS-picked folder into the absolute snapshot path on disk. The
+      // browser can't reveal a filesystem path (security) and the HF cache
+      // stores real bytes in blobs/ with symlinked snapshots — so we resolve by
+      // the folder's NAME (which encodes the repo id), never by reading files.
+      // No upload, no dependence on the cache's symlink layout, and a
+      // just-downloaded model resolves before it's ever been indexed.
+      if ((url.pathname === "/api/quantize/resolve-folder" || url.pathname === "/api/model/resolve-folder") && request.method === "POST") {
+        const body = (await request.json().catch(() => ({}))) as { folder_name?: string; rel_path?: string };
+        const { statSync, readdirSync, readFileSync } = await import("node:fs");
+        const { join } = await import("node:path");
+        const { homedir } = await import("node:os");
+        const hubRoot = process.env.HF_HUB_CACHE
+          ?? (process.env.HF_HOME ? join(process.env.HF_HOME, "hub") : join(homedir(), ".cache/huggingface/hub"));
+        const roots = [hubRoot, join(homedir(), ".cache/mlx-bun")];
+        const hasConfig = (d: string) => { try { return statSync(join(d, "config.json")).isFile(); } catch { return false; } };
+        // basename of the picked folder; the rel path's last-but-one segment is
+        // the <hash> dir if they drilled into a snapshot.
+        const folder = (body.folder_name ?? "").replace(/\\/g, "/").split("/").filter(Boolean).pop() ?? "";
+        const relSegs = (body.rel_path ?? "").replace(/\\/g, "/").split("/").filter(Boolean);
+        const configDir = relSegs.length >= 2 ? relSegs[relSegs.length - 2] : "";
+        const repoIdOf = (modelsDir: string) => modelsDir.slice("models--".length).replaceAll("--", "/");
+        // Pick a repo dir's canonical snapshot: the one refs/main points at, else
+        // the first snapshot that actually has a config.json.
+        const pickSnapshot = (repoDir: string): string | null => {
+          const snaps = join(repoDir, "snapshots");
+          let head = "";
+          try { head = readFileSync(join(repoDir, "refs", "main"), "utf8").trim(); } catch {}
+          if (head && hasConfig(join(snaps, head))) return join(snaps, head);
+          try { for (const h of readdirSync(snaps)) if (hasConfig(join(snaps, h))) return join(snaps, h); } catch {}
+          return null;
+        };
+
+        // (1) HF-cache folder picked — its name IS the repo id.
+        if (folder.startsWith("models--")) {
+          const p = pickSnapshot(join(hubRoot, folder));
+          if (p) return Response.json({ ok: true, path: p, repo_id: repoIdOf(folder) });
+        }
+        // (2) A bare <hash> snapshot folder picked directly (or carried in rel).
+        const hashDir = configDir || folder;
+        if (hashDir) {
+          try {
+            for (const repo of readdirSync(hubRoot)) {
+              if (!repo.startsWith("models--")) continue;
+              const cand = join(hubRoot, repo, "snapshots", hashDir);
+              if (hasConfig(cand)) return Response.json({ ok: true, path: cand, repo_id: repoIdOf(repo) });
+            }
+          } catch {}
+        }
+        // (3) Flat local snapshot under a known root (folder name == model dir).
+        for (const root of roots) {
+          if (folder && hasConfig(join(root, folder))) return Response.json({ ok: true, path: join(root, folder) });
+        }
+        // (4) Registry rescan — last resort (also covers customized cache layouts).
+        const { Registry } = await import("./registry");
+        const reg = new Registry();
+        try {
+          await reg.scan();
+          const all = reg.list();
+          const rec = (folder.startsWith("models--")
+            ? all.find((m) => m.repoId === repoIdOf(folder))
+            : undefined)
+            ?? all.find((m) => m.path.split("/").pop() === hashDir)
+            ?? all.find((m) => m.repoId.split("/").pop() === folder);
+          if (rec) return Response.json({ ok: true, path: rec.path, repo_id: rec.repoId, model_type: rec.modelType });
+        } finally {
+          reg.close();
+        }
+        return Response.json({
+          ok: false,
+          error: "Couldn't locate this folder on disk — paste the path instead.",
+        });
+      }
       if (url.pathname === "/api/quantize/submit" && request.method === "POST") {
         const body = (await request.json().catch(() => ({}))) as {
           model_id?: string; bits?: number; group_size?: number;
@@ -1758,18 +1893,52 @@ export function createServer(
         const store = await ensureJobs();
         const { submitSubprocess } = await import("./jobs");
         const { homedir } = await import("node:os");
+        const { join } = await import("node:path");
+        const { mkdirSync, writeFileSync } = await import("node:fs");
+        const { createHash } = await import("node:crypto");
         const bits = body.bits ?? 4, gs = body.group_size ?? 64;
-        const base = body.model_id.split("/").filter(Boolean).at(-1)!.replace(/[^a-z0-9_.-]/gi, "");
-        // Mixed-precision (sensitivity+knapsack) names the dir by target bpw.
+        // Derive org/name from the source so the quant is named readably (the
+        // source is usually an HF-cache snapshot path whose basename is a hash).
+        const snapMatch = body.model_id.match(/(models--[^/]+)\/snapshots\//);
+        let org = "local", name: string;
+        if (snapMatch) {
+          const parts = snapMatch[1].split("--"); // ["models", org, ...name]
+          org = parts[1] ?? "local";
+          name = parts.slice(2).join("--");
+        } else if (body.model_id.includes("/") && !body.model_id.startsWith("/") && !body.model_id.startsWith("~")) {
+          const seg = body.model_id.split("/"); // a repo id "org/name"
+          org = seg[0]!; name = seg.slice(1).join("-");
+        } else {
+          name = body.model_id.split("/").filter(Boolean).at(-1) ?? "model"; // a bare path
+        }
+        name = (name || "model").replace(/[^a-z0-9_.-]/gi, "");
+        org = (org || "local").replace(/[^a-z0-9_.-]/gi, "");
         const suffix = body.target_bpw ? `mixed-${body.target_bpw}bpw` : `${bits}bit`;
-        const outDir = `${homedir()}/.cache/mlx-bun/quants/${base}-OptiQ-${suffix}`;
+        // Write the quant INTO the HF hub cache as a normal models--org--name/
+        // snapshots/<hash> entry, so the standard registry scan + every other
+        // tool discovers it alongside downloaded models. refs/main makes it a
+        // well-formed cache entry.
+        const quantRepo = `${name}-OptiQ-${suffix}`;
+        const hubRoot = process.env.HF_HUB_CACHE
+          ?? (process.env.HF_HOME ? join(process.env.HF_HOME, "hub") : join(homedir(), ".cache/huggingface/hub"));
+        const repoDir = join(hubRoot, `models--${org}--${quantRepo}`);
+        const snapHash = createHash("sha1").update(`${org}/${quantRepo}`).digest("hex");
+        const outDir = join(repoDir, "snapshots", snapHash);
+        try {
+          mkdirSync(join(repoDir, "refs"), { recursive: true });
+          writeFileSync(join(repoDir, "refs", "main"), snapHash);
+        } catch {}
         const { jobId } = submitSubprocess(store, "quantize", {
           model_id: body.model_id, out_dir: outDir, bits, group_size: gs,
           // forwarded to the mixed-precision path when target_bpw is set
           target_bpw: body.target_bpw, candidate_bits: body.candidate_bits,
           reference: body.reference, calibration_mix: body.calibration_mix,
           n_calibration: body.n_calibration,
-        }, outDir);
+        }, outDir, {
+          // The model is in the HF cache now; drop the Library cache so the next
+          // poll rescans and shows it — no `mlx-bun scan`, no restart.
+          onComplete: () => { libraryCache = null; },
+        });
         return Response.json({ ok: true, job_id: jobId, output_dir: outDir });
       }
 

package/src/tokenizer.ts

@@ -6,7 +6,10 @@
 import { Tokenizer } from "@huggingface/tokenizers";
 
 export interface LoadedTokenizer {
-  encode(text: string): number[];
+  /** Encode text → token ids. `addSpecialTokens` (default true) prepends the BOS;
+   *  pass FALSE for chat-template output, which already emits the BOS as text (else
+   *  you get a corrupting double-BOS `<s><s>…`). */
+  encode(text: string, addSpecialTokens?: boolean): number[];
   decode(ids: number[], skipSpecialTokens?: boolean): string;
   readonly bosTokenId: number | null;
   readonly eosTokenId: number | null;
@@ -28,7 +31,8 @@ export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer>
   };
 
   return {
-    encode: (text) => tok.encode(text).ids.map(Number),
+    encode: (text, addSpecialTokens = true) =>
+      tok.encode(text, { add_special_tokens: addSpecialTokens }).ids.map(Number),
     decode: (ids, skipSpecialTokens = false) =>
       // python's decode([]) === ""; the JS lib throws on empty input
       ids.length === 0 ? "" : tok.decode(ids, { skip_special_tokens: skipSpecialTokens }),

package/src/train/dataset.ts

@@ -269,14 +269,22 @@ export function tokenizePair(
   tok: LoadedTokenizer,
   maxLength: number,
 ): { ids: number[]; promptLen: number } {
-  const promptIds = tok.encode(prompt);
-  let fullIds = tok.encode(prompt + response);
+  // A chat template emits a literal BOS, and the tokenizer's post-processor
+  // prepends one too → a duplicate leading BOS. Inference strips it
+  // (server.ts: "template includes <bos>; tokenizer also prepends one"), so
+  // training must match or it learns on a token stream the model never sees.
+  // No-op for raw (untemplated) prompts, which carry a single BOS.
+  const bos = tok.bosTokenId;
+  const dedupeBos = (ids: number[]): number[] =>
+    bos != null && ids.length >= 2 && ids[0] === bos && ids[1] === bos ? ids.slice(1) : ids;
+  const promptIds = dedupeBos(tok.encode(prompt));
+  let fullIds = dedupeBos(tok.encode(prompt + response));
   let promptLen: number;
   if (isPrefix(fullIds, promptIds)) {
     promptLen = promptIds.length;
   } else {
     promptLen = promptIds.length;
-    const responseIds = tok.encode(response);
+    const responseIds = tok.encode(response, false); // no extra BOS mid-stream
     fullIds = [...promptIds, ...responseIds];
   }
   if (fullIds.length > maxLength) {

package/src/train/flash-cce.ts

@@ -927,7 +927,13 @@ function fwdKernel(): MetalKernel {
 }
 
 function u32(values: number[]): MlxArray {
-  return MlxArray.fromView(new Uint8Array(new Uint32Array(values).buffer.slice(0)), [values.length], Dtype.uint32);
+  // COPY into an mlx-owned array (fromBytesCopy = mlx_array_new_data), NOT the
+  // pinning fromView. These tiny kernel-arg arrays (shape, targets) feed a LAZY
+  // kernel and are disposed before it evaluates; pinning leaks one host buffer
+  // per call, never released (unbounded `pinned` map + a deferred native
+  // use-after-free under memory pressure — the seg/accum crash). A few-int copy
+  // is free; correctness is not.
+  return MlxArray.fromBytesCopy(new Uint8Array(new Uint32Array(values).buffer), [values.length], Dtype.uint32);
 }
 
 /** Flash-CCE forward: per-token response log-prob [M] (f32) + lse [M] (f32),
@@ -1064,12 +1070,19 @@ export function flashCceForward(
     // reduction is simpler and just as exact.)
     const gMax = ops.maxAxis(bMax, 1, false); // [M]
     const gMaxCol = ops.reshape(gMax, [M, 1]);
-    const shifted = ops.exp(ops.sub(bMax, gMaxCol)); // [M, NBLK]
-    const gSum = ops.sumAxis(ops.mul(bSum, shifted), 1, false); // [M]
-    const lse = ops.add(gMax, ops.log(gSum)); // [M]
+    // Bind every nested intermediate so its MlxArray wrapper gets disposed — a
+    // bare `ops.exp(ops.sub(...))` leaks the inner `sub` (stays live in `active`
+    // until the GC finalizer backstop), and this runs TWICE per step (chosen +
+    // rejected). gMax was also missing from the dispose list below.
+    const sub = ops.sub(bMax, gMaxCol);
+    const shifted = ops.exp(sub); // [M, NBLK]
+    const prod = ops.mul(bSum, shifted);
+    const gSum = ops.sumAxis(prod, 1, false); // [M]
+    const logG = ops.log(gSum);
+    const lse = ops.add(gMax, logG); // [M]
     const tgt = ops.sumAxis(bTgt, 1, false); // [M] (owning block only is nonzero)
     const logp = ops.sub(tgt, lse);
-    for (const a of [gMaxCol, shifted, gSum, tgt]) a.dispose();
+    for (const a of [gMax, gMaxCol, sub, shifted, prod, gSum, logG, tgt]) a.dispose();
     const blockMax = bMax!; bMax = undefined; // hand off to caller (backward block-skip)
     return { logp, lse, blockMax };
   } finally {

package/src/train/job.ts

@@ -39,6 +39,8 @@ interface FinetuneSubmit {
   mlp_split?: boolean;
   segment_size?: number;
   save_checkpoints?: boolean;
+  grad_clip_norm?: number;
+  val_max_examples?: number;
   dpo_beta?: number;
   dpo_warmup_iters?: number;
   dpo_lr_schedule?: "constant" | "cosine";
@@ -58,6 +60,14 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
   if (!c.data_dir) throw new Error("finetune job: missing data_dir");
   if (!c.adapter_path) throw new Error("finetune job: missing adapter_path");
 
+  // For ORPO the full L3 stack is the validated production recipe (what the
+  // `mlx-bun train` CLI turns on): flash-CCE head + prefix-sharing + segmented
+  // backward + token-chunked head. DEFAULT_TRAIN_CONFIG keeps them OFF (they're
+  // opt-in for SFT), so default them ON here when method=orpo — an explicit
+  // value from the caller (CLI/API) still wins via ??. finetuneRunner drops the
+  // quantized-head pieces if the base turns out to be unquantized.
+  const isOrpo = (c.method ?? "sft") === "orpo";
+
   const cfg: TrainConfig = {
     ...DEFAULT_TRAIN_CONFIG,
     method: c.method ?? "sft",
@@ -81,18 +91,20 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
     loraPlusRatio: c.lora_plus_ratio ?? DEFAULT_TRAIN_CONFIG.loraPlusRatio,
     gradCheckpoint: c.grad_checkpoint ?? DEFAULT_TRAIN_CONFIG.gradCheckpoint,
     mlpSplit: c.mlp_split ?? DEFAULT_TRAIN_CONFIG.mlpSplit,
-    segmentSize: c.segment_size ?? DEFAULT_TRAIN_CONFIG.segmentSize,
+    segmentSize: c.segment_size ?? (isOrpo ? 2 : DEFAULT_TRAIN_CONFIG.segmentSize),
     saveCheckpoints: c.save_checkpoints ?? DEFAULT_TRAIN_CONFIG.saveCheckpoints,
+    gradClipNorm: c.grad_clip_norm ?? DEFAULT_TRAIN_CONFIG.gradClipNorm,
+    valMaxExamples: c.val_max_examples ?? DEFAULT_TRAIN_CONFIG.valMaxExamples,
     dpoBeta: c.dpo_beta ?? DEFAULT_TRAIN_CONFIG.dpoBeta,
     dpoWarmupIters: c.dpo_warmup_iters ?? DEFAULT_TRAIN_CONFIG.dpoWarmupIters,
     dpoLrSchedule: c.dpo_lr_schedule ?? DEFAULT_TRAIN_CONFIG.dpoLrSchedule,
     orpoLambda: c.orpo_lambda ?? DEFAULT_TRAIN_CONFIG.orpoLambda,
     orpoWarmupIters: c.orpo_warmup_iters ?? DEFAULT_TRAIN_CONFIG.orpoWarmupIters,
     orpoLrSchedule: c.orpo_lr_schedule ?? DEFAULT_TRAIN_CONFIG.orpoLrSchedule,
-    orpoChunkSize: c.orpo_chunk_size ?? DEFAULT_TRAIN_CONFIG.orpoChunkSize,
+    orpoChunkSize: c.orpo_chunk_size ?? (isOrpo ? 512 : DEFAULT_TRAIN_CONFIG.orpoChunkSize),
     orpoFusedCe: c.orpo_fused_ce ?? DEFAULT_TRAIN_CONFIG.orpoFusedCe,
-    orpoFlashCe: c.orpo_flash_ce ?? DEFAULT_TRAIN_CONFIG.orpoFlashCe,
-    orpoPrefixShared: c.orpo_prefix_shared ?? DEFAULT_TRAIN_CONFIG.orpoPrefixShared,
+    orpoFlashCe: c.orpo_flash_ce ?? (isOrpo ? true : DEFAULT_TRAIN_CONFIG.orpoFlashCe),
+    orpoPrefixShared: c.orpo_prefix_shared ?? (isOrpo ? true : DEFAULT_TRAIN_CONFIG.orpoPrefixShared),
     warmStartAdapter: c.warm_start_adapter ?? DEFAULT_TRAIN_CONFIG.warmStartAdapter,
     adapterPath: c.adapter_path,
     baseModel: c.model_dir,
@@ -103,6 +115,18 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
 export const finetuneRunner: JobRunner = async (emit, config) => {
   const { modelDir, dataDir, cfg } = parseConfig(config);
 
+  // Gemma/e4b needs these training env flags set BEFORE the model loads — same
+  // as the `mlx-bun train` CLI. The subprocess job path doesn't inherit them
+  // from the server, so without this a web-submitted ORPO/SFT run on Gemma uses
+  // the wrong kernels. ??= so an explicit env (or the CLI) still wins.
+  try {
+    const isGemma = (await Bun.file(`${modelDir}/config.json`).text()).toLowerCase().includes("gemma");
+    if (isGemma) {
+      process.env.MLX_BUN_PERF_KERNEL ??= "0";
+      process.env.MLX_BUN_FUSED_GELU ??= "0";
+    }
+  } catch {}
+
   emit({ type: "stage", stage: "load", progress: 0.01, message: `loading model ${modelDir}` });
 
   const { loadModelConfig } = await import("../config");
@@ -112,6 +136,19 @@ export const finetuneRunner: JobRunner = async (emit, config) => {
   const { ChatTemplate } = await import("../chat-template");
 
   const modelConfig = await loadModelConfig(modelDir);
+
+  // The flash-CCE / fused-CE heads dequantize a QUANTIZED head in-kernel; on an
+  // unquantized (bf16) base they can't run, so fall back to the full-logits head.
+  // (Quantize the base first to get the large-vocab memory win — see the
+  // Quantize tab.)
+  if (cfg.method === "orpo" && !modelConfig.quantization && (cfg.orpoFlashCe || cfg.orpoFusedCe)) {
+    emit({
+      type: "stage", stage: "load",
+      message: "unquantized base — flash/fused-CE head needs a quantized head; using full-logits head",
+    });
+    cfg.orpoFlashCe = false;
+    cfg.orpoFusedCe = false;
+  }
   const weights = await Weights.open(modelDir);
   const model = createModel(weights, modelConfig);
   const tok = await loadTokenizer(modelDir);

package/src/train/loss.ts

@@ -839,7 +839,8 @@ function makeVocabBlockedHeadVjp(
     const raw = ops.quantizedMatmul(hc, wb.w, wb.scales, wb.biases, head.spec, true); // [Cc, Vb]
     if (cap === null) return { lb: raw, sech2: null };
     const capArr = ops.scalarLike(cap, raw);
-    const th = ops.tanh(ops.div(raw, capArr));
+    const div = ops.div(raw, capArr);
+    const th = ops.tanh(div); div.dispose(); // bind the nested div so its wrapper is freed
     raw.dispose();
     const lb = ops.mul(th, capArr); // cap·tanh(raw/cap)
     const th2 = ops.mul(th, th); th.dispose();
@@ -964,7 +965,10 @@ function makeVocabBlockedHeadVjp(
           const colFlat = ops.arange(v0, v1, 1, Dtype.int32);
           const col = ops.reshape(colFlat, [1, Vb]); colFlat.dispose();
           const tcol = MlxArray.fromInt32(targetsHost.slice(c0, c1), [Cc, 1]);
-          const eqMask = ops.logicalAnd(ops.greaterEqual(col, tcol), ops.greaterEqual(tcol, col));
+          const ge1 = ops.greaterEqual(col, tcol);
+          const ge2 = ops.greaterEqual(tcol, col);
+          const eqMask = ops.logicalAnd(ge1, ge2); // bind the nested compares so their wrappers free
+          ge1.dispose(); ge2.dispose();
           col.dispose(); tcol.dispose();
           const onehot = eqMask.astype(softmax.dtype); eqMask.dispose();
           let g = ops.sub(onehot, softmax); // onehot − softmax  [Cc, Vb]
@@ -1000,7 +1004,8 @@ function makeVocabBlockedHeadVjp(
 /** flash-CCE head as a CustomVjp: forward runs the Metal kernel → Σ logp (stashing
  *  `lse` for the backward, avoiding a forward recompute); backward runs the kernel's
  *  `dh` for unit cotangent then scales by the scalar cotangent `c` (loss = Σ logp →
- *  ∂loss/∂logp_t = c for all t). The coeff filter is on by default in the backward. */
+ *  ∂loss/∂logp_t = c for all t). The coeff filter + vocab-block skip are OFF by
+ *  default (env-gated MLX_BUN_CCE_BWD_*_EPS) — the backward computes the EXACT dh. */
 function makeFlashCceHeadVjp(
   head: HeadQuant, targetsHost: Int32Array, M: number,
   sink: Array<{ dispose(): void }>,
@@ -1031,7 +1036,8 @@ function makeFlashCceHeadVjp(
       const c = cots[0]!;      // [1] = ∂loss/∂(Σ logp); uniform over tokens
       // dh for unit cotangent, then × c (avoids a host read of the cotangent).
       const ones = new Array(M).fill(1) as number[];
-      // filter + block-skip default-on (both approximations the training path accepts).
+      // filter + block-skip are OFF by default (env-gated) → EXACT dh. savedBlockMax is
+      // passed so the lossless vocab-block skip can be enabled via env without re-plumbing.
       const dhUnit = flashCceBackward(hR, fch, targets, savedLse!, ones, undefined, savedBlockMax!); // [M, hidden] f32
       const cM = c.dtype === dhUnit.dtype ? c : c.astype(dhUnit.dtype);
       const dh = ops.mul(dhUnit, cM); // broadcast [M,hidden]·[1]