mlx-bun 0.0.4 → 0.0.5

package/src/eval/ifeval.ts

@@ -0,0 +1,138 @@
+// IFEval — programmatically-verifiable instruction-following scoring (Zhou et
+// al. 2023, "Instruction-Following Evaluation for Large Language Models"). Each
+// instance pairs a prompt with one or more *verifiable* instructions; a response
+// is scored by deterministic checks (no LLM judge), so before/after ORPO
+// comparisons are exact and reproducible.
+//
+// This implements the high-frequency instruction families (a working subset of
+// the ~25 in the reference suite, extensible via INSTRUCTIONS). Strict scoring:
+// an instance passes (strict) iff every instruction passes; we also report the
+// per-instruction pass rate. The official metrics are prompt-level and
+// instruction-level accuracy — both derived from `scoreInstance`.
+
+export interface IFEvalInstance {
+  /** The prompt text (used by instructions that reference it, e.g. "repeat the request"). */
+  prompt: string;
+  /** Instruction ids, e.g. "keywords:existence". */
+  instruction_id_list: string[];
+  /** Per-instruction kwargs, aligned with instruction_id_list. */
+  kwargs: Array<Record<string, unknown>>;
+}
+
+export interface InstanceResult {
+  /** One pass/fail per instruction (aligned with instruction_id_list). */
+  perInstruction: boolean[];
+  /** Strict: all instructions passed. */
+  followedAll: boolean;
+}
+
+type Checker = (response: string, kw: Record<string, unknown>, inst: IFEvalInstance) => boolean;
+
+const words = (s: string): string[] => s.trim().split(/\s+/).filter(Boolean);
+const sentences = (s: string): string[] =>
+  s.split(/(?<=[.!?])\s+/).map((x) => x.trim()).filter(Boolean);
+const paragraphs = (s: string): string[] =>
+  s.split(/\n\s*\n/).map((x) => x.trim()).filter(Boolean);
+const norm = (s: string): string => s.toLowerCase();
+
+const cmp = (n: number, rel: string | undefined, target: number): boolean => {
+  switch (rel) {
+    case "at least": return n >= target;
+    case "at most": return n <= target;
+    case "less than": return n < target;
+    case "exactly": return n === target;
+    default: return n >= target; // IFEval default for "around"/unspecified is lenient ≥
+  }
+};
+
+/** The verifiable-instruction registry. Keyed by IFEval instruction_id. */
+export const INSTRUCTIONS: Record<string, Checker> = {
+  // ---- keywords ----
+  "keywords:existence": (r, kw) =>
+    (kw.keywords as string[]).every((k) => norm(r).includes(norm(k))),
+  "keywords:frequency": (r, kw) => {
+    const k = norm(kw.keyword as string);
+    const count = norm(r).split(k).length - 1;
+    return cmp(count, kw.relation as string, kw.frequency as number);
+  },
+  "keywords:forbidden_words": (r, kw) =>
+    !(kw.forbidden_words as string[]).some((k) => norm(r).includes(norm(k))),
+
+  // ---- length constraints ----
+  "length_constraints:number_words": (r, kw) =>
+    cmp(words(r).length, kw.relation as string, kw.num_words as number),
+  "length_constraints:number_sentences": (r, kw) =>
+    cmp(sentences(r).length, kw.relation as string, kw.num_sentences as number),
+  "length_constraints:number_paragraphs": (r, kw) =>
+    cmp(paragraphs(r).length, kw.relation as string, kw.num_paragraphs as number),
+
+  // ---- detectable format ----
+  "detectable_format:number_bullet_lists": (r, kw) => {
+    const bullets = r.split("\n").filter((l) => /^\s*[*\-]\s+/.test(l)).length;
+    return cmp(bullets, kw.relation as string, kw.num_bullets as number);
+  },
+  "detectable_format:number_highlighted_sections": (r, kw) => {
+    const hi = (r.match(/\*[^*\n]+\*/g) ?? []).length;
+    return cmp(hi, kw.relation as string, kw.num_highlights as number);
+  },
+  "detectable_format:title": (r) => /<<[^>\n]+>>/.test(r),
+  "detectable_format:json_format": (r) => {
+    const t = r.trim().replace(/^```(?:json)?\s*|\s*```$/g, "").trim();
+    try { JSON.parse(t); return true; } catch { return false; }
+  },
+
+  // ---- case ----
+  "change_case:english_lowercase": (r) => r === r.toLowerCase(),
+  "change_case:english_capital": (r) => r === r.toUpperCase(),
+
+  // ---- start / end ----
+  "startend:end_checker": (r, kw) =>
+    norm(r.trimEnd()).endsWith(norm((kw.end_phrase as string).trim())),
+  "startend:quotation": (r) => {
+    const t = r.trim();
+    return t.startsWith('"') && t.endsWith('"') && t.length >= 2;
+  },
+
+  // ---- punctuation ----
+  "punctuation:no_comma": (r) => !r.includes(","),
+};
+
+/** Score one IFEval instance (strict). Unknown instruction ids count as failed
+ *  (so coverage gaps never inflate the score) — filter the dataset to the
+ *  supported ids, or extend INSTRUCTIONS, to avoid that. */
+export function scoreInstance(instance: IFEvalInstance, response: string): InstanceResult {
+  const perInstruction = instance.instruction_id_list.map((id, i) => {
+    const check = INSTRUCTIONS[id];
+    if (!check) return false;
+    try { return check(response, instance.kwargs[i] ?? {}, instance); } catch { return false; }
+  });
+  return { perInstruction, followedAll: perInstruction.every(Boolean) };
+}
+
+export interface IFEvalReport {
+  n: number;
+  /** Fraction of instances where ALL instructions were followed. */
+  promptAccuracy: number;
+  /** Fraction of individual instructions followed. */
+  instructionAccuracy: number;
+}
+
+/** Aggregate the official prompt-level + instruction-level accuracies over a set
+ *  of (instance, response) pairs. */
+export function aggregate(pairs: Array<{ instance: IFEvalInstance; response: string }>): IFEvalReport {
+  let prompts = 0, instrFollowed = 0, instrTotal = 0;
+  for (const { instance, response } of pairs) {
+    const r = scoreInstance(instance, response);
+    if (r.followedAll) prompts++;
+    instrFollowed += r.perInstruction.filter(Boolean).length;
+    instrTotal += r.perInstruction.length;
+  }
+  return {
+    n: pairs.length,
+    promptAccuracy: pairs.length ? prompts / pairs.length : 0,
+    instructionAccuracy: instrTotal ? instrFollowed / instrTotal : 0,
+  };
+}
+
+/** Instruction ids this implementation supports (filter datasets to these). */
+export const SUPPORTED_INSTRUCTIONS = new Set(Object.keys(INSTRUCTIONS));

