mlx-bun 0.0.8 → 0.0.10

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`; 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 + 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** |
 | **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 |
@@ -36,7 +36,7 @@ isolation + a single GPU lease across concurrent server requests).
 ### Quick start (HTTP API)
 
 ```bash
-curl -s localhost:8090/api/finetune/submit -X POST -H 'content-type: application/json' -d '{
+curl -s localhost:8080/api/finetune/submit -X POST -H 'content-type: application/json' -d '{
   "model_dir": "/path/to/snapshot",
   "data_dir":  "/path/to/dataset",        // dir with train.jsonl (+ optional valid.jsonl)
   "adapter_path": "/path/to/output-adapter",
@@ -81,7 +81,7 @@ differ from the trainer/API defaults in the table below:
 
 The script hard-codes `method=sft`, `batch_size=1`, `steps_per_report=1`, and
 uses the default `ops.sdpa` training attention (set `MLX_BUN_TRAIN_ATTN=flash`
-to override — but flash crashes e4b at multi-K; see Methodology).
+to override — experimental, much slower; see Methodology).
 
 ## Data formats
 
@@ -112,12 +112,19 @@ Probe a file before submitting: `POST /api/finetune/inspect-dataset` with
   and a preference term, `L = L_NLL(chosen) + λ·L_OR`, with
   `L_OR = -log σ(log_odds)` and
   `log_odds = (ℓ_w − ℓ_r) − (log1mexp(ℓ_w) − log1mexp(ℓ_r))`, where `ℓ` is the
