mlx-bun 0.0.4 → 0.0.6

package/src/train/job.ts

@@ -17,7 +17,7 @@ interface FinetuneSubmit {
   model_dir: string;
   data_dir: string;
   adapter_path: string;
-  method?: "sft" | "dpo";
+  method?: "sft" | "dpo" | "orpo";
   rank?: number;
   scale?: number;
   rank_scaling?: TrainConfig["rankScaling"];
@@ -32,12 +32,26 @@ interface FinetuneSubmit {
   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;
+  warm_start_adapter?: string;
 }
 
 function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir: string; cfg: TrainConfig } {
@@ -46,6 +60,14 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
   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";
+
   const cfg: TrainConfig = {
     ...DEFAULT_TRAIN_CONFIG,
     method: c.method ?? "sft",
@@ -55,7 +77,7 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
     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),
+    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,
@@ -64,12 +86,26 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
     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,
-    segmentSize: c.segment_size ?? DEFAULT_TRAIN_CONFIG.segmentSize,
+    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),
+    warmStartAdapter: c.warm_start_adapter ?? DEFAULT_TRAIN_CONFIG.warmStartAdapter,
     adapterPath: c.adapter_path,
     baseModel: c.model_dir,
   };
@@ -79,6 +115,18 @@ function parseConfig(raw: Record<string, unknown>): { modelDir: string; dataDir:
 export const finetuneRunner: JobRunner = async (emit, config) => {
   const { modelDir, dataDir, cfg } = parseConfig(config);
 
+  // Gemma/e4b needs these training env flags set BEFORE the model loads — same
+  // as the `mlx-bun train` CLI. The subprocess job path doesn't inherit them
+  // from the server, so without this a web-submitted ORPO/SFT run on Gemma uses
+  // the wrong kernels. ??= so an explicit env (or the CLI) still wins.
+  try {
+    const isGemma = (await Bun.file(`${modelDir}/config.json`).text()).toLowerCase().includes("gemma");
+    if (isGemma) {
+      process.env.MLX_BUN_PERF_KERNEL ??= "0";
+      process.env.MLX_BUN_FUSED_GELU ??= "0";
+    }
+  } catch {}
+
   emit({ type: "stage", stage: "load", progress: 0.01, message: `loading model ${modelDir}` });
 
   const { loadModelConfig } = await import("../config");
@@ -88,6 +136,19 @@ export const finetuneRunner: JobRunner = async (emit, config) => {
   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);

package/src/train/lora-params.ts

@@ -16,6 +16,7 @@ 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");
 
@@ -40,6 +41,7 @@ export function buildTrainableLora(
   ranks: Map<string, number>,
   scale: number,
   seed: number,
+  rsLora = false,
 ): TrainableLora {
   const linears = model.loraTargets();
   const key = ops.randomKey(BigInt(seed >>> 0));
@@ -61,21 +63,62 @@ export function buildTrainableLora(
     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 } });
+    // 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;
-  for (const t of lora.targets) {
+  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];
 }
 
@@ -83,6 +126,8 @@ export function attachForTraining(model: RuntimeModel, lora: TrainableLora, adap
 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). */
@@ -115,6 +160,9 @@ export interface SaveAdapterConfig {
   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;
 }
 
 /** Save the trained adapter into `dir` in the AdapterManager.mount format:
@@ -156,9 +204,10 @@ export function saveAdapter(
     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 },
+    lora_parameters: { scale: cfg.scale, rank: cfg.rank, rs_lora: cfg.rsLora ?? false },
     applied_ranks: appliedRanks,
     source_model: cfg.baseModel,
   };