package/src/eval/ultrafeedback.ts

@@ -0,0 +1,60 @@
+// UltraFeedback (binarized) → mlx-bun preference rows. The binarized schema
+// (HuggingFaceH4/ultrafeedback_binarized) gives each row a `prompt` plus
+// `chosen` / `rejected` as message lists; ORPO wants {prompt, chosen, rejected}
+// where chosen/rejected are the assistant response *strings*. Pure transform so
+// it is unit-testable without downloading the dataset.
+
+export interface PreferenceRow {
+  prompt: string;
+  chosen: string;
+  rejected: string;
+}
+
+type Msg = { role?: string; content?: string };
+
+/** Last assistant message content in a binarized message list (the response). */
+function lastAssistant(messages: unknown): string | null {
+  if (!Array.isArray(messages)) return null;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const m = messages[i] as Msg;
+    if (m && m.role === "assistant" && typeof m.content === "string" && m.content.trim())
+      return m.content;
+  }
+  return null;
+}
+
+/** Convert one binarized UltraFeedback row → a preference row, or null if it is
+ *  unusable (missing fields, or chosen == rejected so there is no signal). */
+export function toPreferenceRow(row: Record<string, unknown>): PreferenceRow | null {
+  const prompt = typeof row.prompt === "string" ? row.prompt : null;
+  const chosen = lastAssistant(row.chosen);
+  const rejected = lastAssistant(row.rejected);
+  if (!prompt || !chosen || !rejected) return null;
+  if (chosen.trim() === rejected.trim()) return null;
+  return { prompt, chosen, rejected };
+}
+
+/** Rough token estimate (≈4 chars/token) for a cheap length pre-filter before
+ *  tokenization — keeps the curated set within a target context without loading
+ *  a tokenizer. */
+export function approxTokens(text: string): number {
+  return Math.ceil(text.length / 4);
+}
+
+/** Filter + convert a stream of binarized rows. Drops unusable rows and ones
+ *  whose prompt+chosen or prompt+rejected approx-exceeds `maxApproxTokens`. */
+export function curate(
+  rows: Record<string, unknown>[],
+  maxApproxTokens = 2048,
+): PreferenceRow[] {
+  const out: PreferenceRow[] = [];
+  for (const row of rows) {
+    const p = toPreferenceRow(row);
+    if (!p) continue;
+    const base = approxTokens(p.prompt);
+    if (base + approxTokens(p.chosen) > maxApproxTokens) continue;
+    if (base + approxTokens(p.rejected) > maxApproxTokens) continue;
+    out.push(p);
+  }
+  return out;
+}

