mlx-bun 0.0.10 → 0.0.12

package/docs/reference/training.md

@@ -22,7 +22,7 @@ Training is reachable several ways — all of them drive the **same** runner
 
 | Path | How | Use when |
 |---|---|---|
-| **CLI** | `mlx-bun train <model> --data <dir> [options]` — foreground, full ORPO stack on by default; auto-detects e4b/Gemma + sets its env; `--save-every`, `--resume`, `--dry-run`, `--sft-scope full\|response`, `--grad-accum`, `--grad-clip`, `--seed`, `--val-size`; streams loss to the terminal. `mlx-bun help train` for all flags. Run long jobs detached: `nohup mlx-bun train … &` | **The default for CLI users** |
+| **CLI** | `mlx-bun train <model> --data <dir> [options]` — foreground, full ORPO stack on by default; auto-detects e4b/Gemma to pick sensible `--seq`/adapter-path defaults; no env flags needed; `--save-every`, `--resume`, `--dry-run`, `--sft-scope full\|response`, `--grad-accum`, `--grad-clip`, `--seed`, `--val-size`; streams loss to the terminal. `mlx-bun help train` for all flags. Run long jobs detached: `nohup mlx-bun train … &` | **The default for CLI users** |
 | **Web UI** | `mlx-bun serve`, open `/finetune` — pick model → dataset → hyperparameters → train; watch live train/val loss; merge/export the adapter | Interactive, in the browser |
 | **HTTP API** | `POST /api/finetune/submit` (job id + SSE events); `POST /api/finetune/inspect-dataset` to probe a file; `POST /api/finetune/merge` to fold an adapter into base weights | Scripted / remote |
 | **ORPO launcher** | [`scripts/train-orpo.ts`](../../scripts/train-orpo.ts) — the same stack driven by env vars (`MODEL=`, `DATA=`, `RANK=`, `RESUME=`, …). See [orpo-quickstart](./orpo-quickstart.md) | Env-var scripting |
@@ -167,7 +167,7 @@ are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
 | `target_modules` | string[] | `q,k,v,o,gate,up,down _proj` | Which linears get adapters |
 | `num_layers` | int | `-1` | `-1` = all layers; `N` = last N only |
 | `iters` | int >0 | `100` | Total training steps |
-| `learning_rate` | float >0 | `2e-4` (sft) / `5e-5` (dpo) | AdamW LR |
+| `learning_rate` | float >0 | `2e-4` (sft) / `1e-5` (orpo) / `5e-5` (dpo) | AdamW LR |
 | `max_seq_length` | int >0 | `512` | Truncate/pad sequences to this |
 | `batch_size` | int ≥1 | `1` | Rows per step (B=1 is the safe path; B>1 length-sorts + pads to 32) |
 | `grad_accumulation_steps` | int ≥1 | `1` | Accumulate grads over N micro-steps |
@@ -177,10 +177,14 @@ are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
 | `weight_decay` | float ≥0 | `0.01` | AdamW weight decay (β = `[0.9, 0.999]`, fixed) |
 | `lora_dropout` | float [0,1) | `0.0` | LoRA-input dropout (PEFT-style), training-only regularizer for small sets. Recompute-safe (mask keyed by step+layer, so segmented/checkpointed backward reproduces it) |
 | `rs_lora` | bool | `false` | rsLoRA — scale the update by α/√rank instead of α, so `rank_scaling` changes capacity not step size. Recorded in the adapter config; inference applies the same per-layer scale |
-| `lora_plus_ratio` | float ≥1 | `1.0` | LoRA+ — LR multiplier for the B leaves (A stays at base LR). >1 speeds the B-driven early learning (B is zero-init). 1 = off. Wired in the ORPO loop |
+| `lora_plus_ratio` | float ≥1 | `1.0` | LoRA+ — LR multiplier for the B leaves (A stays at base LR). >1 speeds the B-driven early learning (B is zero-init). 1 = off. The same optimizer policy applies to SFT, DPO, and ORPO |
 | `grad_checkpoint` | bool | `false` | Recompute layer activations in backward (memory↔compute; bit-identical) |
