mlx-bun 0.4.0 → 1.0.0

package/src/train/job.ts

@@ -1,205 +0,0 @@
-// 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" | "orpo";
-  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;
-  lora_dropout?: number;
-  rs_lora?: boolean;
-  lora_plus_ratio?: number;
-  grad_checkpoint?: boolean;
-  mlp_split?: boolean;
-  segment_size?: number;
-  save_checkpoints?: boolean;
-  grad_clip_norm?: number;
-  val_max_examples?: number;
-  dpo_beta?: number;
-  dpo_warmup_iters?: number;
-  dpo_lr_schedule?: "constant" | "cosine";
-  orpo_lambda?: number;
-  orpo_warmup_iters?: number;
-  orpo_lr_schedule?: "constant" | "cosine";
-  orpo_chunk_size?: number;
-  orpo_fused_ce?: boolean;
-  orpo_flash_ce?: boolean;
-  orpo_prefix_shared?: boolean;
-  sft_scope?: "full" | "response";
-  warm_start_adapter?: string;
-}
-
-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");
-
-  // For ORPO the full L3 stack is the validated production recipe (what the
-  // `mlx-bun train` CLI turns on): flash-CCE head + prefix-sharing + segmented
-  // backward + token-chunked head. DEFAULT_TRAIN_CONFIG keeps them OFF (they're
-  // opt-in for SFT), so default them ON here when method=orpo — an explicit
-  // value from the caller (CLI/API) still wins via ??. finetuneRunner drops the
-  // quantized-head pieces if the base turns out to be unquantized.
-  const isOrpo = (c.method ?? "sft") === "orpo";
-
-  // ORPO L_SFT scope: "full" (default; paper/TRL-faithful) | "response"
-  // (pre-2026-07 behavior). Validate here so a typo fails at submit, not mid-run.
-  const sftScope = c.sft_scope ?? DEFAULT_TRAIN_CONFIG.sftScope;
-  if (sftScope !== "full" && sftScope !== "response")
-    throw new Error(`finetune job: sft_scope must be "full" or "response" (got ${JSON.stringify(c.sft_scope)})`);
-
-  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 : c.method === "orpo" ? 1e-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,
-    loraDropout: c.lora_dropout ?? DEFAULT_TRAIN_CONFIG.loraDropout,
-    rsLora: c.rs_lora ?? DEFAULT_TRAIN_CONFIG.rsLora,
-    loraPlusRatio: c.lora_plus_ratio ?? DEFAULT_TRAIN_CONFIG.loraPlusRatio,
-    gradCheckpoint: c.grad_checkpoint ?? DEFAULT_TRAIN_CONFIG.gradCheckpoint,
-    mlpSplit: c.mlp_split ?? DEFAULT_TRAIN_CONFIG.mlpSplit,
-    segmentSize: c.segment_size ?? (isOrpo ? 2 : DEFAULT_TRAIN_CONFIG.segmentSize),
-    saveCheckpoints: c.save_checkpoints ?? DEFAULT_TRAIN_CONFIG.saveCheckpoints,
-    gradClipNorm: c.grad_clip_norm ?? DEFAULT_TRAIN_CONFIG.gradClipNorm,
-    valMaxExamples: c.val_max_examples ?? DEFAULT_TRAIN_CONFIG.valMaxExamples,
-    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,
-    orpoLambda: c.orpo_lambda ?? DEFAULT_TRAIN_CONFIG.orpoLambda,
-    orpoWarmupIters: c.orpo_warmup_iters ?? DEFAULT_TRAIN_CONFIG.orpoWarmupIters,
-    orpoLrSchedule: c.orpo_lr_schedule ?? DEFAULT_TRAIN_CONFIG.orpoLrSchedule,
-    orpoChunkSize: c.orpo_chunk_size ?? (isOrpo ? 512 : DEFAULT_TRAIN_CONFIG.orpoChunkSize),
-    orpoFusedCe: c.orpo_fused_ce ?? DEFAULT_TRAIN_CONFIG.orpoFusedCe,
-    orpoFlashCe: c.orpo_flash_ce ?? (isOrpo ? true : DEFAULT_TRAIN_CONFIG.orpoFlashCe),
-    orpoPrefixShared: c.orpo_prefix_shared ?? (isOrpo ? true : DEFAULT_TRAIN_CONFIG.orpoPrefixShared),
-    sftScope,
-    warmStartAdapter: c.warm_start_adapter ?? DEFAULT_TRAIN_CONFIG.warmStartAdapter,
-    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);
-
-  // The flash-CCE / fused-CE heads dequantize a QUANTIZED head in-kernel; on an
-  // unquantized (bf16) base they can't run, so fall back to the full-logits head.
-  // (Quantize the base first to get the large-vocab memory win — see the
-  // Quantize tab.)
-  if (cfg.method === "orpo" && !modelConfig.quantization && (cfg.orpoFlashCe || cfg.orpoFusedCe)) {
-    emit({
-      type: "stage", stage: "load",
-      message: "unquantized base — flash/fused-CE head needs a quantized head; using full-logits head",
-    });
-    cfg.orpoFlashCe = false;
-    cfg.orpoFusedCe = false;
-  }
-  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