package/src/lora.ts

@@ -37,7 +37,7 @@ export interface AdapterInfo {
  *  materialized mlx arrays. Tensor names come from our header parser;
  *  arrays from mlx's native loader via the proven map-get pattern
  *  (weights.ts). The map and header mmap are freed before returning. */
-function loadAdapterTensors(file: string): Map<string, MlxArray> {
+export function loadAdapterTensors(file: string): Map<string, MlxArray> {
   const sf = SafetensorsFile.open(file);
   const names = [...sf.tensors.keys()];
   sf.mmap.unmap();
@@ -74,18 +74,21 @@ function loadAdapterTensors(file: string): Map<string, MlxArray> {
 }
 
 /** Adapter scale from its config: mlx-lm writes lora_parameters.scale;
- *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py. */
-async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null }> {
+ *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py.
+ *  `rsLora` (recorded at train time) means the effective per-layer scale is
+ *  α/√rank — the caller divides by √(that layer's rank). */
+async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null; rsLora: boolean }> {
   for (const name of ["optiq_lora_config.json", "adapter_config.json"]) {
     const f = Bun.file(`${dir}/${name}`);
     if (await f.exists()) {
       const cfg = (await f.json()) as Record<string, any>;
       const lp = cfg.lora_parameters;
+      const rsLora = Boolean(cfg.rs_lora ?? lp?.rs_lora ?? false);
       if (lp && typeof lp === "object")
-        return { scale: Number(lp.scale ?? 20.0), rank: lp.rank ?? null };
+        return { scale: Number(lp.scale ?? 20.0), rank: lp.rank ?? null, rsLora };
       const alpha = Number(cfg.lora_alpha ?? 16);
       const r = Number(cfg.r ?? 8);
-      return { scale: r ? alpha / r : 1.0, rank: r || null };
+      return { scale: r ? alpha / r : 1.0, rank: r || null, rsLora };
     }
   }
   throw new Error(`no adapter_config.json in ${dir}`);
@@ -200,7 +203,7 @@ export class AdapterManager {
     if (!existsSync(path)) throw new Error(`adapter dir not found: ${path}`);
 
     const weightsFile = adapterWeightsFile(path);
-    const { scale, rank: configRank } = await readAdapterScale(path);
+    const { scale, rank: configRank, rsLora } = await readAdapterScale(path);
     const tensors = loadAdapterTensors(weightsFile);
     const targets = this.#model.loraTargets();
 
@@ -243,7 +246,9 @@ export class AdapterManager {
             `vs base [in ${linear.inFeatures}, out ${linear.outFeatures}]; ` +
             `was this adapter trained for a different base model?`,
           );
-        validated.push({ linear, lw: { a, b, scale, rank } });
+        // rsLoRA: effective per-layer scale is α/√rank (matches training).
+        const effScale = rsLora ? scale / Math.sqrt(rank) : scale;
+        validated.push({ linear, lw: { a, b, scale: effScale, rank } });
       }
       if (validated.length === 0)
         throw new Error(

package/src/mlx/metal-kernel.ts

@@ -29,8 +29,16 @@ export interface MetalKernelSpec {
 }
 
 export interface MetalKernelCall {
-  outputs: { shape: number[]; dtype: Dtype }[];
-  grid: [number, number, number];
+  /** Concrete output shapes/dtypes. Provide this OR `outputShapeFn` (the latter
+   *  makes the kernel derive its outputs from the input shapes, like a real MLX
+   *  primitive — required to compose inside an mx.compile'd / shapeless closure,
+   *  where the wrapper must answer "given these input shapes, what comes out?"). */
+  outputs?: { shape: number[]; dtype: Dtype }[];
+  /** Derive output shapes/dtypes from the (current/traced) input shapes. */
+  outputShapeFn?: (inputs: MlxArray[]) => { shape: number[]; dtype: Dtype }[];
+  /** Concrete launch grid, OR derive it from inputs (so it tracks shape under
+   *  compile replay the same way `outputShapeFn` tracks output shapes). */
+  grid: [number, number, number] | ((inputs: MlxArray[]) => [number, number, number]);
   threadGroup: [number, number, number];
   templateInts?: Record<string, number>;
   templateDtypes?: Record<string, Dtype>;
@@ -60,12 +68,18 @@ export class MetalKernel {
     if (this.#disposed) throw new Error("MetalKernel used after dispose");
     const cfg = C.mlx_fast_metal_kernel_config_new();
     try {
-      for (const o of call.outputs) {
+      // Derive outputs + grid from the (current) input shapes when a function is
+      // given — so the SAME kernel object answers correctly for any input shape
+      // during a compile/replay trace, instead of baking in one call's shapes.
+      const outputs = call.outputShapeFn ? call.outputShapeFn(inputs) : call.outputs;
+      if (!outputs) throw new Error("metal_kernel: provide outputs or outputShapeFn");
+      const grid = typeof call.grid === "function" ? call.grid(inputs) : call.grid;
+      for (const o of outputs) {
         const shape = new Int32Array(o.shape);
         if (C.mlx_fast_metal_kernel_config_add_output_arg(cfg, ptr(shape), BigInt(o.shape.length), o.dtype) !== 0)
           throw new Error(`metal_kernel add_output_arg failed: ${takeMlxError() ?? ""}`);
       }
-      C.mlx_fast_metal_kernel_config_set_grid(cfg, call.grid[0], call.grid[1], call.grid[2]);
+      C.mlx_fast_metal_kernel_config_set_grid(cfg, grid[0], grid[1], grid[2]);
       C.mlx_fast_metal_kernel_config_set_thread_group(cfg, call.threadGroup[0], call.threadGroup[1], call.threadGroup[2]);
       if (call.initValue !== undefined)
         C.mlx_fast_metal_kernel_config_set_init_value(cfg, call.initValue);

package/src/model/flash-attention.ts

@@ -406,7 +406,13 @@ for (uint kj = 0u; kj < n_kv_tiles; ++kj) {
     }
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
-    if (is_causal && kv_base_row > q_row_global) {
+    // Skip a KV tile only when it is entirely beyond this Q tile's last row
+    // (fully causally masked). The condition is UNIFORM across the threadgroup
+    // (q_tile_idx·BQ is the same for every thread), so the barrier below is hit
+    // by all threads — using the per-thread q_row_global here would diverge the
+    // barrier (undefined behavior). The per-element mask below still handles the
+    // partial diagonal tile.
+    if (is_causal && kv_base_row >= (q_tile_idx * BQ + BQ)) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         continue;
     }
@@ -533,16 +539,12 @@ export function flashBackward(
       grid: [nQTiles * bq, Hq, B],
       threadGroup: [bq, 1, 1],
     });
-    // The dKV kernel emits dK with its (Tkv, D) axes transposed in the buffer
-    // (verified: dK[i][j] == correct[j][i] at T=D, a [Tkv,D]↔[D,Tkv] swap in
-    // general — scripts/flash-dk-debug.ts JS-FIX 0.07% vs autograd). dV is
-    // unaffected. Recover the correct layout: transpose → materialize → reshape.
-    const dkFixed = ops.reshape(
-      ops.contiguous(ops.transposeAxes(dkv[0]!, [0, 1, 3, 2])),
-      [B, Hkv, Tkv, D],
-    );
-    dkv[0]!.dispose();
-    return [dq[0]!, dkFixed, dkv[1]!];
+    // dK is written by the kernel in the correct [B,Hkv,Tkv,D] layout
+    // (dK_j[d] → [kv_row][d]); no post-transpose. (An earlier "fix" transposed
+    // it on the false belief the buffer was swapped — that transpose is only the
+    // identity at Tkv==D and corrupted dK for Tkv≠D; removed after
+    // scripts/experiments/flash-dkv-debug.ts showed the raw output is correct.)
+    return [dq[0]!, dkv[0]!, dkv[1]!];
   } finally {
     Dvec.dispose();
     shape.dispose();
@@ -607,11 +609,25 @@ export function flashAttention(
 }
 
 // ---------------------------------------------------------------------------
-// Training-attention routing. ops.sdpa's vjp gives a wrong dK (it's an
-// inference op), so the training forward must use a known-correct attention:
-// the flash kernel for causal/full layers (memory-efficient, validated), and a
-// manual materialized SDPA for sliding-window (array-mask) layers. Set by the
-// trainer; null = inference (use ops.sdpa as before).
+// Training-attention routing (opt-in; default null = ops.sdpa).
+//
+// Parity oracle, made explicit:
+//  - DEFAULT (ops.sdpa) is the L1 path: mlx-lm's tuner differentiates the same
+//    mx.fast.scaled_dot_product_attention, so training through it is bit-faithful
+//    to mlx-lm. Exact dQ/dK/dV (finite-difference cross-checked), fast, but
+//    O(L²) backward memory (materializes the score matrix).
+//  - This flash kernel is the L2 path: a port of mlx-optiq's
+//    flash_attention_metal (O(L) backward memory). Its parity oracle is OPTIQ,
+//    not finite differences. Confirmed: optiq's flash dK == ops.sdpa to f16
+//    (scripts/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
+//    (tests/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
+//
+// Two real port bugs were fixed to reach that parity (the original diverged
+// ~100% from optiq on dK): a spurious dK transpose in flashBackward (corrupted
+// dK for Tkv≠D), and a per-thread (divergent) threadgroup_barrier in the dQ
+// causal tile-skip. flash is opt-in (MLX_BUN_TRAIN_ATTN=flash) for memory-bound
+// long context; it is ~30× slower than ops.sdpa. flashSupported gates
+// head_dim/seq; sliding-window (array-mask) layers fall back to manualSdpa.
 // ---------------------------------------------------------------------------
 
 let trainAttnMode: "flash" | "manual" | null = null;

package/src/model/fused-geglu-kernel.ts

@@ -16,6 +16,8 @@
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import { MetalKernel } from "../mlx/metal-kernel";
+import { CustomVjp } from "../mlx/custom-vjp";
+import * as ops from "../mlx/ops";
 
 /** ON by default: the match-compat kernel is BIT-EXACT with the spelled-out
  *  path (kl --decode = 0 on e4b), so per the bit-parity-floor principle it
@@ -28,6 +30,24 @@ export function fusedGeluEnabled(): boolean {
 /** Dispatch counter (a gate test asserts the kernel actually ran). */
 export let fusedGeluCalls = 0;
 
+// Training-mode flag: when set, model GeGLU sites use the DIFFERENTIABLE wrapper
+// (Metal-kernel forward + hand-derived vjp) instead of the plain `fusedGeglu`
+// (a bare CustomKernel with no vjp — feeding it to autograd throws
+// "[Primitive::vjp] Not implemented for CustomKernel"). The differentiable
+// wrapper also recomputes the gelu in the backward from the primal `a` instead
+// of retaining the ~9 forward intermediates, so it lowers the MLP activation
+// memory the backward holds — the dominant resident at long context under
+// ops.sdpa (which never materializes the O(T²) scores). Set by the trainer
+// around the train loop; cleared in its finally. Off → inference behaviour
+// (plain `fusedGeglu`) is byte-for-byte unchanged.
+let _trainingGeglu = false;
+export function setFusedGeluTraining(on: boolean): void {
+  _trainingGeglu = on;
+}
+export function fusedGeluTraining(): boolean {
+  return _trainingGeglu;
+}
+
 // Pointwise: one thread per element. gelu_approx =
 //   0.5 x (1 + tanh(√(2/π)·(x + 0.044715 x³)))  — composed in f32.
 const SOURCE = String.raw`
@@ -85,3 +105,94 @@ export function fusedGeglu(a: MlxArray, b: MlxArray): MlxArray {
   });
   return out!;
 }
+
+// ---------------------------------------------------------------------------
+// Differentiable wrapper: Metal kernel forward + pure-MLX backward
+// ---------------------------------------------------------------------------
+
+// GELU tanh-approx constants
+const SQRT_2_PI = 0.7978845608028654;
+const C_CUBE = 0.044715;
+
+/** Compute gelu_approx(a) = 0.5 * a * (1 + tanh(sqrt(2/pi) * (a + 0.044715*a^3)))
+ *  and gelu'(a) using the MLX ops layer. All intermediates are disposed
+ *  before returning; caller owns the two returned arrays. */
+function geluAndGrad(a: MlxArray): { gelu: MlxArray; geluGrad: MlxArray } {
+  // a^3 = a * a * a
+  const a2 = ops.mul(a, a);
+  const a3 = ops.mul(a2, a);
+  // inner = a + 0.044715 * a^3
+  const cx3 = ops.mulScalar(a3, C_CUBE);
+  const inner = ops.add(a, cx3);
+  // z = sqrt(2/pi) * inner
+  const z = ops.mulScalar(inner, SQRT_2_PI);
+  // t = tanh(z)
+  const t = ops.tanh(z);
+  // gelu = 0.5 * a * (1 + t)
+  const one = ops.scalarLike(1, a);
+  const t1 = ops.add(one, t);
+  const halfA = ops.mulScalar(a, 0.5);
+  const gelu = ops.mul(halfA, t1);
+
+  // gelu'(a) = 0.5*(1+t) + 0.5*a*(1-t^2)*dz/da
+  // dz/da = sqrt(2/pi) * (1 + 3*0.044715*a^2)
+  // 1 - t^2 = sech^2(z)
+  const halfT1 = ops.mulScalar(t1, 0.5);            // 0.5*(1+t)
+  const t2 = ops.mul(t, t);                         // t^2
+  const sech2 = ops.sub(one, t2);                   // 1 - t^2
+  const halfASech2 = ops.mul(halfA, sech2);         // 0.5*a*(1-t^2)
+  // dz/da
+  const a2scaled = ops.mulScalar(a2, 3 * C_CUBE);   // 3*0.044715*a^2
+  const dzda = ops.add(one, a2scaled);               // 1 + 3*0.044715*a^2
+  const dzdaScaled = ops.mulScalar(dzda, SQRT_2_PI);// sqrt(2/pi)*dz_unnorm
+  // chain: 0.5*a*(1-t^2)*dz/da
+  const chainTerm = ops.mul(halfASech2, dzdaScaled);
+  const geluGrad = ops.add(halfT1, chainTerm);       // 0.5*(1+t) + chain
+
+  // dispose all intermediates
+  for (const arr of [a2, a3, cx3, inner, z, t, one, t1, halfA, halfT1, t2, sech2,
+                      halfASech2, a2scaled, dzda, dzdaScaled, chainTerm]) {
+    arr.dispose();
+  }
+
+  return { gelu, geluGrad };
+}
+
+let _differentiableKernel: CustomVjp | null = null;
+
+function getDifferentiableKernel(): CustomVjp {
+  if (!_differentiableKernel) {
+    _differentiableKernel = new CustomVjp(
+      // forward: use the Metal kernel
+      ([a, b]: MlxArray[]) => [fusedGeglu(a!, b!)],
+      // backward
+      (primals: MlxArray[], cotangents: MlxArray[], _outputs: MlxArray[]) => {
+        const [a, b] = primals as [MlxArray, MlxArray];
+        const [dc] = cotangents as [MlxArray];
+
+        const { gelu, geluGrad } = geluAndGrad(a);
+
+        // grad_b = dc * gelu(a)
+        const grad_b = ops.mul(dc, gelu);
+
+        // grad_a = dc * b * gelu'(a)
+        const dcb = ops.mul(dc, b);
+        const grad_a = ops.mul(dcb, geluGrad);
+
+        gelu.dispose();
+        geluGrad.dispose();
+        dcb.dispose();
+
+        return [grad_a, grad_b];
+      },
+    );
+  }
+  return _differentiableKernel;
+}
+
+/** gelu_approx(a) · b with a hand-derived vjp. Uses the Metal kernel forward,
+ *  MLX-ops backward. For training only (inference uses the non-differentiable
+ *  fusedGeglu). */
+export function fusedGegluDifferentiable(a: MlxArray, b: MlxArray): MlxArray {
+  return getDifferentiableKernel().apply([a, b])[0]!;
+}

package/src/model/gemma4-base.ts

@@ -46,6 +46,31 @@ export interface LoraWeights {
  *  (PLAN Phase 8 decision). Set/restored by generate(). */
 export class LoraState {
   active: string[] = [];
+  /** Training-only LoRA-input dropout. `rate` is the drop probability; `seed` is
+   *  set per micro-step by the trainer and is CONSTANT across that step's forward
+   *  and any recompute (segmented / gradient-checkpoint), so each layer's mask —
+   *  keyed by (seed, the linear's dropoutId) — is reproduced exactly in the
+   *  backward. `seed === null` (the default, and for inference) disables it. */
+  dropoutRate = 0;
+  dropoutSeed: number | null = null;
+}
+
+/** Inverted LoRA-input dropout, keyed by (seed, id) so it is deterministic —
+ *  the same (seed, id, shape) reproduces the mask, which is what makes the
+ *  segmented / gradient-checkpoint recompute correct. kept ⇒ x/(1-p),
+ *  dropped ⇒ 0 (preserves the expectation). Caller owns the result. */
+export function loraInputDropout(x: MlxArray, p: number, seed: number, id: number): MlxArray {
+  const key = ops.randomKey(BigInt(seed) * 100003n + BigInt(id));
+  const u = ops.randomUniform(x.shape, Dtype.float32, 0, 1, key); // [shape] in [0,1)
+  key.dispose();
+  const pArr = ops.scalarLike(p, u);
+  const keep = ops.less(pArr, u); // u > p ⇒ keep (bool [shape])
+  const scaleX = ops.scalarLike(1 / (1 - p), x);
+  const zeroX = ops.scalarLike(0, x);
+  const mask = ops.where(keep, scaleX, zeroX); // x.dtype
+  const xd = ops.mul(x, mask);
+  for (const a of [u, pArr, keep, scaleX, zeroX, mask]) a.dispose();
+  return xd;
 }
 
 export class QuantizedLinear {
@@ -53,6 +78,9 @@ export class QuantizedLinear {
   adapters: Map<string, LoraWeights> | null = null;
   /** Shared per-model active state (wired by AdapterManager.mount). */
   loraState: LoraState | null = null;
+  /** Stable per-target index for keying training-only LoRA dropout (set by
+   *  attachForTraining). Gives each adapted linear an independent dropout mask. */
+  dropoutId = 0;
 
   constructor(
     readonly w: MlxArray,
@@ -92,10 +120,19 @@ export class QuantizedLinear {
     // cast form is what the adapters were trained behind.)
     const st = this.loraState;
     if (st && st.active.length > 0 && this.adapters && this.adapters.size > 0) {
+      // Training-only LoRA-input dropout (PEFT applies dropout to x before A;
+      // the base quantized path is untouched). Keyed by (seed, dropoutId) so the
+      // backward recompute reproduces the exact mask.
+      let xLora = x;
+      let xDrop: MlxArray | null = null;
+      if (st.dropoutSeed !== null && st.dropoutRate > 0) {
+        xDrop = loraInputDropout(x, st.dropoutRate, st.dropoutSeed, this.dropoutId);
+        xLora = xDrop;
+      }
       for (const id of st.active) {
         const lw = this.adapters.get(id);
         if (!lw) continue;
-        const xa = ops.matmul(x, lw.a);
+        const xa = ops.matmul(xLora, lw.a);
         const z = ops.matmul(xa, lw.b);
         xa.dispose();
         const zs = ops.mulScalar(z, lw.scale);
@@ -105,6 +142,7 @@ export class QuantizedLinear {
         out = disposing(out, ops.add(out, zc));
         zc.dispose();
       }
+      xDrop?.dispose();
     }
     return out;
   }