mlx-bun 0.0.4

package/src/train/forward.ts

@@ -0,0 +1,229 @@
+// Training forward pass: full-sequence logits with no KV reuse.
+//
+// B=1 (no padding): trainForward runs the model's stock forwardHidden over a
+// fresh cache exactly as before — bit-identical to the single-example trainer.
+//
+// B>1 (padded batch): every row in the batch is padded to a common length L,
+// but the rows have different *true* lengths. A plain causal mask would let a
+// real query attend to padded keys (and let padded queries pollute the cache
+// view), so for batched forwards we build an explicit [B, 1, L, L] boolean
+// mask that is the causal mask AND-ed with a per-row "key is within this row's
+// valid length" mask, and route it through the model's attention via a cache
+// wrapper whose makeMask returns it. The wrapper delegates every KV operation
+// to a real cache (KVCache/RotatingKVCache are already shape-generic over B),
+// so only the mask differs from the stock path. Window/offset handling is
+// preserved: training always runs offset-0 full-sequence forwards, so the
+// causal-with-window mask is built at offset 0, exactly matching what the
+// underlying cache's makeMask would have produced for the unpadded case.
+
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { createCausalMask, type Cache, type Mask } from "../model/gemma4-base";
+import type { RuntimeModel } from "../model/factory";
+
+/** Stateless cache for the training forward. Training is always a single
+ *  offset-0 full-sequence pass, so `updateAndFetch` is a pure pass-through (no
+ *  buffer, no state mutation) and `makeMask` is the offset-0 causal/windowed
+ *  mask — identical to KVCache at offset 0. The statelessness is REQUIRED for
+ *  gradient checkpointing: mlx_checkpoint re-runs the layer closure during the
+ *  backward recompute, and a stateful (appending) cache would corrupt on the
+ *  second run. */
+export class TrainingCache implements Cache {
+  offset = 0;
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    // Return fresh views: the caller disposes its input k/v right after (its
+    // contract with the real KVCache, which returns buffer slices). A
+    // full-range slice is a cheap view node; mlx keeps the source alive
+    // through it, so disposing the inputs is safe.
+    return [k.slice([0, 0, 0, 0], k.shape), v.slice([0, 0, 0, 0], v.shape)];
+  }
+  makeMask(N: number, windowSize: number | null): Mask {
+    if (N === 1) return { mode: "", arr: null };
+    if (windowSize === null || N <= windowSize) return { mode: "causal", arr: null };
+    return { mode: "array", arr: createCausalMask(N, 0, windowSize) };
+  }
+  state(): MlxArray[] {
+    return [];
+  }
+  isTrimmable(): boolean {
+    return true;
+  }
+  trim(_n: number): void {
+    /* offset is pinned at 0 */
+  }
+  dispose(): void {
+    /* owns no arrays */
+  }
+}
+
+/** Run a full-sequence forward for training.
+ *  @param ids int32 array [B, L].
+ *  @returns logits [B, L, V] (caller owns; dispose when done).
+ *
+ *  When `validLengths` is omitted (or every row is full-length L) the batch
+ *  has no padding and the model's stock mask path is used unchanged — the
+ *  B=1 result is bit-identical to the original trainer. When some row is
+ *  shorter than L, a padding-aware batched mask is injected. */
+export function trainForwardHidden(
+  model: RuntimeModel,
+  ids: MlxArray,
+  validLengths?: number[],
+): MlxArray {
+  const [B, L] = ids.shape as [number, number];
+  const needsPadMask =
+    validLengths !== undefined && validLengths.some((v) => v < L);
+
+  if (needsPadMask) {
+    // B>1 padded path: dense cache wrapped with the per-row padding mask
+    // (unchanged; not yet checkpointed — B=1 is the checkpointing target).
+    const realCache = model.makeCache();
+    const cache = wrapWithBatchedMask(realCache, B, L, validLengths!);
+    try {
+      return model.forwardHidden(ids, cache);
+    } finally {
+      for (const c of realCache) c.dispose();
+      for (const c of cache) (c as BatchedMaskCache).disposeOwnMask();
+    }
+  }
+
+  // B=1 / no padding: stateless offset-0 caches (one per donor layer), so the
+  // forward is pure and safe to recompute under gradient checkpointing.
+  const probe = model.makeCache();
+  const cache: Cache[] = probe.map(() => new TrainingCache());
+  for (const c of probe) c.dispose();
+  try {
+    return model.forwardHidden(ids, cache); // [B, L, hidden]
+  } finally {
+    for (const c of cache) c.dispose();
+  }
+}
+
+/** Full-sequence forward returning logits [B, L, V]. Convenience over
+ *  trainForwardHidden + logitsFromHidden. The SFT B=1 path deliberately skips
+ *  this and applies the LM head only at response positions (responseOnlyCe)
+ *  to avoid materializing the [B, L, vocab] logits — the dominant training
+ *  memory term at long context with a 262k vocab. */
+export function trainForward(
+  model: RuntimeModel,
+  ids: MlxArray,
+  validLengths?: number[],
+): MlxArray {
+  const h = trainForwardHidden(model, ids, validLengths);
+  const logits = model.logitsFromHidden(h);
+  h.dispose();
+  return logits;
+}
+
+/** Build the padding-aware additive-style boolean mask for one window type.
+ *  Shape `[B, 1, L, L]`, true where query i may attend to key j:
+ *    allow[b,0,i,j] = (j <= i)               (causal, offset 0)
+ *                     AND (j < validLen[b])   (key within row b's true length)
+ *                     [AND (i < j + window)   if a sliding window applies]
+ *  Built from createCausalMask(L, 0, window) (the model's own [L, L] bool
+ *  matrix) broadcast over B, masked by a per-row [B,1,1,L] key-validity mask.
+ *  Caller owns the returned array. */
+export function buildBatchedPadMask(
+  B: number,
+  L: number,
+  validLengths: number[],
+  windowSize: number | null,
+): MlxArray {
+  // [L, L] bool causal (+window) matrix at offset 0 — exactly the model's.
+  const causal2d = createCausalMask(L, 0, windowSize); // [L, L] bool
+  const causal4d = ops.reshape(causal2d, [1, 1, L, L]);
+  causal2d.dispose();
+
+  // Per-row key validity: keyValid[b, 0, 0, j] = j < validLengths[b].
+  // Build as int32 {0,1} then compare to get bool, broadcasting over q (i).
+  const keyValidData = new Int32Array(B * L);
+  for (let b = 0; b < B; b++) {
+    const vlen = validLengths[b]!;
+    for (let j = 0; j < L; j++) keyValidData[b * L + j] = j < vlen ? 1 : 0;
+  }
+  const keyValidI = MlxArray.fromInt32(keyValidData, [B, 1, 1, L]);
+  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
+  const keyValid = ops.less(zero, keyValidI); // 0 < x  → bool true where x==1
+  zero.dispose();
+  keyValidI.dispose();
+
+  const allow = ops.logicalAnd(causal4d, keyValid); // broadcasts [1,1,L,L] & [B,1,1,L]
+  causal4d.dispose();
+  keyValid.dispose();
+  return allow; // [B, 1, L, L] bool
+}
+
+/** Cache wrapper that forwards every KV op to a real cache but overrides
+ *  makeMask to return a precomputed padding-aware batched mask. One mask is
+ *  built per distinct window size (null = full attention, else the sliding
+ *  window) and shared across all caches of that type; ownership of the mask
+ *  arrays stays with the wrapper (disposeOwnMask), NOT the model — the model's
+ *  forwardLayers only disposes the mask it pulled from makeMask, so makeMask
+ *  hands back a fresh view each call and keeps the canonical copy itself. */
+class BatchedMaskCache implements Cache {
+  constructor(
+    private readonly inner: Cache,
+    private readonly window: number | null,
+    private readonly maskStore: Map<number | null, MlxArray>,
+  ) {}
+
+  get offset(): number {
+    return this.inner.offset;
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    return this.inner.updateAndFetch(k, v);
+  }
+
+  makeMask(N: number, _windowSize: number | null): Mask {
+    const canonical = this.maskStore.get(this.window)!;
+    // Hand back a fresh full-range view so the model can dispose it without
+    // touching our canonical copy. slice over all 4 axes = cheap view node.
+    const [B, , L] = canonical.shape as [number, number, number, number];
+    const view = canonical.slice([0, 0, 0, 0], [B, 1, L, L]);
+    return { mode: "array", arr: view };
+  }
+
+  state(): MlxArray[] {
+    return this.inner.state();
+  }
+  isTrimmable(): boolean {
+    return this.inner.isTrimmable();
+  }
+  trim(n: number): void {
+    this.inner.trim(n);
+  }
+  dispose(): void {
+    this.inner.dispose();
+  }
+  /** Dispose the canonical mask copies (shared store; idempotent via clear). */
+  disposeOwnMask(): void {
+    for (const m of this.maskStore.values()) m.dispose();
+    this.maskStore.clear();
+  }
+}
+
+/** Wrap each real cache in a BatchedMaskCache, precomputing one padding-aware
+ *  mask per distinct window. Caches without a sliding window get the
+ *  full-attention (null) mask. */
+function wrapWithBatchedMask(
+  realCache: Cache[],
+  B: number,
+  L: number,
+  validLengths: number[],
+): Cache[] {
+  const windows = realCache.map((c) => cacheWindow(c));
+  const distinct = new Set(windows);
+  const store = new Map<number | null, MlxArray>();
+  for (const w of distinct) store.set(w, buildBatchedPadMask(B, L, validLengths, w));
+  return realCache.map((c, i) => new BatchedMaskCache(c, windows[i]!, store));
+}
+
+/** The sliding-window size a cache enforces, or null for full attention.
+ *  RotatingKVCache stores its window as `maxSize`; KVCache has none. Mirrors
+ *  gemma4 forwardLayers, where sliding layers use windowSize == the rotating
+ *  cache's maxSize. (Quantized KV caches are not used on the training path.) */
+function cacheWindow(c: Cache): number | null {
+  const maxSize = (c as { maxSize?: number }).maxSize;
+  return typeof maxSize === "number" ? maxSize : null;
+}