@@ -1,254 +0,0 @@
-// 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";
-import { loadAdapterTensors } from "../lora";
-
-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,
-  rsLora = false,
-): 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);
-    // rsLoRA: α/√rank so per-layer rank scaling changes capacity, not step size.
-    const lwScale = rsLora ? scale / Math.sqrt(rank) : scale;
-    targets.push({ modulePath, linear, lw: { a, b, scale: lwScale, rank } });
-  });
-  subkeys.dispose();
-
-  return { adapterId: "train", scale, targets };
-}
-
-/** Warm-start: overwrite a freshly-built LoRA's A/B leaves with the weights saved in
- *  `dir/adapters.safetensors` (saveAdapter's `${modulePath}.lora_a/b` names), so a run
- *  continues from a checkpoint's WEIGHTS. The optimizer state and LR schedule restart —
- *  Adam re-warms its moments in a few steps — so this is a weights warm-start, not a
- *  bit-perfect resume. The checkpoint's rank/targets must match this LoRA (shapes are
- *  checked). MUST be called BEFORE the optimizer is built so it tracks the loaded
- *  handles. Returns the number of targets loaded. */
-export function warmStartFromAdapter(lora: TrainableLora, dir: string): number {
-  const file = `${dir}/adapters.safetensors`;
-  const tensors = loadAdapterTensors(file); // caller owns; disposed below
-  const adopted = new Set<MlxArray>();
-  try {
-    // Two phases so the swap is atomic: validate EVERY target first (throws before
-    // anything is disposed), then dispose+adopt. A mid-loop failure must not leave
-    // lora.targets half-replaced (some checkpoint leaves, some fresh init).
-    const pending: Array<{ t: (typeof lora.targets)[number]; a: MlxArray; b: MlxArray }> = [];
-    for (const t of lora.targets) {
-      const a = tensors.get(`${t.modulePath}.lora_a`);
-      const b = tensors.get(`${t.modulePath}.lora_b`);
-      if (!a || !b) throw new Error(`warm-start: ${t.modulePath} missing lora_a/lora_b in ${file}`);
-      if (a.dtype !== Dtype.float32 || b.dtype !== Dtype.float32)
-        throw new Error(`warm-start: ${t.modulePath} tensors must be f32`);
-      if (a.shape[0] !== t.lw.a.shape[0] || a.shape[1] !== t.lw.a.shape[1] ||
-          b.shape[0] !== t.lw.b.shape[0] || b.shape[1] !== t.lw.b.shape[1])
-        throw new Error(`warm-start: shape mismatch at ${t.modulePath} — the checkpoint's rank/targets must match this run`);
-      pending.push({ t, a, b });
-    }
-    for (const { t, a, b } of pending) {
-      t.lw.a.dispose(); t.lw.b.dispose(); // free the fresh random/zero init
-      t.lw.a = a; t.lw.b = b;             // adopt the checkpoint weights as the live leaves
-      adopted.add(a); adopted.add(b);
-    }
-  } finally {
-    for (const [, arr] of tensors) if (!adopted.has(arr)) arr.dispose();
-  }
-  return lora.targets.length;
-}
-
-/** 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;
-  lora.targets.forEach((t, i) => {
-    (t.linear.adapters ??= new Map()).set(adapterId, t.lw);
-    t.linear.loraState = model.loraState;
-    t.linear.dropoutId = i; // stable per-target key for recompute-safe dropout
-  });
-  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);
-  model.loraState.dropoutRate = 0; // disable training dropout for inference
-  model.loraState.dropoutSeed = null;
-}
-
-/** 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;
-  /** rsLoRA: if true the effective per-layer scale is α/√rank (the loader must
-   *  apply the same, so it's recorded in optiq_lora_config.json). */
-  rsLora?: boolean;
-}
-
-export interface SaveAdapterOptions {
-  /** Override config-file writes for deterministic durability tests. */
-  writeFile?: (path: string, data: string) => Promise<unknown>;
-}
-
-/** 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 async function saveAdapter(
-  lora: TrainableLora,
-  dir: string,
-  cfg: SaveAdapterConfig,
-  appliedRanks: Record<string, number>,
-  options: SaveAdapterOptions = {},
-): Promise<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,
-    rs_lora: cfg.rsLora ?? false,
-    target_modules: cfg.targetModules,
-    num_layers: cfg.numLayers,
-    lora_parameters: { scale: cfg.scale, rank: cfg.rank, rs_lora: cfg.rsLora ?? false },
-    applied_ranks: appliedRanks,
-    source_model: cfg.baseModel,
-  };
-  // 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,
-    // mlx-bun's `scale` is the numerator before rsLoRA's √rank division.
-    // Ordinary PEFT needs alpha/r = scale; rsLoRA needs alpha/√r = scale/√r.
-    lora_alpha: cfg.rsLora ? cfg.scale : Math.round(cfg.rank * cfg.scale),
-    lora_dropout: 0.0,
-    peft_type: "LORA",
-    r: cfg.rank,
-    target_modules: cfg.targetModules,
-    task_type: "CAUSAL_LM",
-    use_rslora: cfg.rsLora ?? false,
-    optiq: { rank_scaling: cfg.rankScaling, applied_ranks: appliedRanks },
-  };
-  const writeFile = options.writeFile ?? ((path: string, data: string) => Bun.write(path, data));
-  await Promise.all([
-    writeFile(`${dir}/optiq_lora_config.json`, JSON.stringify(optiqCfg, null, 2) + "\n"),
-    writeFile(`${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();
-  }
-}