mlx-bun 0.0.4

package/src/train/trainer.ts

@@ -0,0 +1,600 @@
+// LoRA trainer: SFT (masked cross-entropy) and DPO loops over a single
+// quantized base, differentiating only the LoRA A/B leaves.
+//
+// The value_and_grad closure differentiates [...flatA, ...flatB]. Inside the
+// closure we swap the incoming primal handles into the LoraWeights so the
+// model's forward builds its graph against the differentiated leaves; after
+// AdamW.step replaces a leaf, the writeback callback pushes the new handle
+// back into the LoraWeights so the next forward (and the next closure swap)
+// sees it. Base quantized weights are never in argnums — they stay frozen.
+
+import { MlxArray } from "../mlx/array";
+import { clearCache, peakMemory, activeMemory, cacheMemory, resetPeakMemory } from "../mlx/ffi";
+
+const MEM_LOG = process.env.MLX_BUN_MEM_LOG === "1";
+const gbStr = (b: number) => `${(b / 1e9).toFixed(2)}GB`;
+import { ValueAndGrad } from "../mlx/autograd";
+import * as ops from "../mlx/ops";
+import { evalAll } from "../mlx/ops";
+import type { RuntimeModel } from "../model/factory";
+import { Gemma4Model, type GradCheckpointCtx } from "../model/gemma4";
+import { MiniCPM5Model } from "../model/minicpm5";
+import { setTrainingAttn } from "../model/flash-attention";
+import { SegmentedBackward, SegmentedBackwardGemma4, planSegmentsBySize } from "./segmented";
+import type { LoadedTokenizer } from "../tokenizer";
+import type { ChatTemplate } from "../chat-template";
+import type { Emit } from "../jobs/types";
+import {
+  loadSftDataset, iterateSftBatches, loadDpoDataset, iterateDpoBatches,
+  rowLength, type SftBatch, type DpoBatch,
+} from "./dataset";
+import {
+  buildTrainableLora, attachForTraining, detachTraining, flatParams,
+  saveAdapter, disposeLora, type TrainableLora, type SaveAdapterConfig,
+} from "./lora-params";
+import { resolveRanks, DEFAULT_TARGET_MODULES, type RankScaling } from "./rank";
+import { sftLoss, dpoLoss, dpoRefLogps, dpoMetrics } from "./loss";
+import { AdamW, warmupCosineSchedule } from "./optimizer";
+import { writeFileSync } from "node:fs";
+
+export interface TrainConfig {
+  method: "sft" | "dpo";
+  rank: number;
+  scale: number;
+  rankScaling: RankScaling;
+  targetModules: string[];
+  numLayers: number;
+  iters: number;
+  learningRate: number;
+  maxSeqLen: number;
+  batchSize: number;
+  gradAccumSteps: number;
+  seed: number;
+  stepsPerReport: number;
+  stepsPerEval: number;
+  betas: [number, number];
+  weightDecay: number;
+  /** Gradient checkpointing: recompute layer activations in the backward pass
+   *  to bound memory at long context (mlx-lm grad_checkpoint; default in optiq).
+   *  Numerically identical to off — pure memory↔compute trade. */
+  gradCheckpoint: boolean;
+  /** Segmented backward: layers per segment (0 = off). Streams the SFT backward
+   *  segment-by-segment so only one segment's activations are live at a time —
+   *  beats per-layer checkpointing at long context, where naive checkpointing
+   *  holds every layer's recompute activations at once and crashes. B=1 SFT
+   *  only (the responseOnlyCe path); mutually exclusive with gradCheckpoint.
+   *  Numerically identical to off. See docs/design/segmented-backward-training.md. */
+  segmentSize: number;
+  /** Keep ALL eval-step checkpoints (full mountable adapters under
+   *  checkpoints/step-<NNNNN>-val<loss>/) + write metrics.json. Off by default
+   *  so server finetune jobs and tests stay clean; the chunk-finetune script
+   *  turns it on so the best-ON-TASK checkpoint can be chosen by the real eval
+   *  afterward (val is only a proxy — we keep them all and let the eval pick). */
+  saveCheckpoints: boolean;
+  // DPO
+  dpoBeta: number;
+  dpoWarmupIters: number;
+  dpoLrSchedule: "constant" | "cosine";
+  // output
+  adapterPath: string;
+  baseModel: string;
+}
+
+export interface TrainResult {
+  adapterPath: string;
+  appliedRanks: Record<string, number>;
+  numIters: number;
+}
+
+export const DEFAULT_TRAIN_CONFIG: TrainConfig = {
+  method: "sft",
+  rank: 8,
+  scale: 1.0,
+  rankScaling: "by_bits",
+  targetModules: [...DEFAULT_TARGET_MODULES],
+  numLayers: -1,
+  iters: 100,
+  learningRate: 2e-4,
+  maxSeqLen: 512,
+  batchSize: 1,
+  gradAccumSteps: 1,
+  seed: 0,
+  stepsPerReport: 10,
+  stepsPerEval: 50,
+  betas: [0.9, 0.999],
+  weightDecay: 0.01,
+  gradCheckpoint: false, // default off until validated; flip to true (optiq default) after
+  segmentSize: 0, // default off; >0 enables segmented backward (MiniCPM5 SFT B=1)
+  saveCheckpoints: false, // off by default; scripts opt in (keep-all-checkpoints + metrics.json)
+
+  dpoBeta: 0.1,
+  dpoWarmupIters: 0,
+  dpoLrSchedule: "cosine",
+  adapterPath: "adapters",
+  baseModel: "",
+};
+
+/** Run a LoRA fine-tune. Emits per-step train metrics, periodic val metrics,
+ *  and a final stage:done with the adapter path + applied ranks. Saves the
+ *  adapter (last + best-on-val) and returns where it landed. */
+export async function trainLora(
+  model: RuntimeModel,
+  tok: LoadedTokenizer,
+  tmpl: ChatTemplate,
+  dataDir: string,
+  cfg: TrainConfig,
+  emit: Emit,
+): Promise<TrainResult> {
+  emit({ type: "stage", stage: "setup", progress: 0.02, message: "resolving ranks" });
+
+  const ranks = resolveRanks(model, {
+    rank: cfg.rank,
+    rankScaling: cfg.rankScaling,
+    targetModules: cfg.targetModules,
+    numLayers: cfg.numLayers,
+  });
+  const appliedRanks: Record<string, number> = Object.fromEntries(ranks);
+
+  const lora = buildTrainableLora(model, ranks, cfg.scale, cfg.seed);
+  attachForTraining(model, lora, "train");
+
+  const saveCfg: SaveAdapterConfig = {
+    rank: cfg.rank,
+    scale: cfg.scale,
+    rankScaling: cfg.rankScaling,
+    targetModules: cfg.targetModules,
+    numLayers: cfg.numLayers,
+    method: cfg.method,
+    baseModel: cfg.baseModel,
+  };
+
+  // Keep-ALL-checkpoints + a durable metrics record. At each eval step we
+  // snapshot the adapter to checkpoints/step-<NNNNN>-val<loss>/ — a full,
+  // mountable adapter — so the best-ON-TASK checkpoint can be chosen by the
+  // real eval afterward (val is only a proxy for task quality, so we do NOT
+  // pick for the user). The loss/val trajectory is written to metrics.json so
+  // runs are comparable and survive (the tee'd stdout log is ephemeral). We
+  // wrap `emit` rather than touching the loops: it already streams train/val
+  // metrics, and the live LoRA leaves are the current step's weights at emit
+  // time, so saving here captures exactly that checkpoint.
+  const trainHistory: { step: number; loss: number }[] = [];
+  const valHistory: { step: number; loss: number; checkpoint: string | null }[] = [];
+  const startedAt = Date.now();
+  const collect: Emit = (e) => {
+    if (e.type === "metric" && e.kind === "train") {
+      trainHistory.push({ step: e.step, loss: e.loss });
+    } else if (e.type === "metric" && e.kind === "val") {
+      let checkpoint: string | null = null;
+      if (cfg.saveCheckpoints) {
+        const tag = `step-${String(e.step).padStart(5, "0")}-val${e.loss.toFixed(4)}`;
+        checkpoint = `${cfg.adapterPath}/checkpoints/${tag}`;
+        saveAdapter(lora, checkpoint, saveCfg, appliedRanks);
+        emit({ type: "stage", stage: "checkpoint", progress: e.step / cfg.iters,
+               message: `checkpoint ${tag}` });
+      }
+      valHistory.push({ step: e.step, loss: e.loss, checkpoint });
+    }
+    emit(e);
+  };
+
+  try {
+    const result =
+      cfg.method === "dpo"
+        ? await dpoLoop(model, tok, tmpl, dataDir, cfg, lora, collect)
+        : await sftLoop(model, tok, tmpl, dataDir, cfg, lora, collect);
+
+    // Final save (last adapter).
+    detachTraining(model, lora);
+    saveAdapter(lora, cfg.adapterPath, saveCfg, appliedRanks);
+
+    // Durable, structured run record alongside the adapter (only when we kept
+    // checkpoints). Pick the best-val one for convenience, but ship ALL of them
+    // — the real eval decides; val is only a proxy.
+    if (cfg.saveCheckpoints) {
+      let bestVal: { step: number; loss: number; checkpoint: string | null } | null = null;
+      for (const v of valHistory) if (bestVal === null || v.loss < bestVal.loss) bestVal = v;
+      writeFileSync(`${cfg.adapterPath}/metrics.json`, JSON.stringify({
+        config: {
+          method: cfg.method, rank: cfg.rank, scale: cfg.scale,
+          learningRate: cfg.learningRate, maxSeqLen: cfg.maxSeqLen,
+          segmentSize: cfg.segmentSize, iters: cfg.iters,
+          numLayers: cfg.numLayers, baseModel: cfg.baseModel,
+        },
+        wallSeconds: (Date.now() - startedAt) / 1000,
+        peakGb: peakMemory() / 1e9,
+        finalTrainLoss: trainHistory.at(-1)?.loss ?? null,
+        finalValLoss: valHistory.at(-1)?.loss ?? null,
+        bestVal,
+        valTrajectory: valHistory,
+        trainLosses: trainHistory,
+      }, null, 2));
+    }
+
+    emit({
+      type: "stage",
+      stage: "done",
+      progress: 1,
+      message: `adapter saved to ${cfg.adapterPath}`,
+      adapter_path: cfg.adapterPath,
+      applied_ranks: appliedRanks,
+    });
+    return { adapterPath: cfg.adapterPath, appliedRanks, numIters: result.numIters };
+  } finally {
+    disposeLora(lora);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// SFT loop
+// ---------------------------------------------------------------------------
+
+async function sftLoop(
+  model: RuntimeModel,
+  tok: LoadedTokenizer,
+  tmpl: ChatTemplate,
+  dataDir: string,
+  cfg: TrainConfig,
+  lora: TrainableLora,
+  emit: Emit,
+): Promise<{ numIters: number }> {
+  emit({ type: "stage", stage: "data", progress: 0.05, message: "loading SFT dataset" });
+  const train = await loadSftDataset(`${dataDir}/train.jsonl`, tok, tmpl);
+  const valid = (await fileExists(`${dataDir}/valid.jsonl`))
+    ? await loadSftDataset(`${dataDir}/valid.jsonl`, tok, tmpl)
+    : [];
+
+  const params = flatParams(lora);
+  const opt = new AdamW(
+    params,
+    { lr: cfg.learningRate, betas: cfg.betas, eps: 1e-8, weightDecay: cfg.weightDecay },
+    (i, p) => writeParam(lora, i, p),
+  );
+
+  // Segmented backward: stream the SFT backward segment-by-segment (only one
+  // segment's activations live at a time). Phase A is MiniCPM5 SFT B=1 only;
+  // it replaces the single value_and_grad below and is mutually exclusive with
+  // gradient checkpointing (see docs/design/segmented-backward-training.md).
+  const useSegmented = cfg.segmentSize > 0;
+  let segmented: SegmentedBackward | SegmentedBackwardGemma4 | null = null;
+  if (useSegmented) {
+    if (model instanceof MiniCPM5Model)
+      segmented = new SegmentedBackward(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
+    else if (model instanceof Gemma4Model)
+      segmented = new SegmentedBackwardGemma4(model, lora, planSegmentsBySize(model.layers.length, cfg.segmentSize));
+    else
+      throw new Error("segmented backward (segmentSize > 0) is only wired for MiniCPM5 and Gemma4");
+    emit({ type: "stage", stage: "setup", progress: 0.04,
+      message: `segmented backward: ${Math.ceil(model.layers.length / cfg.segmentSize)} segments of <=${cfg.segmentSize} layers` });
+  }
+
+  // The value_and_grad loss closure: temporarily swap the differentiated
+  // primals into the LoraWeights so the forward graph differentiates them,
+  // build the SFT loss against the active batch, then restore the original
+  // leaf wrappers (the primal wrappers get disposed by the autograd closure;
+  // the graph already captured their handles). Not built on the segmented path
+  // (segmentedSftGrads runs its own per-segment value_and_grads).
+  let currentBatch: SftBatch | null = null;
+  const vag = useSegmented
+    ? null
+    : new ValueAndGrad((primals) => {
+        const saved = swapPrimals(lora, primals);
+        try {
+          return sftLoss(model, currentBatch!);
+        } finally {
+          restorePrimals(lora, saved);
+        }
+      }, params.map((_, i) => i));
+
+  // Gradient checkpointing: group each layer's LoRA weights so forwardLayers
+  // can thread them as explicit checkpoint inputs (see Gemma4Model.gradCkpt).
+  // Checkpoints are created per forward and disposed after each step's backward.
+  let ckptCtx: GradCheckpointCtx | null = null;
+  if (!useSegmented && cfg.gradCheckpoint && model instanceof Gemma4Model) {
+    const byLayer: GradCheckpointCtx["byLayer"] = new Map();
+    for (const t of lora.targets) {
+      const m = t.modulePath.match(/\.layers\.(\d+)\./);
+      if (!m) continue;
+      const li = Number(m[1]);
+      (byLayer.get(li) ?? byLayer.set(li, []).get(li)!).push(t.lw);
+    }
+    ckptCtx = { byLayer, keepAlive: [] };
+    model.gradCkpt = ckptCtx;
+  }
+
+  // Training attention. Default = mlx fused SDPA (ops.sdpa): verified correct
+  // dQ/dK/dV vs a hand-rolled reference (scripts/sdpa-vs-manual.ts: 0.00%) and
+  // fast (one fused kernel). flash is opt-in via MLX_BUN_TRAIN_ATTN=flash for an
+  // O(L) memory backward at very long context, BUT its dK kernel currently has
+  // a bug (~100% off) and is ~30x slower here — so it is NOT the default.
+  // Cleared in finally.
+  if (process.env.MLX_BUN_TRAIN_ATTN === "flash") setTrainingAttn("flash");
+
+  // Pad with the eos id when available (a real sentinel token); padded
+  // positions are excluded from both the loss and attention masks, so the
+  // exact pad value never affects the result — it just must be a valid id.
+  const padId = tok.eosTokenId ?? 0;
+  const batches = iterateSftBatches(train, cfg.batchSize, cfg.maxSeqLen, cfg.seed, true, padId);
+  const t0 = Date.now();
+
+  try {
+    for (let step = 1; step <= cfg.iters; step++) {
+      currentBatch = batches.next().value as SftBatch;
+      const tokensThisStep = countResponseTokens(currentBatch);
+
+      if (MEM_LOG) resetPeakMemory();
+      const { value, grads } = segmented
+        ? segmented.step(currentBatch!)
+        : vag!.apply(flatParams(lora));
+      // Evaluate the loss AND grads in ONE mlx_eval so the backward shares the
+      // forward pass. Evaluating `value` on its own first (e.g. via toFloat32)
+      // materializes then frees the forward, forcing the backward to recompute
+      // and hold ALL layer activations — which silently defeats gradient
+      // checkpointing (verified: separate-eval ~0% memory savings, combined-eval
+      // restores it). Keep this single combined eval before reading either.
+      evalAll([value, ...grads]);
+      const lossVal = value.toFloat32()[0]!;
+      value.dispose();
+
+      const gradNorm = globalNorm(grads); // grads already materialized above
+      opt.step(grads);
+      opt.evalState();
+      if (MEM_LOG) {
+        const pk = peakMemory(), act = activeMemory(), ch = cacheMemory();
+        clearCache();
+        console.log(`  [mem] step ${step}: PEAK(live)=${gbStr(pk)} active=${gbStr(act)} cache=${gbStr(ch)} -> after clearCache active=${gbStr(activeMemory())}`);
+      } else {
+        clearCache();
+      }
+      if (ckptCtx) {
+        // Grads are materialized (globalNorm/evalState), so the recompute is
+        // done — free this step's checkpoint closures before the next forward.
+        for (const ck of ckptCtx.keepAlive) ck.dispose();
+        ckptCtx.keepAlive.length = 0;
+      }
+
+      if (step % cfg.stepsPerReport === 0 || step === 1) {
+        const elapsed = (Date.now() - t0) / 1000;
+        emit({
+          type: "metric",
+          kind: "train",
+          step,
+          loss: lossVal,
+          grad_norm: gradNorm,
+          learning_rate: opt.lr,
+          tokens_per_sec: elapsed > 0 ? (step * tokensThisStep) / elapsed : 0,
+          progress: step / cfg.iters,
+          message: `step ${step}/${cfg.iters} loss=${lossVal.toFixed(4)}`,
+        });
+      }
+
+      if (valid.length > 0 && step % cfg.stepsPerEval === 0) {
+        // Eval is forward-only (no backward), so disable checkpointing for it —
+        // it adds no benefit and a forward-only pass fits without it.
+        if (ckptCtx) (model as Gemma4Model).gradCkpt = null;
+        const vLoss = evalSftLoss(model, valid, cfg);
+        if (ckptCtx) (model as Gemma4Model).gradCkpt = ckptCtx;
+        emit({ type: "metric", kind: "val", step, loss: vLoss, progress: step / cfg.iters });
+      }
+    }
+  } finally {
+    setTrainingAttn(null);
+    vag?.dispose();
+    segmented?.dispose();
+    opt.dispose();
+    if (ckptCtx) {
+      (model as Gemma4Model).gradCkpt = null;
+      for (const ck of ckptCtx.keepAlive) ck.dispose();
+    }
+  }
+  return { numIters: cfg.iters };
+}
+
+function evalSftLoss(model: RuntimeModel, valid: ReturnType<typeof Array.prototype.slice>, cfg: TrainConfig): number {
+  let total = 0;
+  let n = 0;
+  for (const ex of valid as { ids: number[]; promptLen: number }[]) {
+    let ids = ex.ids;
+    let promptLen = ex.promptLen;
+    if (ids.length < 2) continue;
+    if (ids.length > cfg.maxSeqLen) {
+      ids = ids.slice(0, cfg.maxSeqLen);
+      promptLen = Math.min(promptLen, Math.max(0, ids.length - 1));
+    }
+    const batch: SftBatch = { ids: [ids], promptLens: [promptLen] };
+    try {
+      const loss = sftLoss(model, batch);
+      total += loss.toFloat32()[0]!;
+      loss.dispose();
+      n++;
+    } catch {
+      // skip examples with no response tokens
+    }
+    clearCache();
+  }
+  return n > 0 ? total / n : 0;
+}
+
+// ---------------------------------------------------------------------------
+// DPO loop
+// ---------------------------------------------------------------------------
+
+async function dpoLoop(
+  model: RuntimeModel,
+  tok: LoadedTokenizer,
+  tmpl: ChatTemplate,
+  dataDir: string,
+  cfg: TrainConfig,
+  lora: TrainableLora,
+  emit: Emit,
+): Promise<{ numIters: number }> {
+  emit({ type: "stage", stage: "data", progress: 0.05, message: "loading DPO dataset" });
+  const train = await loadDpoDataset(`${dataDir}/train.jsonl`, tok, tmpl, cfg.maxSeqLen);
+  const valid = (await fileExists(`${dataDir}/valid.jsonl`))
+    ? await loadDpoDataset(`${dataDir}/valid.jsonl`, tok, tmpl, cfg.maxSeqLen)
+    : [];
+
+  const params = flatParams(lora);
+  const opt = new AdamW(
+    params,
+    { lr: cfg.learningRate, betas: cfg.betas, eps: 1e-8, weightDecay: cfg.weightDecay },
+    (i, p) => writeParam(lora, i, p),
+  );
+
+  const schedule =
+    cfg.dpoLrSchedule === "cosine" || cfg.dpoWarmupIters > 0
+      ? warmupCosineSchedule(cfg.learningRate, cfg.dpoWarmupIters, cfg.iters)
+      : null;
+
+  // Reference log-probs are computed each step OUTSIDE value_and_grad (so the
+  // reference activations free before the policy forwards), then read as
+  // detached constants inside the closure.
+  let currentBatch: DpoBatch | null = null;
+  let refChosen: MlxArray | null = null;
+  let refRejected: MlxArray | null = null;
+
+  const vag = new ValueAndGrad((primals) => {
+    const saved = swapPrimals(lora, primals);
+    try {
+      return dpoLoss(model, currentBatch!, cfg.dpoBeta, refChosen!, refRejected!);
+    } finally {
+      restorePrimals(lora, saved);
+    }
+  }, params.map((_, i) => i));
+
+  const padId = tok.eosTokenId ?? 0;
+  const batches = iterateDpoBatches(train, cfg.batchSize, cfg.seed, true, padId);
+  const t0 = Date.now();
+
+  try {
+    for (let step = 1; step <= cfg.iters; step++) {
+      currentBatch = batches.next().value as DpoBatch;
+      if (schedule) opt.lr = schedule(step);
+
+      const refs = dpoRefLogps(model, lora, currentBatch);
+      refChosen = refs.refChosen;
+      refRejected = refs.refRejected;
+
+      const { value, grads } = vag.apply(flatParams(lora));
+      const lossVal = value.toFloat32()[0]!;
+      value.dispose();
+
+      const gradNorm = globalNorm(grads);
+      opt.step(grads);
+      opt.evalState();
+
+      refChosen.dispose();
+      refRejected.dispose();
+      refChosen = refRejected = null;
+      clearCache();
+
+      if (step % cfg.stepsPerReport === 0 || step === 1) {
+        const m = dpoMetrics(model, lora, currentBatch, cfg.dpoBeta);
+        const elapsed = (Date.now() - t0) / 1000;
+        emit({
+          type: "metric",
+          kind: "train",
+          step,
+          loss: lossVal,
+          grad_norm: gradNorm,
+          learning_rate: opt.lr,
+          accuracy: m.accuracy,
+          margin: m.margin,
+          tokens_per_sec: 0,
+          progress: step / cfg.iters,
+          message: `step ${step}/${cfg.iters} loss=${lossVal.toFixed(4)} acc=${m.accuracy.toFixed(2)} margin=${m.margin.toFixed(3)}`,
+        });
+        clearCache();
+      }
+
+      if (valid.length > 0 && step % cfg.stepsPerEval === 0) {
+        let vl = 0, va = 0, vm = 0, vn = 0;
+        for (const ex of valid) {
+          const vb: DpoBatch = {
+            chosenIds: [ex.chosenIds], rejectedIds: [ex.rejectedIds],
+            chosenMask: [ex.chosenMask], rejectedMask: [ex.rejectedMask],
+          };
+          const m = dpoMetrics(model, lora, vb, cfg.dpoBeta);
+          vl += m.loss; va += m.accuracy; vm += m.margin; vn++;
+          clearCache();
+        }
+        emit({
+          type: "metric", kind: "val", step,
+          loss: vn ? vl / vn : 0,
+          accuracy: vn ? va / vn : 0,
+          margin: vn ? vm / vn : 0,
+          progress: step / cfg.iters,
+        });
+      }
+    }
+  } finally {
+    vag.dispose();
+    opt.dispose();
+    refChosen?.dispose();
+    refRejected?.dispose();
+  }
+  return { numIters: cfg.iters };
+}
+
+// ---------------------------------------------------------------------------
+// helpers
+// ---------------------------------------------------------------------------
+
+/** Swap the incoming value_and_grad primals into the LoraWeights so the
+ *  forward graph differentiates them. Order: [...A, ...B] over sorted targets.
+ *  Returns the displaced original leaves so the caller can restore them (the
+ *  primal wrappers are disposed by the autograd closure after this returns). */
+function swapPrimals(lora: TrainableLora, primals: MlxArray[]): MlxArray[] {
+  const n = lora.targets.length;
+  const saved: MlxArray[] = [];
+  for (let i = 0; i < n; i++) {
+    saved.push(lora.targets[i]!.lw.a, lora.targets[i]!.lw.b);
+    lora.targets[i]!.lw.a = primals[i]!;
+    lora.targets[i]!.lw.b = primals[n + i]!;
+  }
+  return saved;
+}
+
+/** Restore the original leaves displaced by swapPrimals (saved order is
+ *  [a0, b0, a1, b1, ...]). */
+function restorePrimals(lora: TrainableLora, saved: MlxArray[]): void {
+  for (let i = 0; i < lora.targets.length; i++) {
+    lora.targets[i]!.lw.a = saved[2 * i]!;
+    lora.targets[i]!.lw.b = saved[2 * i + 1]!;
+  }
+}
+
+/** Write a replaced param (AdamW output) back into its LoraWeights leaf. */
+function writeParam(lora: TrainableLora, i: number, p: MlxArray): void {
+  const n = lora.targets.length;
+  if (i < n) lora.targets[i]!.lw.a = p;
+  else lora.targets[i - n]!.lw.b = p;
+}
+
+/** L2 norm over all grads (also forces eval). */
+function globalNorm(grads: MlxArray[]): number {
+  const partial: MlxArray[] = grads.map((g) => {
+    const sq = ops.square(g);
+    const flat = ops.reshape(sq, [g.size]);
+    const sum = ops.sumAxis(flat, 0, false);
+    sq.dispose();
+    flat.dispose();
+    return sum;
+  });
+  evalAll(partial);
+  let sumSq = 0;
+  for (const p of partial) {
+    sumSq += p.toFloat32()[0]!;
+    p.dispose();
+  }
+  return Math.sqrt(sumSq);
+}
+
+function countResponseTokens(batch: SftBatch): number {
+  const promptLen = batch.promptLens[0]!;
+  const len = batch.ids[0]!.length;
+  return Math.max(0, len - Math.max(promptLen, 1));
+}
+
+async function fileExists(path: string): Promise<boolean> {
+  return Bun.file(path).exists();
+}