package/src/train/index.ts

@@ -0,0 +1,43 @@
+// Public surface for native LoRA + DPO fine-tuning.
+
+export { ValueAndGrad } from "../mlx/autograd";
+
+export {
+  loadSftDataset, iterateSftBatches, encodeSftRow, sftRowFormat,
+  loadDpoDataset, iterateDpoBatches, encodeDpoRow, tokenizePair,
+  probeFormat, rowLength, dpoChosenLength, dpoRejectedLength, PAD_TO,
+  type SftExample, type SftBatch, type SftRowFormat,
+  type DpoExample, type DpoBatch, type DatasetFormat,
+} from "./dataset";
+
+export { trainForward, buildBatchedPadMask } from "./forward";
+
+export {
+  sftLoss, dpoLoss, dpoRefLogps, dpoMetrics,
+} from "./loss";
+
+export {
+  AdamW, warmupCosineSchedule, type AdamWOptions,
+} from "./optimizer";
+
+export {
+  buildTrainableLora, attachForTraining, detachTraining, setLoraScale,
+  flatA, flatB, flatParams, saveAdapter, disposeLora,
+  type TrainableLora, type TrainableTarget, type SaveAdapterConfig,
+} from "./lora-params";
+
+export {
+  resolveRanks, DEFAULT_TARGET_MODULES,
+  type RankScaling, type ResolveRanksOptions,
+} from "./rank";
+
+export {
+  trainLora, DEFAULT_TRAIN_CONFIG,
+  type TrainConfig, type TrainResult,
+} from "./trainer";
+
+export { mergeAdapters, type MergeStats } from "./merge";
+
+export { exportAdapter, type ExportManifest } from "./export";
+
+export { finetuneRunner, inspectDataset } from "./job";

