mlx-bun 0.0.9 → 0.0.11

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 |
@@ -132,6 +132,11 @@ piece independently toggled and each falling back cleanly:
 - **`orpo_flash_ce`** — the [M,vocab]-free flash-CCE Metal head (steel GEMM fwd
   **and** bwd; e4b backward 754 ms, peak 0.93 GB flat at M=8192). Implies the
   fused head. Falls back to the MLX fused head when the tiling isn't clean.
+  **Auto-dispatch by M** (2026-07-02): batches with M below `MLX_BUN_FLASH_MIN_M`
+  (default 1024) take the exact MLX fused head instead — it is faster at short
+  M (e4b M=64 1.6×, M=512 1.13×; CPM5 at every M), while flash wins memory
+  everywhere and time at long M on e4b (M≥2048; 8K: 10.7 vs 13.2 s with the
+  filter/blockMax defaults). `MLX_BUN_FLASH_MIN_M=0` always honors flash.
 - **`segment_size`** — segmented backward (gradient-checkpointed layer
   activations) for long context; the head is ~free in memory, the activations
   are the wall.
@@ -162,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 |
@@ -174,8 +179,12 @@ are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
 | `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 |
 | `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) |
@@ -184,7 +193,7 @@ are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
 | `orpo_lr_schedule` | `constant` \| `cosine` | `cosine` | ORPO LR schedule (orpo only) |
 | `orpo_chunk_size` | int | `0` (off) | Token-chunk the response head, bounding the `[M,vocab]` logits to `[chunk,vocab]` (defaults to 512 when a fused head is on). Exact (orpo only) |
 | `orpo_fused_ce` | bool | `false` | Fused linear-CE head: one CustomVjp with an analytic softmax−onehot backward (MLX `quantizedMatmul` both ways) — no autograd through the head, no retained `[M,vocab]`. Exact; `[chunk,vocab]` transient (orpo only) |
-| `orpo_flash_ce` | bool | `false` | Route the fused head through the **flash-CCE Metal kernel** (verbatim MLX steel GEMM + ORPO epilogue, fwd **and** bwd): neither `[M,vocab]` nor a dequantized head touches HBM → `[M,vocab]`-free, fastest on large vocab (e4b bwd 754 ms, 0.93 GB flat @ M=8192). Implies the fused head. The Apple-CCE coeff filter / blockMax skip are opt-in via `MLX_BUN_CCE_BWD_FILTER_EPS` / `_BLOCK_EPS` (default off — exact). B=1 (orpo only) |
+| `orpo_flash_ce` | bool | `false` | Route the fused head through the **flash-CCE Metal kernel** (verbatim MLX steel GEMM + ORPO epilogue, fwd **and** bwd): neither `[M,vocab]` nor a dequantized head touches HBM → `[M,vocab]`-free (e4b 0.93 GB flat @ M=8192). Implies the fused head. The Apple-CCE coeff filter + blockMax skip default ON at 1e-5 (2026-07-02; combined bwd 1.71× CPM5 / 3.16× e4b vs exact — `MLX_BUN_CCE_BWD_FILTER_EPS=0` / `_BLOCK_EPS=0` restore exact). Batches with M < `MLX_BUN_FLASH_MIN_M` (default 1024) auto-take the exact fused head (faster at short M; `=0` always honors flash). B=1 (orpo only) |
 | `orpo_prefix_shared` | bool | `false` | Shared prompt-prefix: one forward over `[prompt; chosen; rejected]` with a block-sparse mask + block-wise RoPE, so the shared prompt is encoded **once** (a big win when the prompt dominates). Falls back to two-forward per row on prompt mismatch. Composes with the flash head and the segmented backward (MiniCPM5 + e4b). B=1 (orpo only) |
 | `sft_scope` | `full` \| `response` | `full` | Scope of ORPO's chosen-NLL term. `full` (paper/TRL-faithful): token-mean CE over the **full prompt+response** (only padding excluded), from the same chosen forward. `response`: `L_NLL = −ℓ_w` — reproduces pre-2026-07 runs bit-exactly. The odds-ratio `ℓ_w/ℓ_r` are response-only in both modes. Applies to every ORPO path (naive/chunked/fused/flash/prefix/segmented) (orpo only). CLI spelling: `mlx-bun train … --sft-scope full\|response` |
 
@@ -192,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)
 
@@ -210,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
@@ -228,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
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mlx-bun",
-  "version": "0.0.9",
+  "version": "0.0.11",
   "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);
       }