+| `mlp_split` | bool | `false` | Gradient-checkpoint MLP splitting (Gemma4-only, requires `grad_checkpoint`) |
 | `segment_size` | int | `0` (off) | `>0` enables segmented backward — layers per segment (see below) |
 | `save_checkpoints` | bool | `false` | Save every eval-step checkpoint + write `metrics.json` |
+| `grad_clip_norm` | float ≥0 | `1.0` | Gradient-norm clipping (on by default); `0` = off |
+| `val_max_examples` | int >0 | `256` | Caps the validation-set size used per eval |
+| `warm_start_adapter` | string | `""` (off) | Adapter dir to warm-start/resume LoRA weights from (the API field behind the CLI's `--resume`) |
 | `dpo_beta` | float >0 | `0.1` | DPO strength (dpo only) |
 | `dpo_warmup_iters` | int ≥0 | `0` | DPO LR warmup (dpo only) |
 | `dpo_lr_schedule` | `constant` \| `cosine` | `cosine` | DPO LR schedule (dpo only) |
@@ -197,17 +201,14 @@ are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
 
 | Env var | Set for training | Default | Why |
 |---|---|---|---|
-| `MLX_BUN_FUSED_GELU` | **`0` (required for Gemma)** | on | The fused GeGLU is a **CustomKernel with no gradient (vjp)** ([`fused-geglu-kernel.ts`](../../src/model/fused-geglu-kernel.ts)); the Gemma forward uses it ([`gemma4.ts:277`](../../src/model/gemma4.ts)), so a Gemma (e4b/12B/26B) backward fails unless it's off. MiniCPM5 (Llama-arch SwiGLU) never hits it, so the `.ts` script's MiniCPM5 default doesn't need it. |
-| `MLX_BUN_PERF_KERNEL` | **`0` for training** | on | The fused quantized-decode kernel likewise has no vjp. It only fires at decode L=1 (rare in the L>1 training forward), but the e4b recipe sets `0` to be safe. |
 | `MLX_BUN_TRAIN_ATTN` | **leave unset** | unset → `ops.sdpa` | Default `ops.sdpa` **is** mlx's fused flash-attention kernel — the correct, working path. `flash` selects a *different* hand-rolled O(L)-memory kernel: its two port bugs (dK buffer-transpose, dQ divergent causal barrier) are fixed and FD-validated (`tests/flash-attention.test.ts`, T≤256), but it is ~30× slower and the historical e4b ≥2K crash has **not been re-validated at that scale** since the fix — leave unset unless memory-bound and re-verified. |
 | `MLX_BUN_MEM_LOG` | `1` to profile | off | Print per-step peak/active/cache memory |
 
-> Important: the trainer itself ([`trainer.ts`](../../src/train/trainer.ts))
-> only reads `MLX_BUN_TRAIN_ATTN` and `MLX_BUN_MEM_LOG` — but the **model
-> forward** it runs reads `MLX_BUN_FUSED_GELU` / `MLX_BUN_PERF_KERNEL`. The
-> trainer does **not** disable those itself, so **the caller must** export
-> them. The [e4b recipe](#what-we-actually-run-the-e4b-recipe) does this; if
-> you train a Gemma model by hand, set `MLX_BUN_FUSED_GELU=0` yourself.
+> Training needs **no flag sanitization**: the old requirement to export
+> `MLX_BUN_FUSED_GELU=0` / `MLX_BUN_PERF_KERNEL=0` (the no-vjp custom
+> kernels) is gone — those kernels and their env vars were deleted
+> 2026-07-05 ([unified-engine-frontier-plan.md](../design/unified-engine-frontier-plan.md)).
+> Exporting them in an old recipe is harmless but does nothing.
 
 ## What we actually run (the e4b recipe)
 
@@ -215,7 +216,7 @@ Everything above is the full surface (*what you can do*). In practice the
 fine-tune we run is [`scripts/ft-e4b-v2.sh`](../../scripts/ft-e4b-v2.sh):
 e4b (gemma-4-e4b-it-OptiQ-4bit, pinned snapshot) on the lucien `chunk-v2-500`
 curated set (450 train convs) through the segmented-backward trainer. It wraps
-`chunk-finetune.ts` with the e4b-required env and a two-step workflow:
+`chunk-finetune.ts` with a two-step workflow:
 
 ```bash
 scripts/ft-e4b-v2.sh probe   # 2-iter memory/stability check (~1 min) — RUN FIRST
@@ -233,12 +234,10 @@ What the recipe pins (and why it differs from the bare defaults):
 | `SEG` | `4` | **segmented backward**, 4 layers/segment — so 8K-ctx activations fit |
 | `RANK` / `SCALE` / `LR` | `16` / `20` / `1e-5` | task-tuned |
 | `ITERS` | `2` (probe) / `900` (train) | ~2 epochs over 450 examples |
-| `MLX_BUN_PERF_KERNEL` / `MLX_BUN_FUSED_GELU` | `0` / `0` | **required** — the fused kernels have no vjp (see env table) |
 | attention | default `ops.sdpa` | mlx's fused flash kernel; **not** `MLX_BUN_TRAIN_ATTN=flash` (~30× slower; e4b ≥2K unvalidated since the kernel fix) |
 
-The two non-negotiables for e4b: **segmented backward** (`SEG>0`, so the
-long-context activations fit) and the **fused kernels off** (so the backward
-has gradients). Always run `probe` before `train`.
+The non-negotiable for e4b: **segmented backward** (`SEG>0`, so the
+long-context activations fit). Always run `probe` before `train`.
 
 ## Methodology
 
@@ -296,6 +295,15 @@ A finished run writes a PEFT-compatible adapter directory:
 - `optiq_lora_config.json` — mlx-bun/optiq adapter metadata (per-layer ranks)
 - `adapter_config.json` — PEFT-compatible config
 
+The save completes only after both config files are written, including
+eval-step checkpoints, so a successful process exit always leaves an
+immediately mountable directory. Serving accepts both mlx-lm tensor layout
+(`[in, rank]` / `[rank, out]`) and standard PEFT layout
+(`[rank, in]` / `[out, rank]`, including `base_model.model.*` names).
+Standard PEFT `use_rslora` metadata is honored as
+`lora_alpha / sqrt(rank)`, and saved rsLoRA adapters emit that metadata so
+their training and serving scales agree.
+
 When `save_checkpoints: true`, each eval step also writes
 `checkpoints/step-<NNNNN>-val<loss>/` and a durable `metrics.json` (config,
 wall seconds, peak GB, final/best train+val loss, full val trajectory).

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mlx-bun",
-  "version": "0.0.10",
+  "version": "0.0.12",
   "description": "Native MLX inference for Bun on Apple Silicon — local LLM server (OpenAI/Anthropic-compatible) and TypeScript library. No Python.",
   "type": "module",
   "license": "MIT",
@@ -55,13 +55,15 @@
   "devDependencies": {
     "@anthropic-ai/sdk": "^0.104.1",
     "@types/bun": "^1.3.14",
+    "happy-dom": "^20.10.6",
     "openai": "^6.42.0"
   },
   "dependencies": {
-    "@earendil-works/pi-coding-agent": "0.80.2",
-    "@earendil-works/pi-ai": "0.80.2",
+    "@earendil-works/pi-ai": "0.80.3",
+    "@earendil-works/pi-coding-agent": "0.80.3",
     "@huggingface/jinja": "^0.5.9",
     "@huggingface/tokenizers": "^0.1.3",
+    "@mlc-ai/web-xgrammar": "^0.1.27",
     "fast-png": "^8.0.0",
     "typebox": "1.1.38"
   }

package/src/anthropic.ts

@@ -150,6 +150,19 @@ export function anthropicToChatBody(body: AnthropicRequest): Record<string, unkn
             content: flattenText(block.content),
           });
           break;
+        case "audio":
+        case "input_audio":
+        case "audio_url":
+          // The Anthropic Messages protocol has no audio input block type.
+          // Reject loudly (the server maps this throw to a 400
+          // invalid_request_error) instead of silently dropping the clip —
+          // the model answering as if no audio was sent is worse than an
+          // error. Audio is served on the OpenAI surface.
+          throw new Error(
+            `"${block.type}" content blocks are not part of the Anthropic ` +
+              `Messages protocol — send audio to the OpenAI endpoint ` +
+              `/v1/chat/completions as an input_audio content part`,
+          );
         default:
           if (block.text) textParts.push(block.text);
       }
@@ -456,58 +469,68 @@ export function translateOpenAiSse(
   const dec = new TextDecoder();
   let buf = "";
   let finalized = false;
+  let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
+  let cancelled = false;
   return new ReadableStream<Uint8Array>({
-    async start(controller) {
+    start(controller) {
       const emit = (frames: string[]) => {
         for (const f of frames) controller.enqueue(enc.encode(f));
       };
-      const reader = upstream.getReader();
-      try {
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) break;
-          buf += dec.decode(value, { stream: true });
-          let sep: number;
-          while ((sep = buf.indexOf("\n\n")) !== -1) {
-            const frame = buf.slice(0, sep);
-            buf = buf.slice(sep + 2);
-            const data = frame
-              .split("\n")
-              .find((l) => l.startsWith("data: "))
-              ?.slice(6);
-            if (data == null) continue;
-            if (data === "[DONE]" || data === '"[DONE]"') {
-              if (!finalized) {
-                finalized = true;
-                emit(translator.finalize());
+      reader = upstream.getReader();
+      void (async () => {
+        try {
+          while (!cancelled) {
+            const { done, value } = await reader!.read();
+            if (done) break;
+            buf += dec.decode(value, { stream: true });
+            let sep: number;
+            while ((sep = buf.indexOf("\n\n")) !== -1) {
+              const frame = buf.slice(0, sep);
+              buf = buf.slice(sep + 2);
+              const data = frame
+                .split("\n")
+                .find((l) => l.startsWith("data: "))
+                ?.slice(6);
+              if (data == null) continue;
+              if (data === "[DONE]" || data === '"[DONE]"') {
+                if (!finalized) {
+                  finalized = true;
+                  emit(translator.finalize());
+                }
+                continue;
               }
-              continue;
-            }
-            const parsed = JSON.parse(data);
-            if (parsed?.error) {
-              emit([
-                sse("error", {
-                  type: "error",
-                  error: { type: "api_error", message: parsed.error.message ?? "generation failed" },
-                }),
-              ]);
-              continue;
+              const parsed = JSON.parse(data);
+              if (parsed?.error) {
+                emit([
+                  sse("error", {
+                    type: "error",
+                    error: { type: "api_error", message: parsed.error.message ?? "generation failed" },
+                  }),
+                ]);
+                continue;
+              }
+              emit(translator.addChunk(parsed));
             }
-            emit(translator.addChunk(parsed));
           }
+          if (!cancelled && !finalized) emit(translator.finalize());
+        } catch (e) {
+          if (!cancelled)
+            emit([
+              sse("error", {
+                type: "error",
+                error: { type: "api_error", message: (e as Error).message },
+              }),
+            ]);
+        } finally {
+          reader?.releaseLock();
+          reader = null;
+          if (!cancelled) controller.close();
         }
-        if (!finalized) emit(translator.finalize());
-      } catch (e) {
-        emit([
-          sse("error", {
-            type: "error",
-            error: { type: "api_error", message: (e as Error).message },
-          }),
-        ]);
-      } finally {
-        reader.releaseLock();
-        controller.close();
-      }
+      })();
+    },
+    async cancel(reason) {
+      cancelled = true;
+      await reader?.cancel(reason);
     },
   });
 }