package/src/train/job.ts

@@ -0,0 +1,148 @@
+// The finetune JobRunner: wires a submitted finetune job to trainLora.
+// Speaks the shared src/jobs/types.ts contract. Registered as kind "finetune"
+// so both submitInProcess and the subprocess job-entry resolver find it
+// (job-entry.ts KIND_MODULES maps finetune → ../train/job.ts).
+
+import type { JobRunner } from "../jobs/types";
+import { registerRunner } from "../jobs/runner";
+import {
+  setWiredLimit, synchronize, maxRecommendedWorkingSetSize,
+} from "../mlx/ffi";
+import { gpuStream } from "../mlx/array";
+import { trainLora, DEFAULT_TRAIN_CONFIG, type TrainConfig } from "./trainer";
+import { probeFormat } from "./dataset";
+
+/** Parse the submit config into a TrainConfig (with model/data paths). */
+interface FinetuneSubmit {
+  model_dir: string;
+  data_dir: string;
+  adapter_path: string;
+  method?: "sft" | "dpo";
+  rank?: number;
+  scale?: number;
+  rank_scaling?: TrainConfig["rankScaling"];
+  target_modules?: string[];
+  num_layers?: number;
+  iters?: number;
+  learning_rate?: number;
+  max_seq_length?: number;
+  batch_size?: number;
+  grad_accumulation_steps?: number;
+  seed?: number;
+  steps_per_report?: number;
+  steps_per_eval?: number;
+  weight_decay?: number;
+  grad_checkpoint?: boolean;
+  segment_size?: number;
+  save_checkpoints?: boolean;
+  dpo_beta?: number;
+  dpo_warmup_iters?: number;
+  dpo_lr_schedule?: "constant" | "cosine";
+}
+
+function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir: string; cfg: TrainConfig } {
+  const c = raw as unknown as FinetuneSubmit;
+  if (!c.model_dir) throw new Error("finetune job: missing model_dir");
+  if (!c.data_dir) throw new Error("finetune job: missing data_dir");
+  if (!c.adapter_path) throw new Error("finetune job: missing adapter_path");
+
+  const cfg: TrainConfig = {
+    ...DEFAULT_TRAIN_CONFIG,
+    method: c.method ?? "sft",
+    rank: c.rank ?? DEFAULT_TRAIN_CONFIG.rank,
+    scale: c.scale ?? DEFAULT_TRAIN_CONFIG.scale,
+    rankScaling: c.rank_scaling ?? DEFAULT_TRAIN_CONFIG.rankScaling,
+    targetModules: c.target_modules ?? DEFAULT_TRAIN_CONFIG.targetModules,
+    numLayers: c.num_layers ?? DEFAULT_TRAIN_CONFIG.numLayers,
+    iters: c.iters ?? DEFAULT_TRAIN_CONFIG.iters,
+    learningRate: c.learning_rate ?? (c.method === "dpo" ? 5e-5 : DEFAULT_TRAIN_CONFIG.learningRate),
+    maxSeqLen: c.max_seq_length ?? DEFAULT_TRAIN_CONFIG.maxSeqLen,
+    batchSize: c.batch_size ?? DEFAULT_TRAIN_CONFIG.batchSize,
+    gradAccumSteps: c.grad_accumulation_steps ?? DEFAULT_TRAIN_CONFIG.gradAccumSteps,
+    seed: c.seed ?? DEFAULT_TRAIN_CONFIG.seed,
+    stepsPerReport: c.steps_per_report ?? DEFAULT_TRAIN_CONFIG.stepsPerReport,
+    stepsPerEval: c.steps_per_eval ?? DEFAULT_TRAIN_CONFIG.stepsPerEval,
+    betas: DEFAULT_TRAIN_CONFIG.betas,
+    weightDecay: c.weight_decay ?? DEFAULT_TRAIN_CONFIG.weightDecay,
+    gradCheckpoint: c.grad_checkpoint ?? DEFAULT_TRAIN_CONFIG.gradCheckpoint,
+    segmentSize: c.segment_size ?? DEFAULT_TRAIN_CONFIG.segmentSize,
+    saveCheckpoints: c.save_checkpoints ?? DEFAULT_TRAIN_CONFIG.saveCheckpoints,
+    dpoBeta: c.dpo_beta ?? DEFAULT_TRAIN_CONFIG.dpoBeta,
+    dpoWarmupIters: c.dpo_warmup_iters ?? DEFAULT_TRAIN_CONFIG.dpoWarmupIters,
+    dpoLrSchedule: c.dpo_lr_schedule ?? DEFAULT_TRAIN_CONFIG.dpoLrSchedule,
+    adapterPath: c.adapter_path,
+    baseModel: c.model_dir,
+  };
+  return { modelDir: c.model_dir, dataDir: c.data_dir, cfg };
+}
+
+export const finetuneRunner: JobRunner = async (emit, config) => {
+  const { modelDir, dataDir, cfg } = parseConfig(config);
+
+  emit({ type: "stage", stage: "load", progress: 0.01, message: `loading model ${modelDir}` });
+
+  const { loadModelConfig } = await import("../config");
+  const { Weights } = await import("../weights");
+  const { createModel } = await import("../model/factory");
+  const { loadTokenizer } = await import("../tokenizer");
+  const { ChatTemplate } = await import("../chat-template");
+
+  const modelConfig = await loadModelConfig(modelDir);
+  const weights = await Weights.open(modelDir);
+  const model = createModel(weights, modelConfig);
+  const tok = await loadTokenizer(modelDir);
+  const tmpl = await ChatTemplate.load(modelDir);
+
+  // Scoped wired limit around the whole run (set → train → synchronize →
+  // restore). Subprocess-isolated so pinning is safe (PLAN Phase 6 rule:
+  // never leave the limit raised process-permanently).
+  const oldLimit = setWiredLimit(maxRecommendedWorkingSetSize());
+  try {
+    const result = await trainLora(model, tok, tmpl, dataDir, cfg, emit);
+    return { outputPath: result.adapterPath };
+  } finally {
+    synchronize(gpuStream);
+    setWiredLimit(oldLimit);
+    weights.dispose();
+  }
+};
+
+/** Inspect a dataset directory before submit: counts + detected format. */
+export async function inspectDataset(dataDir: string): Promise<{
+  ok: boolean;
+  n_train: number;
+  n_valid: number;
+  format: string;
+  error?: string;
+}> {
+  try {
+    const trainPath = `${dataDir}/train.jsonl`;
+    if (!(await Bun.file(trainPath).exists()))
+      return { ok: false, n_train: 0, n_valid: 0, format: "unknown", error: `${trainPath} not found` };
+    const nTrain = await countAndProbe(trainPath);
+    const validPath = `${dataDir}/valid.jsonl`;
+    const nValid = (await Bun.file(validPath).exists()) ? await countAndProbe(validPath) : { n: 0, fmt: "" };
+    return {
+      ok: true,
+      n_train: nTrain.n,
+      n_valid: nValid.n,
+      format: nTrain.fmt,
+    };
+  } catch (e) {
+    return {
+      ok: false, n_train: 0, n_valid: 0, format: "unknown",
+      error: e instanceof Error ? e.message : String(e),
+    };
+  }
+}
+
+async function countAndProbe(path: string): Promise<{ n: number; fmt: string }> {
+  const text = await Bun.file(path).text();
+  const lines = text.split("\n").map((l) => l.trim()).filter(Boolean);
+  let fmt = "unknown";
+  if (lines.length > 0) fmt = probeFormat(JSON.parse(lines[0]!) as Record<string, unknown>);
+  return { n: lines.length, fmt };
+}
+
+// Self-register on import (mirrors the dataset/quantize convention).
+registerRunner("finetune", finetuneRunner);

