mlx-bun 0.4.0 → 1.0.0

package/src/quantize/safetensors-writer.ts

@@ -1,249 +0,0 @@
-// Sharded safetensors writer: takes a list of named mlx arrays and writes
-// them to disk as one or more `model*.safetensors` files via mlx's native
-// saver (mlx_save_safetensors), matching the on-disk layout the loader
-// (src/weights.ts → ShardedSafetensors) reads back.
-//
-// Layout parity with mlx-lm / the MiniCPM5-OptiQ target:
-//   - single shard  → `model.safetensors`, NO index file.
-//   - multiple shards → `model-00001-of-0000N.safetensors`, plus a
-//     `model.safetensors.index.json` { metadata:{total_size,
-//     total_parameters}, weight_map: <name → shard> }.
-//   - metadata map carries {"format":"mlx"} so the file round-trips through
-//     mlx_load_safetensors with the same provenance HF/mlx writes.
-//
-// Memory discipline: each array is eval'd immediately before insertion so its
-// bytes are materialized one tensor at a time (the caller hands us already-
-// disposed-after-write arrays), bounding peak memory to ~one tensor over the
-// resident weight set rather than the whole model.
-
-import { ptr } from "bun:ffi";
-import { renameSync, writeFileSync } from "node:fs";
-import { join } from "node:path";
-import { C, takeMlxError } from "../mlx/ffi";
-import type { MlxArray } from "../mlx/array";
-
-const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
-
-/** Default greedy shard ceiling: 5 GiB (mlx-lm's MAX shard size). */
-export const DEFAULT_SHARD_BYTES = 5 << 30;
-
-export interface NamedTensor {
-  name: string;
-  array: MlxArray;
-}
-
-export interface ShardInfo {
-  /** Filename written into outDir (e.g. "model.safetensors"). */
-  file: string;
-  /** Tensor names placed in this shard, in insertion order. */
-  names: string[];
-  /** Byte size of the tensors in this shard (sum of nbytes). */
-  bytes: number;
-}
-
-export interface SafetensorsIndex {
-  metadata: { total_size: number; total_parameters: number };
-  weight_map: Record<string, string>;
-}
-
-export interface WriteResult {
-  shards: ShardInfo[];
-  totalSize: number;
-  totalParams: number;
-  /** The index object (also written to disk when shards.length > 1). */
-  index: SafetensorsIndex;
-}
-
-export interface WriteOpts {
-  /** Greedy shard ceiling in bytes (default DEFAULT_SHARD_BYTES). */
-  shardBytes?: number;
-}
-
-/** Insert one array into a native string→array map (key cstring, handle u64).
- *  Evals the array first so its bytes are materialized exactly here. */
-function insertArray(map: bigint, name: string, array: MlxArray): void {
-  array.eval();
-  if (C.mlx_map_string_to_array_insert(map, ptr(cstr(name)), array.handle) !== 0)
-    throw new Error(`map insert ${name} failed: ${takeMlxError() ?? ""}`);
-}
-
-/** Build the {"format":"mlx"} metadata map. Caller frees it. */
-function newMetaMap(): bigint {
-  const map = C.mlx_map_string_to_string_new();
-  if (C.mlx_map_string_to_string_insert(map, ptr(cstr("format")), ptr(cstr("mlx"))) !== 0)
-    throw new Error(`metadata insert failed: ${takeMlxError() ?? ""}`);
-  return map;
-}
-
-/** Pack the shard's tensors into a fresh native map and save to `path`. */
-function saveShard(path: string, tensors: NamedTensor[]): void {
-  const arrMap = C.mlx_map_string_to_array_new();
-  const metaMap = newMetaMap();
-  try {
-    for (const { name, array } of tensors) insertArray(arrMap, name, array);
-    if (C.mlx_save_safetensors(ptr(cstr(path)), arrMap, metaMap) !== 0)
-      throw new Error(`mlx_save_safetensors(${path}) failed: ${takeMlxError() ?? ""}`);
-  } finally {
-    C.mlx_map_string_to_array_free(arrMap);
-    C.mlx_map_string_to_string_free(metaMap);
-  }
-}
-
-/** Shard filename for shard i of n (1-based, 5-digit, matches mlx-lm). */
-function shardName(i: number, n: number): string {
-  if (n === 1) return "model.safetensors";
-  const pad = (x: number) => String(x).padStart(5, "0");
-  return `model-${pad(i + 1)}-of-${pad(n)}.safetensors`;
-}
-
-/** Greedily bin tensors into shards no larger than `shardBytes`. A single
- *  tensor larger than the ceiling still gets its own shard (never split). */
-function planShards(tensors: NamedTensor[], shardBytes: number): NamedTensor[][] {
-  const shards: NamedTensor[][] = [];
-  let cur: NamedTensor[] = [];
-  let curBytes = 0;
-  for (const t of tensors) {
-    const nb = t.array.nbytes;
-    if (cur.length > 0 && curBytes + nb > shardBytes) {
-      shards.push(cur);
-      cur = [];
-      curBytes = 0;
-    }
-    cur.push(t);
-    curBytes += nb;
-  }
-  if (cur.length > 0) shards.push(cur);
-  return shards.length > 0 ? shards : [[]];
-}
-
-/**
- * Write `tensors` to `outDir` as sharded safetensors. Returns shard layout,
- * total byte size, total parameter count, and the (single source of truth)
- * index object. An index file is written only when there is more than one
- * shard — matching the loader's expectation that a lone `model.safetensors`
- * carries no sidecar index.
- */
-export function writeShardedSafetensors(
-  outDir: string,
-  tensors: NamedTensor[],
-  opts: WriteOpts = {},
-): WriteResult {
-  const shardBytes = opts.shardBytes ?? DEFAULT_SHARD_BYTES;
-  const plan = planShards(tensors, shardBytes);
-  const n = plan.length;
-
-  const shards: ShardInfo[] = [];
-  const weight_map: Record<string, string> = {};
-  let totalSize = 0;
-  let totalParams = 0;
-
-  for (let i = 0; i < n; i++) {
-    const group = plan[i]!;
-    const file = shardName(i, n);
-    let bytes = 0;
-    for (const t of group) {
-      bytes += t.array.nbytes;
-      totalParams += t.array.size;
-      weight_map[t.name] = file;
-    }
-    totalSize += bytes;
-    saveShard(join(outDir, file), group);
-    shards.push({ file, names: group.map((t) => t.name), bytes });
-  }
-
-  // weight_map sorted by name for a stable, diff-friendly index.
-  const sortedMap: Record<string, string> = {};
-  for (const name of Object.keys(weight_map).sort()) sortedMap[name] = weight_map[name]!;
-
-  const index: SafetensorsIndex = {
-    metadata: { total_size: totalSize, total_parameters: totalParams },
-    weight_map: sortedMap,
-  };
-
-  if (n > 1) {
-    writeFileSync(
-      join(outDir, "model.safetensors.index.json"),
-      JSON.stringify(index, null, 2),
-    );
-  }
-
-  return { shards, totalSize, totalParams, index };
-}
-
-/**
- * Incremental sharded writer for model-scale streaming transforms (27B fold/
- * quantize): tensors are added one at a time; when the current shard reaches
- * the byte ceiling it is SAVED AND ITS ARRAYS DISPOSED immediately, so peak
- * memory is one shard plus the tensor in flight — never the whole model
- * (the whole-list writeShardedSafetensors path holds every array until the
- * caller disposes them, which OOMs a 51 GB model on a 32 GB box).
- *
- * Shards are written under temp names (the final `-of-NNNNN` count is unknown
- * until finish()) and renamed at finish, which also writes the index.
- * add() takes ownership of the array (evaluates it now, disposes after save).
- */
-export class ShardedWriter {
-  readonly #outDir: string;
-  readonly #shardBytes: number;
-  #cur: NamedTensor[] = [];
-  #curBytes = 0;
-  #done: { tmp: string; names: string[]; bytes: number }[] = [];
-  #totalParams = 0;
-
-  constructor(outDir: string, opts: WriteOpts = {}) {
-    this.#outDir = outDir;
-    this.#shardBytes = opts.shardBytes ?? DEFAULT_SHARD_BYTES;
-  }
-
-  add(name: string, array: MlxArray): void {
-    array.eval();
-    const nb = array.nbytes;
-    if (this.#cur.length > 0 && this.#curBytes + nb > this.#shardBytes) this.#flush();
-    this.#cur.push({ name, array });
-    this.#curBytes += nb;
-    this.#totalParams += array.size;
-  }
-
-  #flush(): void {
-    if (this.#cur.length === 0) return;
-    const tmp = `model-shard-${String(this.#done.length + 1).padStart(5, "0")}.tmp.safetensors`;
-    saveShard(join(this.#outDir, tmp), this.#cur);
-    this.#done.push({
-      tmp,
-      names: this.#cur.map((t) => t.name),
-      bytes: this.#curBytes,
-    });
-    for (const t of this.#cur) t.array.dispose();
-    this.#cur = [];
-    this.#curBytes = 0;
-  }
-
-  finish(): WriteResult {
-    this.#flush();
-    const n = this.#done.length;
-    const shards: ShardInfo[] = [];
-    const weight_map: Record<string, string> = {};
-    let totalSize = 0;
-    for (let i = 0; i < n; i++) {
-      const d = this.#done[i]!;
-      const file = shardName(i, n);
-      renameSync(join(this.#outDir, d.tmp), join(this.#outDir, file));
-      for (const name of d.names) weight_map[name] = file;
-      totalSize += d.bytes;
-      shards.push({ file, names: d.names, bytes: d.bytes });
-    }
-    const sortedMap: Record<string, string> = {};
-    for (const name of Object.keys(weight_map).sort()) sortedMap[name] = weight_map[name]!;
-    const index: SafetensorsIndex = {
-      metadata: { total_size: totalSize, total_parameters: this.#totalParams },
-      weight_map: sortedMap,
-    };
-    if (n > 1) {
-      writeFileSync(
-        join(this.#outDir, "model.safetensors.index.json"),
-        JSON.stringify(index, null, 2),
-      );
-    }
-    return { shards, totalSize, totalParams: this.#totalParams, index };
-  }
-}

package/src/quantize/sensitivity.ts

@@ -1,216 +0,0 @@
-// Per-layer quantization sensitivity analysis (native TS port of
-// optiq/core/sensitivity.py `analyze_sensitivity_exact`, uniform-baseline path).
-//
-// For each quantizable Linear in a loaded (already-quantized) model:
-//   1. dequantize its packed weight back to bf16 — this is the bf16 "source"
-//      the OptIQ uniform_4bit path streams off disk; here it's self-contained
-//      because the running model already carries quantized weights.
-//   2. for each candidate bit-width, re-quantize that bf16 source at the new
-//      bits and SWAP it into the running QuantizedLinear (mutate w/scales/biases
-//      + the QuantSpec bits/groupSize).
-//   3. run a no-KV full-sequence forward on each calibration sample, compute
-//      KL(reference || current) (port of `_kl_from_ref`), average over samples.
-//   4. restore the layer's original quantized weight.
-//
-// Reference logits = the UNMODIFIED running model (taken once up front). When a
-// candidate bit equals the layer's baseline bit-width, re-quantizing reproduces
-// exactly what the layer already is, so KL is 0 by construction — we record 0.0
-// and skip the forward passes (port of the OptIQ baseline-bit short-circuit).
-//
-// Faithful to: `_simulate_quantize` (quantize→dequantize via MlxArray ops),
-// `_mutate_quantized_layer_to_bits`, `_kl_from_ref`, the uniform-4bit-baseline
-// path, and the bits==baseline ⇒ 0 short-circuit.
-
-import { Dtype } from "../mlx/ffi";
-import { cpuStream, MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { trainForward } from "../train/forward";
-import type { RuntimeModel } from "../model/factory";
-import type {
-  QuantizedLinear,
-  QuantizedLinearState,
-} from "../model/gemma4-base";
-
-/** Per-layer KL sensitivity at each candidate bit-width (port of
- *  optiq SensitivityResult). */
-export interface SensitivityResult {
-  layerName: string;
-  /** bit-width → mean KL(reference || that-bit) over calibration samples. */
-  sensitivities: Record<number, number>;
-  paramCount: number;
-}
-
-export interface AnalyzeSensitivityOptions {
-  candidateBits?: number[];
-  groupSize?: number;
-}
-
-// --------------------------------------------------------------------------
-// Scoped state exchange for sensitivity probes. QuantizedLinear owns the only
-// mutation capability; inference consumers retain a read-only getter surface.
-// --------------------------------------------------------------------------
-
-/** Snapshot of a QuantizedLinear's quantized state (for restore). */
-type LayerState = QuantizedLinearState;
-
-function captureState(layer: QuantizedLinear): LayerState {
-  return { w: layer.w, scales: layer.scales, biases: layer.biases, spec: layer.spec };
-}
-
-function restoreState(layer: QuantizedLinear, state: LayerState): void {
-  const transient = layer.exchangeQuantizedState(state);
-  // Dispose the transient swapped-in tensors before restoring the originals.
-  if (transient.w !== state.w) transient.w.dispose();
-  if (transient.scales !== state.scales) transient.scales.dispose();
-  if (transient.biases && transient.biases !== state.biases) transient.biases.dispose();
-}
-
-/** Quantize `bf16Weight` at `newBits`/`groupSize` and swap it into `layer`
- *  in place (port of `_mutate_quantized_layer_to_bits`). The previously
- *  swapped-in tensors (if any beyond the captured original) are released by
- *  the caller via restoreState; this only installs the new ones. */
-function mutateLayerToBits(
-  layer: QuantizedLinear,
-  bf16Weight: MlxArray,
-  newBits: number,
-  groupSize: number,
-  original: LayerState,
-): void {
-  const q = ops.quantize(bf16Weight, groupSize, newBits, "affine", cpuStream);
-  const previous = layer.exchangeQuantizedState({
-    w: q.packed,
-    scales: q.scales,
-    biases: q.biases,
-    spec: { bits: newBits, groupSize, mode: "affine" },
-  });
-  // Free any previously swapped-in transient tensors (not the original).
-  if (previous.w !== original.w) previous.w.dispose();
-  if (previous.scales !== original.scales) previous.scales.dispose();
-  if (previous.biases && previous.biases !== original.biases) previous.biases.dispose();
-  ops.evalAll([q.packed, q.scales, q.biases]);
-}
-
-// --------------------------------------------------------------------------
-// KL primitive (port of _kl_from_ref)
-// --------------------------------------------------------------------------
-
-/** KL(reference || current) averaged across batch/seq (port of `_kl_from_ref`):
- *    log_cur = cur - logsumexp(cur); log_ref = ref - logsumexp(ref)
- *    ref_probs = softmax(ref)
- *    mean( sum( ref_probs * (log_ref - log_cur), axis=-1 ) )
- *  Returns the scalar KL as a JS number. Casts to f32 for stability. */
-export function klFromRef(curLogits: MlxArray, refLogits: MlxArray): number {
-  const cur = curLogits.dtype === Dtype.float32 ? curLogits : curLogits.astype(Dtype.float32);
-  const ref = refLogits.dtype === Dtype.float32 ? refLogits : refLogits.astype(Dtype.float32);
-
-  const lseCur = ops.logsumexpAxis(cur, -1, true);
-  const lseRef = ops.logsumexpAxis(ref, -1, true);
-  const logCur = ops.sub(cur, lseCur);
-  const logRef = ops.sub(ref, lseRef);
-  const refProbs = ops.softmaxAxis(ref, -1, true);
-
-  const diff = ops.sub(logRef, logCur);
-  const weighted = ops.mul(refProbs, diff);
-  const summed = ops.sumAxis(weighted, -1, false);
-  const meaned = ops.meanAll(summed, false);
-
-  const val = meaned.toFloat32()[0]!;
-
-  for (const a of [lseCur, lseRef, logCur, logRef, refProbs, diff, weighted, summed, meaned]) {
-    a.dispose();
-  }
-  if (cur !== curLogits) cur.dispose();
-  if (ref !== refLogits) ref.dispose();
-  return val;
-}
-
-// --------------------------------------------------------------------------
-// Public entry: analyzeSensitivityExact
-// --------------------------------------------------------------------------
-
-/** Per-layer KL sensitivity of an already-quantized MLX model.
- *
- *  @param model        loaded RuntimeModel (its QuantizedLinears are the
- *                      uniform baseline; their dequantized weights are the
- *                      self-contained bf16 source).
- *  @param layers       the quantizable linears to probe, keyed by module path
- *                      (typically `model.loraTargets()`).
- *  @param calIds       calibration samples — each a row of int32 token ids.
- *  @param options      candidateBits (default [4,8]) and groupSize (default 64).
- */
-export function analyzeSensitivityExact(
-  model: RuntimeModel,
-  layers: Map<string, QuantizedLinear>,
-  calIds: number[][],
-  options: AnalyzeSensitivityOptions = {},
-  onProgress?: (done: number, total: number, layerName: string) => void,
-): SensitivityResult[] {
-  const candidateBits = [...(options.candidateBits ?? [4, 8])];
-  const groupSize = options.groupSize ?? 64;
-
-  // Calibration sample arrays (int32 [1, L]).
-  const calArrays = calIds.map((ids) =>
-    MlxArray.fromInt32(new Int32Array(ids), [1, ids.length]),
-  );
-
-  try {
-    // Reference logits from the unmodified running model (taken once).
-    const refLogits: MlxArray[] = calArrays.map((ids) => trainForward(model, ids));
-    ops.evalAll(refLogits);
-
-    const entries = [...layers.entries()];
-    const results: SensitivityResult[] = [];
-
-    try {
-      for (let li = 0; li < entries.length; li++) {
-        const [layerName, layer] = entries[li]!;
-        const original = captureState(layer);
-        const baselineBits = layer.spec.bits;
-
-        // Self-contained bf16 source: dequantize the layer's own packed weight.
-        const bf16Source = ops.dequantize(
-          original.w, original.scales, original.biases, original.spec, cpuStream,
-        ).astype(Dtype.bfloat16, cpuStream);
-        ops.evalAll([bf16Source]);
-
-        const sensitivities: Record<number, number> = {};
-        try {
-          for (const bits of candidateBits) {
-            if (bits === baselineBits) {
-              // KL is 0 by construction (re-quantizing to the baseline bit-width
-              // reproduces the running model). Skip the forwards.
-              sensitivities[bits] = 0.0;
-              continue;
-            }
-            mutateLayerToBits(layer, bf16Source, bits, groupSize, original);
-
-            let klSum = 0;
-            for (let i = 0; i < calArrays.length; i++) {
-              const cur = trainForward(model, calArrays[i]!);
-              klSum += klFromRef(cur, refLogits[i]!);
-              cur.dispose();
-            }
-            sensitivities[bits] = klSum / Math.max(calArrays.length, 1);
-          }
-        } finally {
-          // Restore even when quantization or a calibration forward throws.
-          restoreState(layer, original);
-          bf16Source.dispose();
-        }
-
-        // param count = product of the dequantized (full) weight shape.
-        const fullShape = [layer.outFeatures, layer.inFeatures];
-        const paramCount = fullShape.reduce((a, b) => a * b, 1);
-
-        results.push({ layerName, sensitivities, paramCount });
-        onProgress?.(li + 1, entries.length, layerName);
-      }
-    } finally {
-      for (const r of refLogits) r.dispose();
-    }
-
-    return results;
-  } finally {
-    for (const a of calArrays) a.dispose();
-  }
-}