mlx-bun 0.0.9 → 0.0.11

package/src/train/trainer.ts

@@ -20,7 +20,6 @@ import type { RuntimeModel } from "../model/factory";
 import { Gemma4Model, type GradCheckpointCtx } from "../model/gemma4";
 import { MiniCPM5Model } from "../model/minicpm5";
 import { setTrainingAttn } from "../model/flash-attention";
-import { setFusedGeluTraining } from "../model/fused-geglu-kernel";
 import { SegmentedBackward, SegmentedBackwardGemma4, SegmentedBackwardOrpo, SegmentedBackwardOrpoGemma4, SegmentedBackwardOrpoPrefix, SegmentedBackwardOrpoPrefixGemma4, planSegmentsBySize } from "./segmented";
 import type { LoadedTokenizer } from "../tokenizer";
 import type { ChatTemplate } from "../chat-template";
@@ -260,7 +259,6 @@ export async function trainLora(
   // autograd flows and the backward recomputes the gelu from the primal instead
   // of retaining the spelled-out intermediates. No-op for non-Gemma models.
   // Cleared in the finally below.
-  setFusedGeluTraining(true);
 
   const saveCfg: SaveAdapterConfig = {
     rank: cfg.rank,
@@ -383,11 +381,20 @@ export async function trainLora(
     });
     return { adapterPath: cfg.adapterPath, appliedRanks, numIters: result.numIters };
   } finally {
-    setFusedGeluTraining(false);
     disposeLora(lora);
   }
 }
 
+// ---------------------------------------------------------------------------
+// Segment planning: uniform planSegmentsBySize only. A full-attention
+// ISOLATION planner (kernel-review backlog #2 / segmented-backward-training.md
+// §5) was built and A/B-measured 2026-07-02: ZERO peak win — mlx's sdpa
+// backward is O(L²) for EVERY layer (~3.5 GB/layer @8K e4b), sliding included,
+// so the worst segment is set by layer count alone and segment_size is the
+// whole knob (seg1 = 14.59 GB @8K, +3% step time — fits the 24 GB M4 Pro).
+// Evidence: scripts/experiments/seg-isolation-smoke.ts; note in segmented.ts.
+// ---------------------------------------------------------------------------
+
 // ---------------------------------------------------------------------------
 // SFT loop
 // ---------------------------------------------------------------------------
@@ -421,14 +428,14 @@ async function sftLoop(
   const useSegmented = cfg.segmentSize > 0;
   let segmented: SegmentedBackward | SegmentedBackwardGemma4 | null = null;
   if (useSegmented) {
-    if (model instanceof MiniCPM5Model)
-      segmented = new SegmentedBackward(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
-    else if (model instanceof Gemma4Model)
-      segmented = new SegmentedBackwardGemma4(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
-    else
+    if (!(model instanceof MiniCPM5Model) && !(model instanceof Gemma4Model))
       throw new Error("segmented backward (segmentSize > 0) is only wired for MiniCPM5 and Gemma4");
+    const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
+    segmented = model instanceof MiniCPM5Model
+      ? new SegmentedBackward(model, lora, ranges)
+      : new SegmentedBackwardGemma4(model, lora, ranges);
     emit({ type: "stage", stage: "setup", progress: 0.04,
-      message: `segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers` });
+      message: `segmented backward: ${ranges.length} segments of <=${cfg.segmentSize} layers` });
   }
 
   // The value_and_grad loss closure: temporarily swap the differentiated
@@ -479,7 +486,7 @@ async function sftLoop(
   // scripts/experiments/flash-fd-check.ts), though ~30× slower than ops.sdpa.
   // Cleared in finally.
   //
-  // GEMMA GUARD (same convention as the PERF_KERNEL/FUSED_GELU sanitization in
+  // GEMMA GUARD (same convention as the launcher's env 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)
@@ -832,6 +839,8 @@ async function orpoLoop(
   let segmentedOrpo: SegmentedBackwardOrpo | SegmentedBackwardOrpoGemma4 | null = null;
   let segmentedOrpoPrefix: SegmentedBackwardOrpoPrefix | SegmentedBackwardOrpoPrefixGemma4 | null = null;
   if (useSegmented) {
+    if (!(model instanceof MiniCPM5Model) && !(model instanceof Gemma4Model))
+      throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
     const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
     if (useSegmentedPrefix) {
       // The composed prefix-shared segmented backward (MiniCPM5 or Gemma4/e4b), plus a
@@ -845,15 +854,13 @@ async function orpoLoop(
         segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
       }
       emit({ type: "stage", stage: "setup", progress: 0.04,
-        message: `orpo segmented + prefix-share (${model instanceof Gemma4Model ? "e4b" : "MiniCPM5"}): single concat forward streamed over ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers (two-forward fallback on prompt mismatch)` });
+        message: `orpo segmented + prefix-share (${model instanceof Gemma4Model ? "e4b" : "MiniCPM5"}): single concat forward streamed over ${ranges.length} segments of <=${cfg.segmentSize} layers (two-forward fallback on prompt mismatch)` });
     } else if (model instanceof MiniCPM5Model)
       segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
-    else if (model instanceof Gemma4Model)
-      segmentedOrpo = new SegmentedBackwardOrpoGemma4(model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
     else
-      throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
+      segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
     if (!useSegmentedPrefix) emit({ type: "stage", stage: "setup", progress: 0.04,
-      message: `orpo segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers${segFusedChunk > 0 ? ` + fused linear-CE head (${segFusedChunk}/chunk)` : ""}` });
+      message: `orpo segmented backward: ${ranges.length} segments of <=${cfg.segmentSize} layers${segFusedChunk > 0 ? ` + fused linear-CE head (${segFusedChunk}/chunk)` : ""}` });
   }
 
   // Token-chunked head (non-segmented path only): the per-chunk checkpoints are