package/src/train/lora-params.ts

@@ -0,0 +1,194 @@
+// Trainable LoRA parameter management: build A/B leaves, attach them to the
+// model's quantized linears for the forward pass, expose a stable flat order
+// for value_and_grad argnums, and save the adapter in the AdapterManager.mount
+// format (lora_a / lora_b tensor names + optiq_lora_config.json + PEFT
+// adapter_config.json).
+//
+// Init follows mlx-lm LoRALinear.from_base: A ~ uniform(-1/sqrt(in), 1/sqrt(in)),
+// B = zeros, so the LoRA residual starts at exactly zero (the adapted model ==
+// base model at step 0). Trained in f32. (mlx-lm uses uniform, not normal —
+// matching it keeps L1 training faithful to the parent.)
+
+import { mkdirSync } from "node:fs";
+import { ptr } from "bun:ffi";
+import { C, Dtype } from "../mlx/ffi";
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { RuntimeModel } from "../model/factory";
+import type { LoraWeights, QuantizedLinear } from "../model/gemma4-base";
+
+const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
+
+/** One trainable LoRA target: the owning linear + its A/B leaves. */
+export interface TrainableTarget {
+  modulePath: string;
+  linear: QuantizedLinear;
+  lw: LoraWeights;
+}
+
+/** The full set of trainable LoRA leaves for one adapter. */
+export interface TrainableLora {
+  adapterId: string;
+  scale: number;
+  targets: TrainableTarget[];
+}
+
+/** Build trainable A/B leaves for the given ranks (one entry per module path).
+ *  Modules absent from `ranks` are skipped. A is normal-init, B is zeros. */
+export function buildTrainableLora(
+  model: RuntimeModel,
+  ranks: Map<string, number>,
+  scale: number,
+  seed: number,
+): TrainableLora {
+  const linears = model.loraTargets();
+  const key = ops.randomKey(BigInt(seed >>> 0));
+  // One subkey per target for reproducible, independent A inits.
+  const paths = [...ranks.keys()].filter((p) => linears.has(p)).sort();
+  const subkeys = ops.randomSplitNum(key, Math.max(1, paths.length));
+  key.dispose();
+
+  const targets: TrainableTarget[] = [];
+  paths.forEach((modulePath, i) => {
+    const linear = linears.get(modulePath)!;
+    const rank = ranks.get(modulePath)!;
+    const inF = linear.inFeatures;
+    const outF = linear.outFeatures;
+    const subkeyRow = subkeys.slice([i, 0], [i + 1, 2]); // [1,2]
+    const subkey = ops.reshape(subkeyRow, [2]);          // mlx wants a [2] key
+    subkeyRow.dispose();
+    const aScale = 1 / Math.sqrt(inF);
+    const a = ops.randomUniform([inF, rank], Dtype.float32, -aScale, aScale, subkey);
+    subkey.dispose();
+    const b = ops.zeros([rank, outF], Dtype.float32);
+    targets.push({ modulePath, linear, lw: { a, b, scale, rank } });
+  });
+  subkeys.dispose();
+
+  return { adapterId: "train", scale, targets };
+}
+
+/** Attach the trainable leaves to the model so the forward pass picks them up
+ *  via loraState.active. */
+export function attachForTraining(model: RuntimeModel, lora: TrainableLora, adapterId: string): void {
+  lora.adapterId = adapterId;
+  for (const t of lora.targets) {
+    (t.linear.adapters ??= new Map()).set(adapterId, t.lw);
+    t.linear.loraState = model.loraState;
+  }
+  model.loraState.active = [adapterId];
+}
+
+/** Detach the trainable leaves (after training / before save). */
+export function detachTraining(model: RuntimeModel, lora: TrainableLora): void {
+  for (const t of lora.targets) t.linear.adapters?.delete(lora.adapterId);
+  model.loraState.active = model.loraState.active.filter((x) => x !== lora.adapterId);
+}
+
+/** Set the LoRA scale on every trainable leaf (DPO reference forward uses 0). */
+export function setLoraScale(lora: TrainableLora, s: number): void {
+  for (const t of lora.targets) t.lw.scale = s;
+}
+
+/** A leaves in stable (sorted-by-modulePath) order — the order used for
+ *  value_and_grad argnums and AdamW state. */
+export function flatA(lora: TrainableLora): MlxArray[] {
+  return lora.targets.map((t) => t.lw.a);
+}
+
+/** B leaves in the same stable order, appended after all A leaves. */
+export function flatB(lora: TrainableLora): MlxArray[] {
+  return lora.targets.map((t) => t.lw.b);
+}
+
+/** [...A, ...B] — the full differentiable parameter vector. */
+export function flatParams(lora: TrainableLora): MlxArray[] {
+  return [...flatA(lora), ...flatB(lora)];
+}
+
+/** Adapter config written alongside the weights. */
+export interface SaveAdapterConfig {
+  rank: number;
+  scale: number;
+  rankScaling: string;
+  targetModules: string[];
+  numLayers: number;
+  method: string;
+  baseModel: string;
+}
+
+/** Save the trained adapter into `dir` in the AdapterManager.mount format:
+ *  adapters.safetensors (lora_a / lora_b tensor names), optiq_lora_config.json,
+ *  and a PEFT-compatible adapter_config.json. */
+export function saveAdapter(
+  lora: TrainableLora,
+  dir: string,
+  cfg: SaveAdapterConfig,
+  appliedRanks: Record<string, number>,
+): void {
+  mkdirSync(dir, { recursive: true });
+
+  // Build the native string→array map and insert lora_a / lora_b per target.
+  const map = C.mlx_map_string_to_array_new();
+  const meta = C.mlx_map_string_to_string_new();
+  try {
+    for (const t of lora.targets) {
+      // Materialize before save (the leaves are live graph nodes).
+      t.lw.a.eval();
+      t.lw.b.eval();
+      if (C.mlx_map_string_to_array_insert(map, ptr(cstr(`${t.modulePath}.lora_a`)), t.lw.a.handle) !== 0)
+        throw new Error(`map insert ${t.modulePath}.lora_a failed`);
+      if (C.mlx_map_string_to_array_insert(map, ptr(cstr(`${t.modulePath}.lora_b`)), t.lw.b.handle) !== 0)
+        throw new Error(`map insert ${t.modulePath}.lora_b failed`);
+    }
+    const file = `${dir}/adapters.safetensors`;
+    if (C.mlx_save_safetensors(ptr(cstr(file)), map, meta) !== 0)
+      throw new Error(`mlx_save_safetensors(${file}) failed`);
+  } finally {
+    C.mlx_map_string_to_array_free(map);
+    C.mlx_map_string_to_string_free(meta);
+  }
+
+  // optiq_lora_config.json — AdapterManager.readAdapterScale reads
+  // lora_parameters.scale + rank from here first.
+  const optiqCfg = {
+    rank: cfg.rank,
+    scale: cfg.scale,
+    rank_scaling: cfg.rankScaling,
+    method: cfg.method,
+    target_modules: cfg.targetModules,
+    num_layers: cfg.numLayers,
+    lora_parameters: { scale: cfg.scale, rank: cfg.rank },
+    applied_ranks: appliedRanks,
+    source_model: cfg.baseModel,
+  };
+  Bun.write(`${dir}/optiq_lora_config.json`, JSON.stringify(optiqCfg, null, 2) + "\n");
+
+  // PEFT-compatible adapter_config.json (loadable by mlx-lm / PEFT).
+  const peftCfg = {
+    fine_tune_type: "lora",
+    num_layers: cfg.numLayers,
+    lora_parameters: { rank: cfg.rank, scale: cfg.scale, dropout: 0.0, keys: null },
+    base_model_name_or_path: cfg.baseModel,
+    bias: "none",
+    fan_in_fan_out: false,
+    inference_mode: false,
+    init_lora_weights: true,
+    lora_alpha: Math.round(cfg.rank * cfg.scale),
+    lora_dropout: 0.0,
+    peft_type: "LORA",
+    r: cfg.rank,
+    target_modules: cfg.targetModules,
+    task_type: "CAUSAL_LM",
+    optiq: { rank_scaling: cfg.rankScaling, applied_ranks: appliedRanks },
+  };
+  Bun.write(`${dir}/adapter_config.json`, JSON.stringify(peftCfg, null, 2) + "\n");
+}
+
+/** Dispose all trainable leaves (frees A/B). */
+export function disposeLora(lora: TrainableLora): void {
+  for (const t of lora.targets) {
+    t.lw.a.dispose();
+    t.lw.b.dispose();
+  }
+}