mlx-bun 0.0.8 → 0.0.10

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (144) hide show
  1. package/README.md +112 -20
  2. package/docs/design/batching-perf-path.md +232 -0
  3. package/docs/design/batching-v2-plan.md +515 -0
  4. package/docs/design/bucketing-stage.md +196 -0
  5. package/docs/design/compat-cli-surface-design.md +11 -11
  6. package/docs/design/docs-reorg-plan.md +1 -1
  7. package/docs/design/dreaming-nightly-pipeline.md +291 -0
  8. package/docs/design/dspark-speculative-decoding.md +106 -0
  9. package/docs/design/generic-model-support.md +381 -0
  10. package/docs/design/memory-inference-path.md +422 -0
  11. package/docs/design/memory-synthesis.md +287 -0
  12. package/docs/design/minicpm5-decode-megakernel.md +484 -0
  13. package/docs/design/mlx-lm-tool-parity-plan.md +535 -0
  14. package/docs/design/omlx-adoption-map.md +90 -0
  15. package/docs/design/orpo-dynamic-lambda.md +41 -0
  16. package/docs/design/orpo-training.md +35 -5
  17. package/docs/design/parallel-slots.md +12 -1
  18. package/docs/design/parity-tier-dag.md +13 -5
  19. package/docs/design/segmented-backward-training.md +15 -0
  20. package/docs/design/spec-decode-larger-targets.md +229 -0
  21. package/docs/design/ssd-kv-cold-tier.md +158 -0
  22. package/docs/design/tauri-desktop-app.md +2 -2
  23. package/docs/design/the-dreaming-handoff.md +141 -0
  24. package/docs/design/the-dreaming-master-plan.md +446 -0
  25. package/docs/design/write-pipeline-entity.md +468 -0
  26. package/docs/investigations/curve-bisector-routing.md +1 -1
  27. package/docs/investigations/decode-roofline-lookagain.md +348 -0
  28. package/docs/investigations/dspark-handoff.md +58 -0
  29. package/docs/investigations/kernel-perf-review-2026-07.md +171 -0
  30. package/docs/planning/PRODUCT_ROADMAP.md +2 -2
  31. package/docs/planning/curve-sampler-research-plan.md +331 -0
  32. package/docs/planning/memory-docs-and-dag-plan.md +185 -0
  33. package/docs/planning/pi-integration-review.md +178 -0
  34. package/docs/planning/release-notes-v0.0.10.md +57 -0
  35. package/docs/planning/release-notes-v0.0.9.md +122 -0
  36. package/docs/planning/web-ui-pass-plan.md +213 -0
  37. package/docs/planning/website-readme-pass-plan.md +224 -0
  38. package/docs/reference/cli.md +255 -0
  39. package/docs/reference/embedding.md +7 -2
  40. package/docs/reference/library-api.md +22 -3
  41. package/docs/reference/memory.md +5 -5
  42. package/docs/reference/models.md +137 -0
  43. package/docs/reference/orpo-quickstart.md +6 -0
  44. package/docs/reference/server-api.md +187 -9
  45. package/docs/reference/server-config.md +43 -9
  46. package/docs/reference/training.md +31 -12
  47. package/package.json +1 -1
  48. package/src/anthropic.ts +22 -0
  49. package/src/assets/curve-designer.html +1 -1
  50. package/src/chat-template.ts +25 -7
  51. package/src/cli.ts +1077 -90
  52. package/src/config.ts +4 -1
  53. package/src/download.ts +99 -8
  54. package/src/embed.ts +77 -0
  55. package/src/eval/perplexity.ts +164 -0
  56. package/src/eval/runner.ts +15 -6
  57. package/src/eval/tasks/bfcl.ts +1 -1
  58. package/src/generate.ts +149 -18
  59. package/src/harness-pi.ts +43 -17
  60. package/src/index.ts +2 -0
  61. package/src/kv-store.ts +240 -98
  62. package/src/lora.ts +2 -2
  63. package/src/memory/article.ts +558 -0
  64. package/src/memory/chunk-validation.ts +150 -0
  65. package/src/memory/chunk.ts +305 -12
  66. package/src/memory/cluster.ts +344 -17
  67. package/src/memory/coverage.ts +134 -0
  68. package/src/memory/crosslink.ts +0 -0
  69. package/src/memory/db.ts +184 -45
  70. package/src/memory/entity.ts +203 -0
  71. package/src/memory/gate.ts +187 -0
  72. package/src/memory/history.ts +152 -0
  73. package/src/memory/ingest.ts +324 -11
  74. package/src/memory/model.ts +321 -0
  75. package/src/memory/normalize.ts +671 -0
  76. package/src/memory/parse.ts +22 -0
  77. package/src/memory/pipeline.ts +242 -30
  78. package/src/memory/prompts.ts +42 -0
  79. package/src/memory/query.ts +334 -0
  80. package/src/memory/reconcile.ts +734 -0
  81. package/src/memory/redirect.ts +108 -0
  82. package/src/memory/reindex.ts +325 -0
  83. package/src/memory/resolve.ts +307 -0
  84. package/src/memory/route.ts +272 -0
  85. package/src/memory/schedule.ts +6 -4
  86. package/src/memory/stages.ts +774 -0
  87. package/src/memory/synthesize.ts +1629 -20
  88. package/src/memory/tools.ts +446 -96
  89. package/src/memory/vault.ts +462 -35
  90. package/src/memory/wikify.ts +1057 -16
  91. package/src/mlx/array.ts +21 -0
  92. package/src/mlx/ffi.ts +3 -0
  93. package/src/mlx/ops.ts +72 -0
  94. package/src/model/compiled-decode.ts +43 -8
  95. package/src/model/factory.ts +29 -8
  96. package/src/model/flash-attention.ts +4 -2
  97. package/src/model/fused-decode-kernel.ts +14 -7
  98. package/src/model/fused-mlp-kernel.ts +169 -0
  99. package/src/model/fused-swiglu-kernel.ts +240 -0
  100. package/src/model/gemma4-base.ts +96 -7
  101. package/src/model/gemma4.ts +33 -2
  102. package/src/model/minicpm5.ts +106 -8
  103. package/src/model/qwen3-delta.ts +48 -1
  104. package/src/model/qwen3.ts +276 -0
  105. package/src/model/qwen3_5.ts +11 -2
  106. package/src/model/steel-linear-kernel.ts +96 -0
  107. package/src/model/support.ts +38 -8
  108. package/src/model/universal/archs.ts +332 -0
  109. package/src/model/universal/dense.ts +576 -0
  110. package/src/model/universal/modules.ts +163 -0
  111. package/src/model/universal/rope.ts +366 -0
  112. package/src/pi-provider.ts +1 -1
  113. package/src/pi-terminal.ts +29 -8
  114. package/src/pi-web.ts +19 -2
  115. package/src/prompt-cache.ts +55 -9
  116. package/src/quantize/config-writer.ts +1 -1
  117. package/src/registry.ts +219 -8
  118. package/src/responses.ts +23 -0
  119. package/src/sampler.ts +220 -9
  120. package/src/serve/batch-scheduler.ts +307 -83
  121. package/src/serve/generation-gateway.ts +156 -29
  122. package/src/server.ts +746 -73
  123. package/src/spec/drafter.ts +81 -49
  124. package/src/spec/dspark/data-dflash.ts +156 -0
  125. package/src/spec/dspark/data.ts +201 -0
  126. package/src/spec/dspark/generate-dflash.ts +193 -0
  127. package/src/spec/dspark/generate.ts +290 -0
  128. package/src/spec/dspark/loss.ts +213 -0
  129. package/src/spec/dspark/module-dflash.ts +361 -0
  130. package/src/spec/dspark/module.ts +540 -0
  131. package/src/spec/dspark/sample.ts +112 -0
  132. package/src/spec/generate.ts +40 -19
  133. package/src/ssd-cache.ts +213 -0
  134. package/src/status-page.html +1 -1
  135. package/src/tokenizer.ts +7 -0
  136. package/src/train/flash-cce.ts +81 -23
  137. package/src/train/fuse.ts +243 -0
  138. package/src/train/index.ts +3 -0
  139. package/src/train/job.ts +8 -0
  140. package/src/train/loss.ts +341 -40
  141. package/src/train/prefix-shared.ts +62 -11
  142. package/src/train/segmented.ts +215 -55
  143. package/src/train/trainer.ts +112 -38
  144. package/src/web/skills/memory/SKILL.md +40 -17