-  **length-normalized** (mean over response tokens) log-prob. Uses the same
+  **length-normalized** (mean over response tokens) log-prob. The SFT term's
+  scope is `sft_scope`: **`full`** (default; paper/TRL-faithful) computes
+  `L_NLL` as the token-mean cross-entropy over the **full prompt+response** of
+  the chosen sequence (only padding excluded — TRL's `chosen_nll_loss`),
+  from the same chosen forward (no extra forward); **`response`** uses
+  `L_NLL = −ℓ_w` (response-only — the pre-2026-07 mlx-bun behavior, kept for
+  reproducing old runs). The odds-ratio `ℓ_w/ℓ_r` stay response-only in **both**
+  modes (that also matches TRL). Uses the same
   `{prompt, chosen, rejected}` preference data as DPO but needs **no reference
   model** (2 forwards/step vs DPO's 4). Default LR `1e-5` (lower than DPO — the
   loss carries a full SFT NLL term a high LR destabilizes). Tune with
-  `orpo_lambda`, `orpo_warmup_iters`, `orpo_lr_schedule`. Design + optimization
-  roadmap: [orpo-training](../design/orpo-training.md).
+  `orpo_lambda`, `orpo_warmup_iters`, `orpo_lr_schedule`, `sft_scope`. Design +
+  optimization roadmap: [orpo-training](../design/orpo-training.md).
 
 The ORPO head + long-context machinery compose into one stack (all B=1), each
 piece independently toggled and each falling back cleanly:
@@ -125,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.
@@ -177,8 +189,9 @@ 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` |
 
 ### Environment variables (training)
 
@@ -186,7 +199,7 @@ are `DEFAULT_TRAIN_CONFIG` (trainer.ts:89).
 |---|---|---|---|
 | `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 custom kernel that crashes e4b at multi-K; do not set it. |
+| `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))
@@ -221,7 +234,7 @@ What the recipe pins (and why it differs from the bare defaults):
 | `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` (that one crashes e4b) |
+| 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
@@ -265,9 +278,15 @@ At long `max_seq_length`, activation memory dominates. Two levers:
 ### Training attention kernel (`MLX_BUN_TRAIN_ATTN`)
 - **default `ops.sdpa`** — mlx's fused SDPA; correct (0.00% vs autograd),
   O(L²) backward memory. Use this.
-- **`flash`** — opt-in O(L) memory path, but the hand-rolled dK kernel is
-  slow and **crashes e4b at multi-K (≥2K)**; do not use it for e4b LoRA
-  training. Detail in segmented-backward-training §6.
+- **`flash`** — opt-in O(L)-memory path. Its two port bugs (dK
+  buffer-transpose, dQ divergent causal barrier) were **fixed** and
+  FD-validated (`scripts/experiments/flash-fd-check.ts`; regression test
+  `tests/flash-attention.test.ts`, T≤256), but it remains ~30× slower than
+  `ops.sdpa`, and the historical e4b multi-K (≥2K) SIGTRAP has **not been
+  re-validated at that scale since the fix**
+  (`scripts/experiments/segmented-grad-test-e4b.ts` is the repro harness).
+  Prefer the default; treat `flash` as experimental on e4b long context.
+  Detail in segmented-backward-training §6.
 
 ## Outputs
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mlx-bun",
-  "version": "0.0.8",
+  "version": "0.0.10",
   "description": "Native MLX inference for Bun on Apple Silicon — local LLM server (OpenAI/Anthropic-compatible) and TypeScript library. No Python.",
   "type": "module",
   "license": "MIT",

package/src/anthropic.ts

@@ -50,6 +50,18 @@ export interface AnthropicRequest {
   top_k?: number;
   stop_sequences?: string[];
   stream?: boolean;
+  /** mlx-lm sampler/penalty extensions — not part of the Anthropic protocol,
+   *  accepted as pass-through extras with the same names as our /v1 chat
+   *  surface. Anthropic has NO logit_bias, so none is accepted here. */
+  min_p?: number;
+  xtc_probability?: number;
+  xtc_threshold?: number;
+  repetition_penalty?: number;
+  repetition_context_size?: number;
+  presence_penalty?: number;
+  presence_context_size?: number;
+  frequency_penalty?: number;
+  frequency_context_size?: number;
   tools?: Array<{
     name?: string;
     description?: string;
@@ -167,6 +179,16 @@ export function anthropicToChatBody(body: AnthropicRequest): Record<string, unkn
   if (body.top_k != null) oai.top_k = body.top_k;
   if (body.stop_sequences?.length) oai.stop = body.stop_sequences;
   if (body.stream) oai.stream = true;
+  // mlx-lm sampler/penalty extensions: straight pass-through (same wire names
+  // on the chat surface). No logit_bias — it doesn't exist in this protocol.
+  for (const k of [
+    "min_p", "xtc_probability", "xtc_threshold",
+    "repetition_penalty", "repetition_context_size",
+    "presence_penalty", "presence_context_size",
+    "frequency_penalty", "frequency_context_size",
+  ] as const) {
+    if (body[k] != null) oai[k] = body[k];
+  }
 
   // Client tools only (entries with name + input_schema); server-tool
   // types (web_search etc.) have no local meaning and are dropped.

package/src/assets/curve-designer.html

@@ -89,7 +89,7 @@ button.ghost{background:var(--card);border:1px solid var(--bd);border-radius:8px
 <script>
 const ln=Math.log;
 const UMIN=ln(1e-6); // bottom-left / plot floor (matches src/curve-sampler.ts CURVE_UMIN)
-const ENGINE_BASE = location.protocol.startsWith("http") ? location.origin : "http://localhost:8090";
+const ENGINE_BASE = location.protocol.startsWith("http") ? location.origin : "http://localhost:8080";
 // ---- curve state: 5 movable control points, sorted ascending in p_in ----
 function identityPoints(){ return [{x_pct:1e-4,y_pct:1e-4},{x_pct:0.1,y_pct:0.1},{x_pct:1,y_pct:1},{x_pct:9,y_pct:9},{x_pct:100,y_pct:100}]; }
 let C = { space:"logprob", points:identityPoints(), monotonic:true };

package/src/chat-template.ts

@@ -131,20 +131,38 @@ export class ChatTemplate {
   /** True when the model has a switchable reasoning channel we can parse. */
   readonly supportsThinking: boolean;
 
-  private constructor(source: string, bosToken: string | null, eosToken: string | null) {
+  private constructor(
+    source: string,
+    bosToken: string | null,
+    eosToken: string | null,
+    forceNoThinking = false,
+  ) {
     this.#template = new Template(source);
     this.#bosToken = bosToken;
     this.#eosToken = eosToken;
     const gatesThinking = source.includes("enable_thinking");
-    this.thinkingFormat = gatesThinking && source.includes("<think>")
-      ? "think-tag"
-      : gatesThinking && source.includes("<|channel>")
-        ? "gemma-channel"
-        : null;
+    // `forceNoThinking` suppresses the switchable channel even when the template
+    // carries the gemma markers. DiffusionGemma ships the shared gemma-family
+    // template (with the `<|channel>thought…<channel|>` reasoning channel), but
+    // its non-autoregressive canvas decode cannot reliably emit the `<channel|>`
+    // close, so an enabled channel never ends (everything is captured as
+    // reasoning, no answer). The OptiQ reference never enables thinking for this
+    // model — it always renders the template's `default(false)` pre-closed empty
+    // channel and decodes the whole canvas as plain text. We match that.
+    this.thinkingFormat = forceNoThinking
+      ? null
+      : gatesThinking && source.includes("<think>")
+        ? "think-tag"
+        : gatesThinking && source.includes("<|channel>")
+          ? "gemma-channel"
+          : null;
     this.supportsThinking = this.thinkingFormat !== null;
   }
 
-  static async load(modelDir: string): Promise<ChatTemplate> {
+  static async load(
+    modelDir: string,
+    opts: { disableThinking?: boolean } = {},
+  ): Promise<ChatTemplate> {
     const config = (await Bun.file(`${modelDir}/tokenizer_config.json`).json()) as Record<string, any>;
     let source: string | undefined = config.chat_template;
     if (!source) {