package/src/vision/prompt.ts

@@ -1,10 +1,21 @@
-// Multimodal prompt assembly: OpenAI image_url content parts → spliced
-// token ids + merged input embeddings + image-token mask.
+// Multimodal prompt assembly: OpenAI image/audio content parts → spliced
+// token ids + merged input embeddings + attention/zeroing masks.
 //
-// The chat template renders one <|image|> token (258880) per image part;
-// we expand each into <boi> + image_token × softTokens + <eoi> (matching
-// optiq's frontend), embed the full sequence, then overwrite the
-// image-token rows with the vision tower's features.
+// The chat template renders one <|image|> token (258880) per image part and
+// one <|audio|> token (258881) per audio part; we expand each in document
+// order — <boi> + image_token × softTokens + <eoi> (matching optiq's vision
+// frontend) and <boa> + audio_token × n + <eoa> with
+// n = min(ceil(duration_ms/40), 750) (processing_gemma4.py) — embed the full
+// sequence, then overwrite the soft-token rows with the towers' features.
+//
+// Attention semantics (audio-input-plan.md §3.3, resolved vs the oracle):
+//   - image runs attend bidirectionally among themselves (bidirMask), BUT
+//     any audio token in the prompt disables the vision overlay entirely —
+//     mixed image+audio prompts run fully causal (bidirMask = null);
+//   - audio runs are strictly causal, always;
+//   - per-layer-input id zeroing (e2b/e4b) applies to the UNION of
+//     multimodal soft tokens (image | audio) — multimodalMask, decoupled
+//     from the bidirectional mask.
 
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
@@ -12,6 +23,10 @@ import * as ops from "../mlx/ops";
 import type { Gemma4Model } from "../model/gemma4";
 import type { ChatTemplate, ChatMessage, ToolDefinition } from "../chat-template";
 import type { LoadedTokenizer } from "../tokenizer";
+import type { AudioTower } from "../audio/conformer";
+import { audioSoftTokenCount, decodeAudio } from "../audio/decode";
+import { extractMelFeatures } from "../audio/features";
+import { fetchMediaBytes } from "../media-fetch";
 
 /** Common contract for both vision towers (encoder-free gemma4_unified in
  *  ./embedder.ts and the SigLIP encoder in ./siglip.ts): preprocess image
@@ -33,12 +48,38 @@ export interface VisionPrompt {
   imageMask: MlxArray;
 }
 
+export interface MultimodalPrompt {
+  ids: number[];
+  /** Unscaled merged embeddings [1, L, hidden] — caller disposes. */
+  embeddings: MlxArray;
+  /** bool [L] image-token mask for the bidirectional attention overlay.
+   *  null when the prompt contains ANY audio (§3.3 Q1: mixed image+audio
+   *  runs fully causal) or no images. Caller disposes when non-null. */
+  bidirMask: MlxArray | null;
+  /** bool [L] union multimodal soft-token mask (image | audio) for
+   *  per-layer-input id zeroing (§3.3 Q2). Caller disposes. */
+  multimodalMask: MlxArray;
+}
+
 export interface VisionTokenIds {
   imageTokenId: number;
   boiTokenId: number;
   eoiTokenId: number;
 }
 
