mlx-bun 0.4.0 → 1.0.0

package/src/train/merge.ts

@@ -1,195 +0,0 @@
-// Merge multiple LoRA adapters into one — rank-concat, mathematically exact.
-// Port of optiq/lora/merge.py.
-//
-// For two adapters on a layer (A1[in,r1] B1[r1,out] scale s1; A2 B2 s2):
-//   A_m = concat([A1, A2], rankAxis)           shape [in, r1+r2]
-//   B_m = concat([s1·B1, s2·B2], rankAxis)      shape [r1+r2, out]
-//   (x @ A_m) @ B_m == s1·(x@A1@B1) + s2·(x@A2@B2)
-// so the merged adapter writes scale=1.0 (per-source scales folded into B).
-//
-// Tensor names follow the AdapterManager.mount format (lora_a / lora_b).
-
-import { mkdirSync, existsSync } from "node:fs";
-import { ptr, read } from "bun:ffi";
-import { C } from "../mlx/ffi";
-import { MlxArray, cpuStream } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { SafetensorsFile } from "../safetensors";
-
-const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
-
-export interface MergeStats {
-  layersMerged: number;
-  layersOnlyInOne: number;
-  totalKeysOut: number;
-  sources: string[];
-  scales: number[];
-}
-
-/** Merge N adapter directories into `outputDir`. `scales` overrides the
- *  per-source scales (else each source's adapter config scale is used). */
-export async function mergeAdapters(
-  adapterDirs: string[],
-  outputDir: string,
-  scales?: number[],
-): Promise<MergeStats> {
-  if (adapterDirs.length < 2) throw new Error("mergeAdapters needs at least two sources");
-
-  const resolvedScales = scales ?? (await Promise.all(adapterDirs.map(resolveScale)));
-  if (resolvedScales.length !== adapterDirs.length)
-    throw new Error("scales length must match adapter count");
-
-  const sources = adapterDirs.map((d) => ({
-    tensors: loadAdapterTensors(resolveSafetensors(d)),
-  }));
-
-  // Universe of module paths across all sources.
-  const allMods = new Set<string>();
-  for (const s of sources)
-    for (const k of s.tensors.keys())
-      if (k.endsWith(".lora_a")) allMods.add(k.slice(0, -".lora_a".length));
-
-  const merged = new Map<string, MlxArray>();
-  const stats: MergeStats = {
-    layersMerged: 0,
-    layersOnlyInOne: 0,
-    totalKeysOut: 0,
-    sources: adapterDirs,
-    scales: resolvedScales,
-  };
-
-  try {
-    for (const mod of [...allMods].sort()) {
-      const aParts: MlxArray[] = [];
-      const bParts: MlxArray[] = [];
-      for (let s = 0; s < sources.length; s++) {
-        const a = sources[s]!.tensors.get(`${mod}.lora_a`);
-        const b = sources[s]!.tensors.get(`${mod}.lora_b`);
-        if (a && b) {
-          aParts.push(a);
-          const scale = resolvedScales[s]!;
-          bParts.push(scale === 1.0 ? b : ops.mulScalar(b, scale));
-        }
-      }
-      if (aParts.length === 0) continue;
-      if (aParts.length === 1) {
-        stats.layersOnlyInOne++;
-        // A is the original; B may be a scaled copy (own it) or the original.
-        merged.set(`${mod}.lora_a`, aParts[0]!);
-        merged.set(`${mod}.lora_b`, bParts[0]!);
-      } else {
-        // rank axis: A is [in, r] → axis 1; B is [r, out] → axis 0.
-        const aMerged = ops.concatAxis(aParts, 1);
-        const bMerged = ops.concatAxis(bParts, 0);
-        merged.set(`${mod}.lora_a`, aMerged);
-        merged.set(`${mod}.lora_b`, bMerged);
-        stats.layersMerged++;
-      }
-    }
-
-    stats.totalKeysOut = merged.size;
-    mkdirSync(outputDir, { recursive: true });
-
-    const map = C.mlx_map_string_to_array_new();
-    const meta = C.mlx_map_string_to_string_new();
-    try {
-      for (const [name, arr] of merged) {
-        arr.eval();
-        if (C.mlx_map_string_to_array_insert(map, ptr(cstr(name)), arr.handle) !== 0)
-          throw new Error(`map insert ${name} failed`);
-      }
-      const file = `${outputDir}/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);
-    }
-
-    // Merged config: scale 1.0 (folded into B), rank = sum of source ranks.
-    let rankSum = 0;
-    for (const d of adapterDirs) rankSum += (await resolveRank(d)) ?? 0;
-    const cfg = {
-      fine_tune_type: "lora",
-      lora_parameters: { scale: 1.0, rank: rankSum || null },
-      r: rankSum || null,
-      peft_type: "LORA",
-      optiq_merge: { sources: adapterDirs, scales: resolvedScales },
-    };
-    await Bun.write(`${outputDir}/adapter_config.json`, JSON.stringify(cfg, null, 2) + "\n");
-    await Bun.write(`${outputDir}/optiq_lora_config.json`, JSON.stringify(cfg, null, 2) + "\n");
-  } finally {
-    // dispose all source + merged arrays (merged copies are owned here).
-    for (const s of sources) for (const a of s.tensors.values()) a.dispose();
-    for (const a of merged.values()) a.dispose();
-  }
-
-  return stats;
-}
-
-function resolveSafetensors(dir: string): string {
-  for (const p of [`${dir}/best/adapters.safetensors`, `${dir}/adapters.safetensors`])
-    if (existsSync(p)) return p;
-  if (existsSync(dir)) return dir; // explicit file path
-  throw new Error(`${dir}: no adapters.safetensors`);
-}
-
-async function resolveScale(dir: string): Promise<number> {
-  const cfg = await readConfig(dir);
-  if (cfg?.lora_parameters?.scale != null) return Number(cfg.lora_parameters.scale);
-  if (cfg?.lora_alpha != null && cfg?.r) return Number(cfg.lora_alpha) / Number(cfg.r);
-  if (cfg?.scale != null) return Number(cfg.scale);
-  return 1.0;
-}
-
-async function resolveRank(dir: string): Promise<number | null> {
-  const cfg = await readConfig(dir);
-  const r = cfg?.lora_parameters?.rank ?? cfg?.r;
-  return r != null ? Number(r) : null;
-}
-
-async function readConfig(dir: string): Promise<any | null> {
-  for (const name of ["adapter_config.json", "optiq_lora_config.json"]) {
-    const f = Bun.file(`${dir}/${name}`);
-    if (await f.exists()) return f.json();
-  }
-  return null;
-}
-
-/** Materialize every tensor from an adapter safetensors file (mirrors the
- *  read idiom in src/lora.ts loadAdapterTensors). */
-function loadAdapterTensors(file: string): Map<string, MlxArray> {
-  const sf = SafetensorsFile.open(file);
-  const names = [...sf.tensors.keys()];
-  sf.mmap.unmap();
-
-  const arrMapSlot = new BigUint64Array([C.mlx_map_string_to_array_new()]);
-  const metaMapSlot = new BigUint64Array([C.mlx_map_string_to_string_new()]);
-  const arrMapPtr = ptr(arrMapSlot);
-  const metaMapPtr = ptr(metaMapSlot);
-  const status = C.mlx_load_safetensors(arrMapPtr, metaMapPtr, ptr(cstr(file)), cpuStream);
-  C.mlx_map_string_to_string_free(read.u64(metaMapPtr, 0));
-  const mapHandle = read.u64(arrMapPtr, 0);
-  if (status !== 0) {
-    C.mlx_map_string_to_array_free(mapHandle);
-    throw new Error(`mlx_load_safetensors(${file}) failed`);
-  }
-  const out = new Map<string, MlxArray>();
-  try {
-    for (const name of names) {
-      const slot = new BigUint64Array([C.mlx_array_new()]);
-      const slotPtr = ptr(slot);
-      if (C.mlx_map_string_to_array_get(slotPtr, mapHandle, ptr(cstr(name))) !== 0)
-        throw new Error(`adapter tensor ${name} missing from native map`);
-      const arr = new MlxArray(read.u64(slotPtr, 0));
-      arr.eval();
-      out.set(name, arr);
-    }
-  } catch (e) {
-    for (const a of out.values()) a.dispose();
-    C.mlx_map_string_to_array_free(mapHandle);
-    throw e;
-  }
-  C.mlx_map_string_to_array_free(mapHandle);
-  return out;
-}

package/src/train/optimizer.ts

@@ -1,183 +0,0 @@
-// AdamW optimizer over a flat list of MlxArray parameters.
-//
-// Port of mlx.optimizers.AdamW: decoupled weight decay (Loshchilov &
-// Hutter), bias-corrected first/second moments. Each step REPLACES every
-// param handle in place (old disposed, new array swapped in) and invokes the
-// `write(i, newParam)` callback so the owning LoraWeights.a/b picks up the new
-// handle for the next forward.
-//
-// Update (per param p with grad g, step t):
-//   p   <- p - lr * wd * p               (decoupled weight decay)
-//   m   <- β1·m + (1-β1)·g
-//   v   <- β2·v + (1-β2)·g²
-//   m̂  =  m / (1 - β1^t)
-//   v̂  =  v / (1 - β2^t)
-//   p   <- p - lr * m̂ / (√v̂ + ε)
-
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-
-export interface AdamWOptions {
-  lr: number;
-  betas?: [number, number];
-  eps?: number;
-  weightDecay?: number;
-  /** Optional per-parameter LR multiplier (effective lr for param i is
-   *  `lr · lrScale[i]`). Used for LoRA+ — the B leaves get a higher LR than A.
-   *  Length must match `params`; missing/undefined entries default to 1. */
-  lrScale?: number[];
-}
-
-export class AdamW {
-  #params: MlxArray[];
-  readonly #write: (i: number, p: MlxArray) => void;
-  readonly #m: MlxArray[];
-  readonly #v: MlxArray[];
-  readonly #b1: number;
-  readonly #b2: number;
-  readonly #eps: number;
-  readonly #wd: number;
-  readonly #lrScale: number[];
-  #lr: number;
-  #t = 0;
-
-  /** @param params live param arrays (e.g. flatParams(lora)).
-   *  @param write callback to push a replaced param back into its owner
-   *               (e.g. lw.a / lw.b). If omitted, only the internal array is
-   *               updated (the caller must read params via `getParam`). */
-  constructor(params: MlxArray[], opts: AdamWOptions, write?: (i: number, p: MlxArray) => void) {
-    this.#params = [...params];
-    this.#write = write ?? (() => {});
-    this.#lr = opts.lr;
-    [this.#b1, this.#b2] = opts.betas ?? [0.9, 0.999];
-    this.#eps = opts.eps ?? 1e-8;
-    this.#wd = opts.weightDecay ?? 0.01;
-    this.#lrScale = opts.lrScale ?? [];
-    // m, v as zeros like each param (handles 0-d params, where ops.zeros's
-    // empty shape buffer would trip ptr()).
-    this.#m = params.map(zerosLikeF32);
-    this.#v = params.map(zerosLikeF32);
-  }
-
-  /** Current parameter array at index i (after the latest step). */
-  getParam(i: number): MlxArray {
-    return this.#params[i]!;
-  }
-
-  set lr(value: number) {
-    this.#lr = value;
-  }
-  get lr(): number {
-    return this.#lr;
-  }
-
-  get step_count(): number {
-    return this.#t;
-  }
-
-  /** Apply one AdamW update. `grads[i]` is dL/dparams[i]. Takes ownership of
-   *  the grads (they are disposed here). */
-  step(grads: MlxArray[]): void {
-    if (grads.length !== this.#params.length)
-      throw new Error(`AdamW.step: ${grads.length} grads for ${this.#params.length} params`);
-    this.#t += 1;
-    const bc1 = 1 - Math.pow(this.#b1, this.#t);
-    const bc2 = 1 - Math.pow(this.#b2, this.#t);
-
-    for (let i = 0; i < this.#params.length; i++) {
-      const p = this.#params[i]!;
-      const lri = this.#lr * (this.#lrScale[i] ?? 1); // per-param lr (LoRA+)
-      const g0 = grads[i]!;
-      const g = g0.dtype === Dtype.float32 ? g0 : g0.astype(Dtype.float32);
-
-      // m <- β1·m + (1-β1)·g
-      const m0 = this.#m[i]!;
-      const mScaled = ops.mulScalar(m0, this.#b1);
-      const gScaled = ops.mulScalar(g, 1 - this.#b1);
-      const m1 = ops.add(mScaled, gScaled);
-      mScaled.dispose();
-      gScaled.dispose();
-      m0.dispose();
-      this.#m[i] = m1;
-
-      // v <- β2·v + (1-β2)·g²
-      const v0 = this.#v[i]!;
-      const g2 = ops.square(g);
-      const vScaled = ops.mulScalar(v0, this.#b2);
-      const g2Scaled = ops.mulScalar(g2, 1 - this.#b2);
-      const v1 = ops.add(vScaled, g2Scaled);
-      g2.dispose();
-      vScaled.dispose();
-      g2Scaled.dispose();
-      v0.dispose();
-      this.#v[i] = v1;
-
-      // p <- p - lr·wd·p   (decoupled weight decay; per-param lr)
-      let pNext = p;
-      if (this.#wd !== 0) {
-        const decay = ops.mulScalar(p, lri * this.#wd);
-        pNext = ops.sub(p, decay);
-        decay.dispose();
-      }
-
-      // m̂ / (√v̂ + ε)
-      const mHat = ops.mulScalar(m1, 1 / bc1);
-      const vHat = ops.mulScalar(v1, 1 / bc2);
-      const sqrtV = ops.sqrt(vHat);
-      const eps = ops.scalarLike(this.#eps, sqrtV);
-      const denom = ops.add(sqrtV, eps);
-      const update = ops.div(mHat, denom);
-      const lrUpdate = ops.mulScalar(update, lri);
-      const pFinal = ops.sub(pNext, lrUpdate);
-
-      for (const a of [mHat, vHat, sqrtV, eps, denom, update, lrUpdate]) a.dispose();
-      if (pNext !== p) pNext.dispose();
-      if (g !== g0) g.dispose();
-      g0.dispose();
-      p.dispose();
-
-      this.#params[i] = pFinal;
-      this.#write(i, pFinal);
-    }
-  }
-
-  /** Eval all optimizer state + params (call once per step after `step`). */
-  evalState(): void {
-    ops.evalAll([...this.#params, ...this.#m, ...this.#v]);
-  }
-
-  dispose(): void {
-    for (const a of this.#m) a.dispose();
-    for (const a of this.#v) a.dispose();
-    // params are owned by their LoraWeights; not disposed here.
-  }
-}
-
-/** Zeros with the same shape as `p`, as float32. Uses mul-by-0 so 0-d
- *  params work (ops.zeros's empty shape buffer trips ptr()). */
-function zerosLikeF32(p: MlxArray): MlxArray {
-  const f = p.dtype === Dtype.float32 ? p : p.astype(Dtype.float32);
-  const z = ops.mulScalar(f, 0);
-  if (f !== p) f.dispose();
-  return z;
-}
-
-/** Linear warmup → cosine decay schedule (to 10% of peak), used by DPO.
- *  Returns lr for a 1-based step. Port of dpo.py lr_schedule. */
-export function warmupCosineSchedule(
-  peakLr: number,
-  warmupIters: number,
-  totalIters: number,
-): (step: number) => number {
-  const minLr = 0.1 * peakLr;
-  const warmup = Math.max(warmupIters, 1);
-  const remaining = Math.max(totalIters - warmupIters, 1);
-  return (step: number) => {
-    // step is 1-based; convert to 0-based for the ramp like the reference.
-    const s = step - 1;
-    if (s < warmup) return Math.min(peakLr, (peakLr * (s + 1)) / warmup);
-    const progress = Math.max(0, Math.min(1, (s - warmup) / remaining));
-    return minLr + 0.5 * (peakLr - minLr) * (1 + Math.cos(Math.PI * progress));
-  };
-}