@@ -34,7 +34,7 @@ import {
34
34
  saveAdapter, disposeLora, warmStartFromAdapter, type TrainableLora, type SaveAdapterConfig,
35
35
  } from "./lora-params";
36
36
  import { resolveRanks, bitsMapFromModel, readPerLayerKl, DEFAULT_TARGET_MODULES, type RankScaling } from "./rank";
37
- import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics, orpoLoss, orpoMetrics, type ChunkCtx } from "./loss";
37
+ import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics, orpoLoss, orpoMetrics, type ChunkCtx, type SftScope } from "./loss";
38
38
  import { orpoLossPrefixShared, orpoLossPrefixSharedGemma, splitPrefixBatch } from "./prefix-shared";
39
39
  import { AdamW, warmupCosineSchedule } from "./optimizer";
40
40
  import { writeFileSync, appendFileSync, mkdirSync } from "node:fs";
@@ -134,6 +134,15 @@ export interface TrainConfig {
134
134
  * differ. Composes with orpoFlashCe/orpoFusedCe (each branch routes through the
135
135
  * [M,V]-free head). See src/train/prefix-shared.ts. */
136
136
  orpoPrefixShared: boolean;
137
+ /** ORPO L_SFT scope (`sft_scope`). "full" (default; paper/TRL-faithful):
138
+ * the chosen-NLL term is the token-mean cross-entropy over the FULL
139
+ * prompt+response (only padding excluded), computed from the same chosen
140
+ * forward — matching TRL's ORPOTrainer chosen_nll_loss. "response":
141
+ * L_NLL = -ℓw (response-only), bit-exact to pre-sft_scope runs. The
142
+ * odds-ratio ℓw/ℓr terms are response-only length-normalized means in BOTH
143
+ * modes (that also matches TRL). Applies to every ORPO path (naive /
144
+ * chunked / fused / flash / prefix-shared / segmented). orpo only. */
145
+ sftScope: SftScope;
137
146
  /** Warm-start: path to an existing adapter dir (with adapters.safetensors) whose
138
147
  * LoRA weights initialize this run instead of the random/zero init. Continues
139
148
  * training from a checkpoint's weights; the optimizer state and LR schedule restart.
@@ -188,6 +197,7 @@ export const DEFAULT_TRAIN_CONFIG: TrainConfig = {
188
197
  orpoFusedCe: false, // default off; true = fused linear-CE head (analytic backward, no [M,vocab])
189
198
  orpoFlashCe: false, // default off; true = flash-CCE Metal-kernel head (implies fused)
190
199
  orpoPrefixShared: false, // default off; true = shared prompt-prefix single forward (B=1)
200
+ sftScope: "full", // paper/TRL-faithful chosen-NLL over prompt+response; "response" = pre-2026-07 behavior
191
201
  warmStartAdapter: "", // default off; path to an adapter dir to warm-start LoRA weights from
192
202
  adapterPath: "adapters",
193
203
  baseModel: "",
@@ -298,6 +308,7 @@ export async function trainLora(
298
308
  iters: cfg.iters, learning_rate: cfg.learningRate, max_seq_length: cfg.maxSeqLen,
299
309
  batch_size: cfg.batchSize, rank: cfg.rank, scale: cfg.scale,
300
310
  orpo_lambda: cfg.method === "orpo" ? cfg.orpoLambda : undefined,
311
+ sft_scope: cfg.method === "orpo" ? cfg.sftScope : undefined,
301
312
  adapter_path: cfg.adapterPath,
302
313
  });
303
314
 
@@ -349,6 +360,7 @@ export async function trainLora(
349
360
  segmentSize: cfg.segmentSize, iters: cfg.iters,
350
361
  numLayers: cfg.numLayers, baseModel: cfg.baseModel,
351
362
  orpoLambda: cfg.method === "orpo" ? cfg.orpoLambda : undefined,
363
+ sftScope: cfg.method === "orpo" ? cfg.sftScope : undefined,
352
364
  },
353
365
  wallSeconds: (Date.now() - startedAt) / 1000,
354
366
  peakGb: peakMemory() / 1e9,
@@ -376,6 +388,16 @@ export async function trainLora(
376
388
  }
377
389
  }
378
390
 
391
+ // ---------------------------------------------------------------------------
392
+ // Segment planning: uniform planSegmentsBySize only. A full-attention
393
+ // ISOLATION planner (kernel-review backlog #2 / segmented-backward-training.md
394
+ // §5) was built and A/B-measured 2026-07-02: ZERO peak win — mlx's sdpa
395
+ // backward is O(L²) for EVERY layer (~3.5 GB/layer @8K e4b), sliding included,
396
+ // so the worst segment is set by layer count alone and segment_size is the
397
+ // whole knob (seg1 = 14.59 GB @8K, +3% step time — fits the 24 GB M4 Pro).
398
+ // Evidence: scripts/experiments/seg-isolation-smoke.ts; note in segmented.ts.
399
+ // ---------------------------------------------------------------------------
400
+
379
401
  // ---------------------------------------------------------------------------
380
402
  // SFT loop
381
403
  // ---------------------------------------------------------------------------
@@ -409,14 +431,14 @@ async function sftLoop(
409
431
  const useSegmented = cfg.segmentSize > 0;
410
432
  let segmented: SegmentedBackward | SegmentedBackwardGemma4 | null = null;
411
433
  if (useSegmented) {
412
- if (model instanceof MiniCPM5Model)
413
- segmented = new SegmentedBackward(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
414
- else if (model instanceof Gemma4Model)
415
- segmented = new SegmentedBackwardGemma4(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
416
- else
434
+ if (!(model instanceof MiniCPM5Model) && !(model instanceof Gemma4Model))
417
435
  throw new Error("segmented backward (segmentSize > 0) is only wired for MiniCPM5 and Gemma4");
436
+ const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
437
+ segmented = model instanceof MiniCPM5Model
438
+ ? new SegmentedBackward(model, lora, ranges)
439
+ : new SegmentedBackwardGemma4(model, lora, ranges);
418
440
  emit({ type: "stage", stage: "setup", progress: 0.04,
419
- message: `segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers` });
441
+ message: `segmented backward: ${ranges.length} segments of <=${cfg.segmentSize} layers` });
420
442
  }
421
443
 
422
444
  // The value_and_grad loss closure: temporarily swap the differentiated
@@ -466,14 +488,40 @@ async function sftLoop(
466
488
  // dQ-causal-barrier bugs were fixed; FD-validated in
467
489
  // scripts/experiments/flash-fd-check.ts), though ~30× slower than ops.sdpa.
468
490
  // Cleared in finally.
469
- if (process.env.MLX_BUN_TRAIN_ATTN === "flash") setTrainingAttn("flash");
491
+ //
492
+ // GEMMA GUARD (same convention as the PERF_KERNEL/FUSED_GELU sanitization in
493
+ // cli.ts / job.ts, but enforced at the trainer): e4b on this path SIGTRAPed
494
+ // (uncatchable native crash) at multi-K sequence lengths (>=2K, reproduced in
495
+ // scripts/experiments/segmented-grad-test-e4b.ts; docs/reference/training.md)
496
+ // and has NOT been re-validated at that scale since the two kernel fixes —
497
+ // the regression tests stop at T<=256. cli.ts even defaults e4b seq to 8192,
498
+ // so a stale `export MLX_BUN_TRAIN_ATTN=flash` from a MiniCPM5 experiment
499
+ // would otherwise ride silently into a crash mid-run. Refuse it for Gemma
500
+ // until the >=2K re-validation lands; MiniCPM5 stays allowed.
501
+ if (process.env.MLX_BUN_TRAIN_ATTN === "flash") {
502
+ if (model instanceof Gemma4Model)
503
+ throw new Error(
504
+ "MLX_BUN_TRAIN_ATTN=flash is disabled for Gemma models: e4b SIGTRAPs on this " +
505
+ "path at seq >= 2048 (docs/reference/training.md) and it has not been " +
506
+ "re-validated at that scale since the kernel fixes. Unset MLX_BUN_TRAIN_ATTN " +
507
+ "(ops.sdpa, the default, is exact and ~30x faster) or train MiniCPM5.",
508
+ );
509
+ emit({ type: "stage", stage: "setup", progress: 0.045,
510
+ message: "training attention: hand-rolled flash kernel (MLX_BUN_TRAIN_ATTN=flash; O(L) memory, ~30x slower than ops.sdpa)" });
511
+ setTrainingAttn("flash");
512
+ }
470
513
 
471
514
  // Pad with the eos id when available (a real sentinel token); padded
472
515
  // positions are excluded from both the loss and attention masks, so the
473
516
  // exact pad value never affects the result — it just must be a valid id.
474
517
  const padId = tok.eosTokenId ?? 0;
475
518
  const batches = iterateSftBatches(train, cfg.batchSize, cfg.maxSeqLen, cfg.seed, true, padId);
476
- const t0 = Date.now();
519
+ // tokens/sec follows mlx-lm's tuner convention: accumulate the response-token
520
+ // count across the REPORT WINDOW and divide by the window's own wall time
521
+ // (validation time excluded) — not one step's count extrapolated over the
522
+ // whole run, which jumps ~10x step-to-step on variable-length data.
523
+ let windowTokens = 0;
524
+ let windowStart = Date.now();
477
525
 
478
526
  try {
479
527
  for (let step = 1; step <= cfg.iters; step++) {
@@ -483,12 +531,11 @@ async function sftLoop(
483
531
  // when 1). afterMicroEval frees each micro-batch's gradient-checkpoint
484
532
  // closures right after its backward so they don't pile up across
485
533
  // accumulation.
486
- let tokensThisStep = 0;
487
534
  const { loss: lossVal, grads } = accumulateStep(
488
535
  cfg.gradAccumSteps,
489
536
  () => {
490
537
  currentBatch = batches.next().value as SftBatch;
491
- tokensThisStep += countResponseTokens(currentBatch);
538
+ windowTokens += countResponseTokens(currentBatch);
492
539
  return segmented ? segmented.step(currentBatch!) : vag!.apply(flatParams(lora));
493
540
  },
494
541
  ckptCtx
@@ -512,7 +559,7 @@ async function sftLoop(
512
559
  }
513
560
 
514
561
  if (step % cfg.stepsPerReport === 0 || step === 1) {
515
- const elapsed = (Date.now() - t0) / 1000;
562
+ const elapsed = (Date.now() - windowStart) / 1000;
516
563
  emit({
517
564
  type: "metric",
518
565
  kind: "train",
@@ -520,19 +567,31 @@ async function sftLoop(
520
567
  loss: lossVal,
521
568
  grad_norm: gradNorm,
522
569
  learning_rate: opt.lr,
523
- tokens_per_sec: elapsed > 0 ? (step * tokensThisStep) / elapsed : 0,
570
+ tokens_per_sec: elapsed > 0 ? windowTokens / elapsed : 0,
524
571
  progress: step / cfg.iters,
525
572
  message: `step ${step}/${cfg.iters} loss=${lossVal.toFixed(4)}`,
526
573
  });
574
+ windowTokens = 0;
575
+ windowStart = Date.now();
527
576
  }
528
577
 
529
578
  if (valid.length > 0 && step % cfg.stepsPerEval === 0) {
530
579
  // Eval is forward-only (no backward), so disable checkpointing for it —
531
580
  // it adds no benefit and a forward-only pass fits without it.
532
581
  if (ckptCtx) (model as Gemma4Model).gradCkpt = null;
533
- const vLoss = evalSftLoss(model, valid, cfg);
582
+ const tVal = Date.now();
583
+ const { loss: vLoss, used, skipped } = evalSftLoss(model, valid, cfg);
534
584
  if (ckptCtx) (model as Gemma4Model).gradCkpt = ckptCtx;
535
- emit({ type: "metric", kind: "val", step, loss: vLoss, progress: step / cfg.iters });
585
+ // A silently-shrinking val set must be visible: rows whose supervised
586
+ // span is empty after truncation are skipped BY DESIGN (mlx-lm drops
587
+ // them at dataset build time) — report how many, and over what n the
588
+ // mean was taken. Any other per-row failure now rethrows (see
589
+ // evalSftLoss).
590
+ emit({ type: "metric", kind: "val", step, loss: vLoss, progress: step / cfg.iters,
591
+ val_rows_used: used, val_rows_skipped: skipped,
592
+ ...(skipped > 0 ? { message: `val loss over ${used} rows (${skipped} skipped: no response tokens to supervise)` } : {}) });
593
+ // validation wall time is not training throughput — keep it out of the window
594
+ windowStart += Date.now() - tVal;
536
595
  }
537
596
  }
538
597
  } finally {
@@ -548,13 +607,16 @@ async function sftLoop(
548
607
  return { numIters: cfg.iters };
549
608
  }
550
609
 
551
- function evalSftLoss(model: RuntimeModel, valid: ReturnType<typeof Array.prototype.slice>, cfg: TrainConfig): number {
610
+ function evalSftLoss(
611
+ model: RuntimeModel, valid: ReturnType<typeof Array.prototype.slice>, cfg: TrainConfig,
612
+ ): { loss: number; used: number; skipped: number } {
552
613
  let total = 0;
553
614
  let n = 0;
615
+ let skipped = 0;
554
616
  for (const ex of valid as { ids: number[]; promptLen: number }[]) {
555
617
  let ids = ex.ids;
556
618
  let promptLen = ex.promptLen;
557
- if (ids.length < 2) continue;
619
+ if (ids.length < 2) { skipped++; continue; }
558
620
  if (ids.length > cfg.maxSeqLen) {
559
621
  ids = ids.slice(0, cfg.maxSeqLen);
560
622
  promptLen = Math.min(promptLen, Math.max(0, ids.length - 1));
@@ -565,12 +627,18 @@ function evalSftLoss(model: RuntimeModel, valid: ReturnType<typeof Array.prototy
565
627
  total += loss.toFloat32()[0]!;
566
628
  loss.dispose();
567
629
  n++;
568
- } catch {
569
- // skip examples with no response tokens
630
+ } catch (e) {
631
+ // ONLY the empty-supervision row is skippable (a truncated row whose
632
+ // response fell entirely past maxSeqLen — mlx-lm drops these at dataset
633
+ // build). Everything else (MLX OOM/shape/vjp errors surface as catchable
634
+ // JS Errors via the ffi handler) must rethrow: swallowing them silently
635
+ // reduces the val set to its easy rows and reports a plausible loss.
636
+ if (e instanceof Error && e.message === "sftLoss: no response tokens to supervise") skipped++;
637
+ else throw e;
570
638
  }
571
639
  clearCache();
572
640
  }
573
- return n > 0 ? total / n : 0;
641
+ return { loss: n > 0 ? total / n : 0, used: n, skipped };
574
642
  }
575
643
 
576
644
  // ---------------------------------------------------------------------------
@@ -774,28 +842,28 @@ async function orpoLoop(
774
842
  let segmentedOrpo: SegmentedBackwardOrpo | SegmentedBackwardOrpoGemma4 | null = null;
775
843
  let segmentedOrpoPrefix: SegmentedBackwardOrpoPrefix | SegmentedBackwardOrpoPrefixGemma4 | null = null;
776
844
  if (useSegmented) {
845
+ if (!(model instanceof MiniCPM5Model) && !(model instanceof Gemma4Model))
846
+ throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
777
847
  const ranges = planSegmentsBySize(model.layers.length, cfg.segmentSize);
778
848
  if (useSegmentedPrefix) {
779
849
  // The composed prefix-shared segmented backward (MiniCPM5 or Gemma4/e4b), plus a
780
850
  // plain segmented-ORPO fallback for rows whose chosen/rejected prompts differ
781
851
  // (splitPrefixBatch returns null -> two-forward segmented step).
782
852
  if (model instanceof MiniCPM5Model) {
783
- segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefix(model, lora, ranges, cfg.orpoLambda, segPrefixChunk);
784
- segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
853
+ segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefix(model, lora, ranges, cfg.orpoLambda, segPrefixChunk, cfg.sftScope);
854
+ segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
785
855
  } else {
786
- segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefixGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segPrefixChunk);
787
- segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk);
856
+ segmentedOrpoPrefix = new SegmentedBackwardOrpoPrefixGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segPrefixChunk, cfg.sftScope);
857
+ segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
788
858
  }
789
859
  emit({ type: "stage", stage: "setup", progress: 0.04,
790
- 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)` });
860
+ 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)` });
791
861
  } else if (model instanceof MiniCPM5Model)
792
- segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
793
- else if (model instanceof Gemma4Model)
794
- segmentedOrpo = new SegmentedBackwardOrpoGemma4(model, lora, ranges, cfg.orpoLambda, segFusedChunk);
862
+ segmentedOrpo = new SegmentedBackwardOrpo(model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
795
863
  else
796
- throw new Error("segmented backward (segmentSize > 0) for ORPO is only wired for MiniCPM5 and Gemma4");
864
+ segmentedOrpo = new SegmentedBackwardOrpoGemma4(model as Gemma4Model, lora, ranges, cfg.orpoLambda, segFusedChunk, cfg.sftScope);
797
865
  if (!useSegmentedPrefix) emit({ type: "stage", stage: "setup", progress: 0.04,
798
- message: `orpo segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers${segFusedChunk > 0 ? ` + fused linear-CE head (${segFusedChunk}/chunk)` : ""}` });
866
+ message: `orpo segmented backward: ${ranges.length} segments of <=${cfg.segmentSize} layers${segFusedChunk > 0 ? ` + fused linear-CE head (${segFusedChunk}/chunk)` : ""}` });
799
867
  }
800
868
 
801
869
  // Token-chunked head (non-segmented path only): the per-chunk checkpoints are
@@ -837,12 +905,12 @@ async function orpoLoop(
837
905
  const split = splitPrefixBatch(currentBatch!);
838
906
  if (split) {
839
907
  if (model instanceof MiniCPM5Model)
840
- return orpoLossPrefixShared(model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk);
841
- return orpoLossPrefixSharedGemma(model as Gemma4Model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk);
908
+ return orpoLossPrefixShared(model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk, cfg.sftScope);
909
+ return orpoLossPrefixSharedGemma(model as Gemma4Model, split.promptIds, split.chosenResp, split.rejectedResp, cfg.orpoLambda, chunk, cfg.sftScope);
842
910
  }
843
911
  notePrefixFallback();
844
912
  }
845
- return orpoLoss(model, currentBatch!, cfg.orpoLambda, chunk);
913
+ return orpoLoss(model, currentBatch!, cfg.orpoLambda, chunk, cfg.sftScope);
846
914
  } finally {
847
915
  restorePrimals(lora, saved);
848
916
  }
@@ -897,7 +965,7 @@ async function orpoLoop(
897
965
  // Diagnostics recompute the two forwards (like dpoMetrics) — cheap at
898
966
  // B=1 (response-only head, no full logits) and gated by stepsPerReport.
899
967
  // Reusing the value_and_grad forward is a shared-with-DPO follow-on.
900
- const m = orpoMetrics(model, currentBatch!, cfg.orpoLambda);
968
+ const m = orpoMetrics(model, currentBatch!, cfg.orpoLambda, cfg.sftScope);
901
969
  const elapsed = (Date.now() - t0) / 1000;
902
970
  emit({
903
971
  type: "metric",
@@ -929,7 +997,7 @@ async function orpoLoop(
929
997
  chosenIds: [ex.chosenIds], rejectedIds: [ex.rejectedIds],
930
998
  chosenMask: [ex.chosenMask], rejectedMask: [ex.rejectedMask],
931
999
  };
932
- const m = orpoMetrics(model, vb, cfg.orpoLambda);
1000
+ const m = orpoMetrics(model, vb, cfg.orpoLambda, cfg.sftScope);
933
1001
  vl += m.loss; va += m.accuracy; vm += m.margin; vn++;
934
1002
  clearCache();
935
1003
  }
@@ -1105,10 +1173,16 @@ function clipGradsByNorm(grads: MlxArray[], normVal: number, maxNorm: number): v
1105
1173
  }
1106
1174
  }
1107
1175
 
1176
+ /** Supervised (response) tokens across ALL rows of the batch — the count
1177
+ * tokens/sec is reported over (mlx-lm's n_tokens counts every row too). */
1108
1178
  function countResponseTokens(batch: SftBatch): number {
1109
- const promptLen = batch.promptLens[0]!;
1110
- const len = batch.ids[0]!.length;
1111
- return Math.max(0, len - Math.max(promptLen, 1));
1179
+ let total = 0;
1180
+ for (let r = 0; r < batch.ids.length; r++) {
1181
+ const promptLen = batch.promptLens[r]!;
1182
+ const len = batch.ids[r]!.length;
1183
+ total += Math.max(0, len - Math.max(promptLen, 1));
1184
+ }
1185
+ return total;
1112
1186
  }
1113
1187
 
1114
1188
  async function fileExists(path: string): Promise<boolean> {
@@ -20,26 +20,49 @@ the user's machine. Nothing about their memory leaves the device.
20
20
 
21
21
  When the answer depends on personal continuity — an explicit memory/recall
22
22
  request, or a named project, person, decision, preference, or history that is
23
- not explained in the current chat — consult memory before answering or asking
24
- the user to re-explain.
23
+ not explained in the current chat — look it up **before** answering or asking
24
+ the user to re-explain. Doing the lookup first, instead of answering from first
25
+ principles, is itself the helpful move.
25
26
 
26
27
  Do **not** use memory for weather, current public facts, generic web research,
27
28
  or ordinary coding/file tasks when the current files or the web tools answer the
28
29
  question. Memory is for user-specific context, not for performing familiarity.
29
30
 
30
- 1. **Search.** Call `memory_search` with the topic (plain phrasing is fine —
31
- a name or key noun). It returns ranked articles and sample lines.
32
- 2. **Read enough.** For short or central articles, call `memory_read`. For long
33
- articles, use `memory_toc` then `memory_section` to read the relevant part.
34
- The article is the answer — read it rather than reconstructing the user's
35
- position from scattered fragments. Treat it as their current view.
36
- 3. **Follow nearby context when useful.** Use `memory_links` if backlinks or
37
- related concepts matter. `memory_list` gives an overview when search comes up
38
- short. mlx-bun reference docs are separate from personal memory; use
39
- `reference_search`, `reference_read`, and `reference_list` for those.
40
- `memory_status` answers setup/schedule/location questions.
41
- 4. **Use it silently, as a prior.** Let what you find shape your response like
42
- a long-time colleague would — don't announce "I searched your memory."
31
+ The read path is **FIND the article → READ it small → answer**. It is
32
+ deterministic navigation — never a vector search, and there is no way to
33
+ query/sort articles by an infobox field (the infobox is facts you *read*, via
34
+ `memory_infobox`, not a filter).
35
+
36
+ 1. **FIND the article.**
37
+ - When the user names a specific thing, person, or project, call
38
+ `memory_resolve` with the name as said ("S5 IIX", "the Lumix"). It returns
39
+ the matched stem, title, kind, and lead — no search, no guessing.
40
+ - For "what do I have about X", call `memory_category` with a `category`
41
+ (e.g. `Lenses`), `type`, or `series`. This is membership, not search: it
42
+ lists the articles that are actually about X.
43
+ 2. **READ it small (the default).** Call `memory_read` to get the article's TOC
44
+ + lead, then `memory_section` to read just the **one** section you need. The
45
+ whole article is the exception — pass `force=true` only when you truly need
46
+ all of it. The article is the user's consolidated current position; read it
47
+ rather than reconstructing their view from scattered fragments.
48
+ 3. **Follow the graph / read the infobox when useful.** `memory_links` shows
49
+ outgoing links grouped by origin (infobox / series / see also / prose) plus
50
+ backlinks — use it to hop to a neighbour (e.g. from a camera's infobox
51
+ `mount: [[L-Mount]]` to the mount article). `memory_infobox` surfaces the
52
+ structured facts to read. `memory_list` is a bare overview;
53
+ `memory_status` answers setup/schedule/location questions. mlx-bun reference
54
+ docs are separate from personal memory; use `reference_search`,
55
+ `reference_read`, and `reference_list` for those.
56
+ 4. **`memory_search` is the LAST resort.** Only fall back to it when
57
+ `memory_resolve` / `memory_category` have failed across a couple of
58
+ phrasings. It is a plain substring scan that doesn't understand meaning.
59
+ 5. **Use it silently, as a continuation.** Let what you find shape your response
60
+ like a long-time colleague who already knows the state of things. Speak as a
61
+ continuation, never as a retrieval report: do not preface answers with "per
62
+ the wiki…", "according to the article…", "I checked your memory…", or
63
+ "welcome back, I see you're working on X". Cite an article by name only if
64
+ the user asks where something came from, or to distinguish "the wiki says X"
65
+ from "I think Y".
43
66
 
44
67
  Use memory as an invisible prior: it should narrow relevance and improve the
45
68
  answer, not become the topic of the answer. Avoid performative recall — do not
@@ -85,8 +108,8 @@ wiki, scheduling nightly synthesis). Confirm before each persistent action.
85
108
 
86
109
  ## What not to do
87
110
 
88
- - Don't claim a fact "from memory" that you didn't read in an article — search
89
- and read first, or say you don't have it.
111
+ - Don't claim a fact "from memory" that you didn't read in an article — find
112
+ and read the article first, or say you don't have it.
90
113
  - Don't write to or edit the wiki during a normal conversation; the read tools
91
114
  are read-only by design. Synthesis is a separate, deliberate step.
92
115
  - Do open the wiki or a specific article in Obsidian/Finder if the user asks to