+export interface AudioTokenIds {
+  audioTokenId: number;
+  boaTokenId: number;
+  eoaTokenId: number;
+}
+
+/** The towers a multimodal prompt may need; pass only what the request uses
+ *  (vision-only and audio-only fall out as special cases). */
+export interface MultimodalTowers<P extends { softTokens: number }> {
+  vision?: { tower: VisionEncoder<P>; tokenIds: VisionTokenIds };
+  audio?: { tower: AudioTower; tokenIds: AudioTokenIds };
+}
+
 /** Extract image bytes from OpenAI-style content parts, rewriting the
  *  parts to the template's {type:"image"} form. */
 export async function extractImages(
@@ -57,7 +98,7 @@ export async function extractImages(
         const url = (part.image_url as { url: string } | undefined)?.url
           ?? (part.image_url as unknown as string);
         if (typeof url !== "string") throw new Error("image_url part missing url");
-        images.push(await fetchImageBytes(url));
+        images.push(await fetchMediaBytes(url, "image"));
         parts.push({ type: "image" });
       } else if (part.type === "image") {
         // already template-form: expect base64 `data` or pass-through marker
@@ -76,83 +117,211 @@ export async function extractImages(
   return { messages: out, images };
 }
 
-async function fetchImageBytes(url: string): Promise<Uint8Array> {
-  if (url.startsWith("data:")) {
-    const comma = url.indexOf(",");
-    if (comma === -1) throw new Error("malformed data: URL");
-    const meta = url.slice(0, comma);
-    const body = url.slice(comma + 1);
-    if (!meta.includes("base64")) throw new Error("data: URL must be base64");
-    return Uint8Array.from(Buffer.from(body, "base64"));
-  }
-  if (url.startsWith("http://") || url.startsWith("https://")) {
-    const res = await fetch(url);
-    if (!res.ok) throw new Error(`image fetch failed: ${res.status} ${url}`);
-    return new Uint8Array(await res.arrayBuffer());
+/** Extract audio bytes from OpenAI-style content parts, rewriting the parts
+ *  to the template's {type:"audio"} form. Accepts the OpenAI-canonical
+ *  {type:"input_audio", input_audio:{data:<b64>, format:...}} plus optiq's
+ *  {type:"audio"} (base64 `data`) and {type:"audio_url"} aliases. */
+export async function extractAudio(
+  messages: ChatMessage[],
+): Promise<{ messages: ChatMessage[]; audio: Uint8Array[] }> {
+  const audio: Uint8Array[] = [];
+  const out: ChatMessage[] = [];
+  for (const m of messages) {
+    if (!Array.isArray(m.content)) {
+      out.push(m);
+      continue;
+    }
+    const parts: Array<Record<string, unknown>> = [];
+    for (const part of m.content) {
+      if (part.type === "input_audio") {
+        const data = (part.input_audio as { data?: string } | undefined)?.data;
+        if (typeof data !== "string")
+          throw new Error("input_audio part requires base64 `input_audio.data`");
+        audio.push(Uint8Array.from(Buffer.from(data, "base64")));
+        parts.push({ type: "audio" });
+      } else if (part.type === "audio") {
+        if (typeof part.data === "string") {
+          audio.push(Uint8Array.from(Buffer.from(part.data, "base64")));
+          parts.push({ type: "audio" });
+        } else {
+          throw new Error("audio part requires base64 `data`");
+        }
+      } else if (part.type === "audio_url") {
+        const url = (part.audio_url as { url: string } | undefined)?.url
+          ?? (part.audio_url as unknown as string);
+        if (typeof url !== "string") throw new Error("audio_url part missing url");
+        audio.push(await fetchMediaBytes(url, "audio"));
+        parts.push({ type: "audio" });
+      } else {
+        parts.push(part);
+      }
+    }
+    out.push({ ...m, content: parts });
   }
-  throw new Error(`unsupported image url scheme: ${url.slice(0, 16)}`);
+  return { messages: out, audio };
 }
 
-export async function buildVisionPrompt<P extends { softTokens: number }>(
+/** Decoded + featurized audio clip, ready to splice. */
+interface PreparedAudio {
+  mel: Float32Array;
+  frames: number;
+  softTokens: number;
+}
+
+export async function buildMultimodalPrompt<P extends { softTokens: number }>(
   model: Gemma4Model,
-  tower: VisionEncoder<P>,
+  towers: MultimodalTowers<P>,
   tokenizer: LoadedTokenizer,
   template: ChatTemplate,
   messages: ChatMessage[],
   images: Uint8Array[],
-  tokenIds: VisionTokenIds,
+  audio: Uint8Array[],
   tools: ToolDefinition[] | null = null,
-): Promise<VisionPrompt> {
+): Promise<MultimodalPrompt> {
+  if (images.length > 0 && !towers.vision)
+    throw new Error("images present but no vision tower configured");
+  if (audio.length > 0 && !towers.audio)
+    throw new Error("audio present but no audio tower configured");
+
   const rendered = template.render(messages, { tools });
   let ids = tokenizer.encode(rendered);
   if (ids[0] === ids[1] && ids[0] === tokenizer.bosTokenId) ids = ids.slice(1);
 
-  // preprocess + embed every image (in order of appearance), via the
-  // tower's own preprocessing (encoder-free vs SigLIP differ here).
-  const pre = await Promise.all(images.map((bytes) => tower.preprocess(bytes)));
+  // preprocess every clip/image up front (splice lengths depend on them);
+  // audio soft-token counts come from the DECODED samples, never metadata.
+  const pre = await Promise.all(
+    images.map((bytes) => towers.vision!.tower.preprocess(bytes)),
+  );
+  const clips: PreparedAudio[] = audio.map((bytes) => {
+    const samples = decodeAudio(bytes);
+    const softTokens = audioSoftTokenCount(samples.length);
+    const m = extractMelFeatures(samples);
+    return { mel: m.features, frames: m.frames, softTokens };
+  });
+
+  const vIds = towers.vision?.tokenIds;
+  const aIds = towers.audio?.tokenIds;
 
-  // expand each <|image|> into <boi> + image_token×soft + <eoi>
+  // expand each <|image|> into <boi> + image_token×soft + <eoi> and each
+  // <|audio|> into <boa> + audio_token×n + <eoa>, in document order
   const spliced: number[] = [];
-  const runs: { start: number; length: number }[] = [];
+  const runs: { start: number; length: number; kind: "image" | "audio"; index: number }[] = [];
   let imgIdx = 0;
+  let audIdx = 0;
   for (const id of ids) {
-    if (id === tokenIds.imageTokenId) {
+    if (vIds && id === vIds.imageTokenId) {
       const p = pre[imgIdx];
       if (!p) throw new Error("more <|image|> markers than images");
-      spliced.push(tokenIds.boiTokenId);
-      runs.push({ start: spliced.length, length: p.softTokens });
-      for (let k = 0; k < p.softTokens; k++) spliced.push(tokenIds.imageTokenId);
-      spliced.push(tokenIds.eoiTokenId);
+      spliced.push(vIds.boiTokenId);
+      runs.push({ start: spliced.length, length: p.softTokens, kind: "image", index: imgIdx });
+      for (let k = 0; k < p.softTokens; k++) spliced.push(vIds.imageTokenId);
+      spliced.push(vIds.eoiTokenId);
       imgIdx++;
+    } else if (aIds && id === aIds.audioTokenId) {
+      const clip = clips[audIdx];
+      if (!clip) throw new Error("more <|audio|> markers than audio clips");
+      spliced.push(aIds.boaTokenId);
+      runs.push({ start: spliced.length, length: clip.softTokens, kind: "audio", index: audIdx });
+      for (let k = 0; k < clip.softTokens; k++) spliced.push(aIds.audioTokenId);
+      spliced.push(aIds.eoaTokenId);
+      audIdx++;
     } else {
       spliced.push(id);
     }
   }
   if (imgIdx !== images.length)
     throw new Error(`${images.length} images but ${imgIdx} <|image|> markers`);
+  if (audIdx !== audio.length)
+    throw new Error(`${audio.length} audio clips but ${audIdx} <|audio|> markers`);
 
-  // embed text tokens, then overwrite image-token rows with vision features
+  // embed text tokens, then overwrite soft-token rows with tower features
   const idsArr = ops.fromInt32(spliced, [1, spliced.length]);
   let embeds = model.embed.encode(idsArr);
   idsArr.dispose();
   const hidden = embeds.shape[2]!;
-  for (let i = 0; i < runs.length; i++) {
-    const { start, length } = runs[i]!;
-    const feats = tower.features(pre[i]!); // [1, soft, hidden]
-    const cast = feats.astype(embeds.dtype);
-    feats.dispose();
+  for (const run of runs) {
+    let cast: MlxArray;
+    if (run.kind === "image") {
+      // vision towers return pre-divided by embed_scale (their oracle's
+      // frontend convention) — cast to the embedding dtype and splice.
+      const feats = towers.vision!.tower.features(pre[run.index]!); // [1, soft, hidden]
+      cast = feats.astype(embeds.dtype);
+      feats.dispose();
+    } else {
+      // audio mirrors gen-e4b-audio-golden.py EXACTLY:
+      //   features.astype(embeds.dtype) / embed_scale
+      // — raw f32 embed_audio output, cast to bf16 FIRST, then divided by a
+      // weak (dtype-following, i.e. bf16) embed_scale scalar. The order is
+      // load-bearing for the bit-exact greedy gate (see AudioTower.features).
+      const clip = clips[run.index]!;
+      const raw = towers.audio!.tower.features(clip.mel, clip.frames, false); // [1, n, hidden] f32
+      if (raw.shape[1] !== clip.softTokens) {
+        const n = raw.shape[1];
+        raw.dispose();
+        embeds.dispose();
+        throw new Error(`audio tower produced ${n} frames but splice expects ${clip.softTokens}`);
+      }
+      const bf = raw.astype(embeds.dtype);
+      raw.dispose();
+      const scale = ops.scalarLike(towers.audio!.tower.embedScale, bf);
+      cast = ops.div(bf, scale);
+      bf.dispose();
+      scale.dispose();
+    }
+    const { start, length } = run;
     const updated = ops.sliceUpdate(embeds, cast, [0, start, 0], [1, start + length, hidden]);
     cast.dispose();
     embeds.dispose();
     embeds = updated;
   }
 
-  const maskInts = new Int32Array(spliced.length);
-  for (let i = 0; i < spliced.length; i++)
-    maskInts[i] = spliced[i] === tokenIds.imageTokenId ? 1 : 0;
-  const maskI32 = MlxArray.fromInt32(maskInts, [spliced.length]);
-  const imageMask = maskI32.astype(Dtype.bool);
-  maskI32.dispose();
+  // masks over the spliced sequence: union multimodal soft tokens (per-layer
+  // id zeroing) always; image-only bidirectional overlay only when NO audio
+  // is present (§3.3 Q1).
+  const mmInts = new Int32Array(spliced.length);
+  for (let i = 0; i < spliced.length; i++) {
+    const id = spliced[i]!;
+    mmInts[i] = (vIds && id === vIds.imageTokenId) || (aIds && id === aIds.audioTokenId) ? 1 : 0;
+  }
+  const mmI32 = MlxArray.fromInt32(mmInts, [spliced.length]);
+  const multimodalMask = mmI32.astype(Dtype.bool);
+  mmI32.dispose();
 
-  return { ids: spliced, embeddings: embeds, imageMask };
+  let bidirMask: MlxArray | null = null;
+  if (images.length > 0 && audio.length === 0) {
+    const imgInts = new Int32Array(spliced.length);
+    for (let i = 0; i < spliced.length; i++)
+      imgInts[i] = spliced[i] === vIds!.imageTokenId ? 1 : 0;
+    const imgI32 = MlxArray.fromInt32(imgInts, [spliced.length]);
+    bidirMask = imgI32.astype(Dtype.bool);
+    imgI32.dispose();
+  }
+
+  return { ids: spliced, embeddings: embeds, bidirMask, multimodalMask };
+}
+
+/** Vision-only prompt (back-compat wrapper over buildMultimodalPrompt —
+ *  behavior byte-identical to the pre-audio builder: the returned imageMask
+ *  serves as BOTH the bidirectional-attention mask and the per-layer-input
+ *  zeroing mask, which forwardEmbeddings derives from `bidir` when no
+ *  separate multimodal mask is passed). */
+export async function buildVisionPrompt<P extends { softTokens: number }>(
+  model: Gemma4Model,
+  tower: VisionEncoder<P>,
+  tokenizer: LoadedTokenizer,
+  template: ChatTemplate,
+  messages: ChatMessage[],
+  images: Uint8Array[],
+  tokenIds: VisionTokenIds,
+  tools: ToolDefinition[] | null = null,
+): Promise<VisionPrompt> {
+  const mp = await buildMultimodalPrompt(
+    model, { vision: { tower, tokenIds } }, tokenizer, template,
+    messages, images, [], tools,
+  );
+  // no audio was passed, so the union mask IS the image mask — reuse it for
+  // the zero-image edge (the original builder returned an all-false mask).
+  const imageMask = mp.bidirMask ?? mp.multimodalMask;
+  if (mp.bidirMask) mp.multimodalMask.dispose();
+  return { ids: mp.ids, embeddings: mp.embeddings, imageMask };
 }