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.
- package/README.md +112 -20
- package/docs/design/batching-perf-path.md +232 -0
- package/docs/design/batching-v2-plan.md +515 -0
- package/docs/design/bucketing-stage.md +196 -0
- package/docs/design/compat-cli-surface-design.md +11 -11
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dreaming-nightly-pipeline.md +291 -0
- package/docs/design/dspark-speculative-decoding.md +106 -0
- package/docs/design/generic-model-support.md +381 -0
- package/docs/design/memory-inference-path.md +422 -0
- package/docs/design/memory-synthesis.md +287 -0
- package/docs/design/minicpm5-decode-megakernel.md +484 -0
- package/docs/design/mlx-lm-tool-parity-plan.md +535 -0
- package/docs/design/omlx-adoption-map.md +90 -0
- package/docs/design/orpo-dynamic-lambda.md +41 -0
- package/docs/design/orpo-training.md +35 -5
- package/docs/design/parallel-slots.md +12 -1
- package/docs/design/parity-tier-dag.md +13 -5
- package/docs/design/segmented-backward-training.md +15 -0
- package/docs/design/spec-decode-larger-targets.md +229 -0
- package/docs/design/ssd-kv-cold-tier.md +158 -0
- package/docs/design/tauri-desktop-app.md +2 -2
- package/docs/design/the-dreaming-handoff.md +141 -0
- package/docs/design/the-dreaming-master-plan.md +446 -0
- package/docs/design/write-pipeline-entity.md +468 -0
- package/docs/investigations/curve-bisector-routing.md +1 -1
- package/docs/investigations/decode-roofline-lookagain.md +348 -0
- package/docs/investigations/dspark-handoff.md +58 -0
- package/docs/investigations/kernel-perf-review-2026-07.md +171 -0
- package/docs/planning/PRODUCT_ROADMAP.md +2 -2
- package/docs/planning/curve-sampler-research-plan.md +331 -0
- package/docs/planning/memory-docs-and-dag-plan.md +185 -0
- package/docs/planning/pi-integration-review.md +178 -0
- package/docs/planning/release-notes-v0.0.10.md +57 -0
- package/docs/planning/release-notes-v0.0.9.md +122 -0
- package/docs/planning/web-ui-pass-plan.md +213 -0
- package/docs/planning/website-readme-pass-plan.md +224 -0
- package/docs/reference/cli.md +255 -0
- package/docs/reference/embedding.md +7 -2
- package/docs/reference/library-api.md +22 -3
- package/docs/reference/memory.md +5 -5
- package/docs/reference/models.md +137 -0
- package/docs/reference/orpo-quickstart.md +6 -0
- package/docs/reference/server-api.md +187 -9
- package/docs/reference/server-config.md +43 -9
- package/docs/reference/training.md +31 -12
- package/package.json +1 -1
- package/src/anthropic.ts +22 -0
- package/src/assets/curve-designer.html +1 -1
- package/src/chat-template.ts +25 -7
- package/src/cli.ts +1077 -90
- package/src/config.ts +4 -1
- package/src/download.ts +99 -8
- package/src/embed.ts +77 -0
- package/src/eval/perplexity.ts +164 -0
- package/src/eval/runner.ts +15 -6
- package/src/eval/tasks/bfcl.ts +1 -1
- package/src/generate.ts +149 -18
- package/src/harness-pi.ts +43 -17
- package/src/index.ts +2 -0
- package/src/kv-store.ts +240 -98
- package/src/lora.ts +2 -2
- package/src/memory/article.ts +558 -0
- package/src/memory/chunk-validation.ts +150 -0
- package/src/memory/chunk.ts +305 -12
- package/src/memory/cluster.ts +344 -17
- package/src/memory/coverage.ts +134 -0
- package/src/memory/crosslink.ts +0 -0
- package/src/memory/db.ts +184 -45
- package/src/memory/entity.ts +203 -0
- package/src/memory/gate.ts +187 -0
- package/src/memory/history.ts +152 -0
- package/src/memory/ingest.ts +324 -11
- package/src/memory/model.ts +321 -0
- package/src/memory/normalize.ts +671 -0
- package/src/memory/parse.ts +22 -0
- package/src/memory/pipeline.ts +242 -30
- package/src/memory/prompts.ts +42 -0
- package/src/memory/query.ts +334 -0
- package/src/memory/reconcile.ts +734 -0
- package/src/memory/redirect.ts +108 -0
- package/src/memory/reindex.ts +325 -0
- package/src/memory/resolve.ts +307 -0
- package/src/memory/route.ts +272 -0
- package/src/memory/schedule.ts +6 -4
- package/src/memory/stages.ts +774 -0
- package/src/memory/synthesize.ts +1629 -20
- package/src/memory/tools.ts +446 -96
- package/src/memory/vault.ts +462 -35
- package/src/memory/wikify.ts +1057 -16
- package/src/mlx/array.ts +21 -0
- package/src/mlx/ffi.ts +3 -0
- package/src/mlx/ops.ts +72 -0
- package/src/model/compiled-decode.ts +43 -8
- package/src/model/factory.ts +29 -8
- package/src/model/flash-attention.ts +4 -2
- package/src/model/fused-decode-kernel.ts +14 -7
- package/src/model/fused-mlp-kernel.ts +169 -0
- package/src/model/fused-swiglu-kernel.ts +240 -0
- package/src/model/gemma4-base.ts +96 -7
- package/src/model/gemma4.ts +33 -2
- package/src/model/minicpm5.ts +106 -8
- package/src/model/qwen3-delta.ts +48 -1
- package/src/model/qwen3.ts +276 -0
- package/src/model/qwen3_5.ts +11 -2
- package/src/model/steel-linear-kernel.ts +96 -0
- package/src/model/support.ts +38 -8
- package/src/model/universal/archs.ts +332 -0
- package/src/model/universal/dense.ts +576 -0
- package/src/model/universal/modules.ts +163 -0
- package/src/model/universal/rope.ts +366 -0
- package/src/pi-provider.ts +1 -1
- package/src/pi-terminal.ts +29 -8
- package/src/pi-web.ts +19 -2
- package/src/prompt-cache.ts +55 -9
- package/src/quantize/config-writer.ts +1 -1
- package/src/registry.ts +219 -8
- package/src/responses.ts +23 -0
- package/src/sampler.ts +220 -9
- package/src/serve/batch-scheduler.ts +307 -83
- package/src/serve/generation-gateway.ts +156 -29
- package/src/server.ts +746 -73
- package/src/spec/drafter.ts +81 -49
- package/src/spec/dspark/data-dflash.ts +156 -0
- package/src/spec/dspark/data.ts +201 -0
- package/src/spec/dspark/generate-dflash.ts +193 -0
- package/src/spec/dspark/generate.ts +290 -0
- package/src/spec/dspark/loss.ts +213 -0
- package/src/spec/dspark/module-dflash.ts +361 -0
- package/src/spec/dspark/module.ts +540 -0
- package/src/spec/dspark/sample.ts +112 -0
- package/src/spec/generate.ts +40 -19
- package/src/ssd-cache.ts +213 -0
- package/src/status-page.html +1 -1
- package/src/tokenizer.ts +7 -0
- package/src/train/flash-cce.ts +81 -23
- package/src/train/fuse.ts +243 -0
- package/src/train/index.ts +3 -0
- package/src/train/job.ts +8 -0
- package/src/train/loss.ts +341 -40
- package/src/train/prefix-shared.ts +62 -11
- package/src/train/segmented.ts +215 -55
- package/src/train/trainer.ts +112 -38
- package/src/web/skills/memory/SKILL.md +40 -17
package/src/train/trainer.ts
CHANGED
|
@@ -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: ${
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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() -
|
|
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 ?
|
|
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
|
|
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
|
-
|
|
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(
|
|
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
|
-
//
|
|
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 ${
|
|
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
|
-
|
|
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: ${
|
|
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
|
-
|
|
1110
|
-
|
|
1111
|
-
|
|
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 —
|
|
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
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
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 —
|
|
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
|