mlx-bun 0.0.4

package/src/quantize/calibration.ts

@@ -0,0 +1,180 @@
+// Calibration data loading for sensitivity analysis.
+//
+// Port of optiq/calibration/datasets.py `load_llm_calibration` /
+// `_load_optiq_calibration`: load a multi-domain JSONL mix, tokenize, build one
+// long token stream, and chunk it into fixed-length sequences. The
+// `calibrationFn` contract is the TS analogue of OptIQ's
+// `list[(args_tuple, kwargs_dict)]`: here, just a list of `number[]` token-id
+// prompts (each a (1, seq_len) row the model forwards).
+//
+// Two deviations from the Python, both forced by our pure-JS tokenizer surface
+// (src/tokenizer.ts exposes plain `encode` only, no `apply_chat_template`):
+//   - `messages` samples use the role-prefixed dumb-concatenation fallback —
+//     exactly OptIQ's `no chat template` branch.
+//   - `add_special_tokens=False` is the OptIQ encode flag; our `encode` already
+//     omits special tokens for raw substrings.
+//
+// Source mix: the bundled `optiq.jsonl` if locatable on disk (the oracle venv's
+// optiq/calibration/data/optiq.jsonl), else a small built-in fallback corpus so
+// the path works with no external data.
+
+import { existsSync } from "node:fs";
+import type { LoadedTokenizer } from "../tokenizer";
+
+/** One JSONL calibration sample (subset of the OptIQ schema we can consume). */
+interface CalibrationSample {
+  domain?: string;
+  text?: string;
+  messages?: Array<{ role: string; content: string }>;
+  tools?: unknown;
+}
+
+/** Candidate on-disk locations for the bundled optiq.jsonl mix. */
+function bundledMixCandidates(): string[] {
+  const home = process.env.HOME ?? "";
+  const rel = "lib/python3.14/site-packages/optiq/calibration/data/optiq.jsonl";
+  return [
+    process.env.MLX_BUN_CALIBRATION_JSONL,
+    `${home}/Code/mlx-lm-example/.venv/${rel}`,
+    `${home}/Code/mlx-lm/.venv/${rel}`,
+  ].filter((p): p is string => !!p);
+}
+
+/** A tiny built-in fallback corpus (used only when no JSONL mix is found).
+ *  Multi-domain prose so groups have meaningful activation diversity. */
+const FALLBACK_CORPUS: string[] = [
+  "The history of computing stretches back to mechanical calculators and the abacus. Charles Babbage designed the Analytical Engine, a general-purpose mechanical computer, in the nineteenth century. Ada Lovelace wrote what is often considered the first algorithm intended for such a machine. The theoretical foundations were later laid by Alan Turing, whose model of computation defined the limits of what machines can compute.",
+  "In physics, the conservation of energy states that the total energy of an isolated system remains constant over time. Energy can neither be created nor destroyed; rather, it transforms from one form to another. Kinetic energy becomes potential energy as a ball rises, and back again as it falls. Thermodynamics extends these ideas to heat, work, and entropy in large systems.",
+  "Functions in programming are reusable blocks of code that perform a specific task. They take inputs, called arguments, and may return a value. Pure functions have no side effects and always return the same output for the same input, which makes them easier to test and reason about. Recursion is a technique where a function calls itself to solve smaller subproblems.",
+  "The water cycle describes the continuous movement of water on, above, and below the surface of the Earth. Water evaporates from oceans and lakes, condenses into clouds, and precipitates as rain or snow. Rivers carry water back to the sea, and groundwater slowly percolates through soil and rock. Plants release water vapor through transpiration, contributing to atmospheric moisture.",
+  "Economic markets coordinate the production and exchange of goods through prices. When demand for a product rises and supply stays fixed, prices tend to increase, signaling producers to make more. Competition drives firms to lower costs and innovate. Externalities, such as pollution, are costs not reflected in market prices and often require regulation to address.",
+  "A neural network is composed of layers of interconnected nodes, each applying a weighted sum followed by a nonlinear activation. During training, the network adjusts its weights to minimize a loss function via gradient descent. Backpropagation efficiently computes gradients by applying the chain rule layer by layer. Deeper networks can represent more complex functions but are harder to train.",
+];
+
+/**
+ * Build a calibration callable from a tokenizer.
+ *
+ * @returns a zero-arg function producing `number[][]` — N token-id rows, each
+ *          of length `seqLen`, ready to forward through a model.
+ */
+export function loadLlmCalibration(
+  tokenizer: LoadedTokenizer,
+  opts: {
+    nSamples?: number;
+    seqLen?: number;
+    seed?: number;
+    mix?: string;
+  } = {},
+): () => number[][] {
+  const nSamples = opts.nSamples ?? 32;
+  const seqLen = opts.seqLen ?? 1024;
+  const seed = opts.seed ?? 42;
+  const mix = opts.mix ?? "optiq";
+
+  const samples = loadSamples(mix);
+
+  // Concatenate every sample's tokens into one long stream.
+  const allTokens: number[] = [];
+  for (const s of samples) {
+    let text: string | undefined;
+    if (s.messages) {
+      // No chat template available in our JS tokenizer — role-prefixed
+      // concatenation (OptIQ's no-chat-template fallback branch).
+      text = s.messages.map((m) => `[${m.role}]: ${m.content}`).join("\n\n");
+    } else if (s.text !== undefined) {
+      text = s.text;
+    } else {
+      continue;
+    }
+    const toks = tokenizer.encode(text);
+    for (const t of toks) allTokens.push(t);
+  }
+
+  if (allTokens.length < seqLen) {
+    throw new Error(
+      `Calibration mix produced only ${allTokens.length} tokens (< seqLen=${seqLen}). ` +
+        `Try a smaller seqLen or a larger mix.`,
+    );
+  }
+
+  // Chunk start offsets, shuffled deterministically (seeded), like OptIQ.
+  const starts: number[] = [];
+  for (let s = 0; s < allTokens.length - seqLen; s += seqLen) starts.push(s);
+  shuffleInPlace(starts, seed);
+
+  const chunks: number[][] = [];
+  for (const start of starts.slice(0, nSamples)) {
+    chunks.push(allTokens.slice(start, start + seqLen));
+  }
+
+  if (chunks.length === 0) {
+    throw new Error("Could not create any calibration sequences.");
+  }
+
+  return () => chunks;
+}
+
+/** Load raw samples from the mix spec ("optiq" → bundled, else a JSONL path,
+ *  else the built-in fallback corpus). */
+function loadSamples(mix: string): CalibrationSample[] {
+  let jsonlPath: string | null = null;
+  if (mix === "optiq") {
+    jsonlPath = bundledMixCandidates().find((p) => existsSync(p)) ?? null;
+  } else if (existsSync(mix)) {
+    jsonlPath = mix;
+  } else {
+    throw new Error(
+      `calibration mix not found at ${mix}. Use mix='optiq' (default) or a ` +
+        `path to a JSONL file.`,
+    );
+  }
+
+  if (jsonlPath) return parseJsonl(jsonlPath);
+
+  // Built-in fallback: synthesize one prose sample per corpus entry.
+  return FALLBACK_CORPUS.map((text) => ({ domain: "prose", text }));
+}
+
+/** Parse a JSONL file into CalibrationSample[] (one object per non-blank line). */
+function parseJsonl(path: string): CalibrationSample[] {
+  // Synchronous read so the caller can stay synchronous; the file is small.
+  const text = readFileSyncUtf8(path);
+  const out: CalibrationSample[] = [];
+  for (const line of text.split("\n")) {
+    const trimmed = line.trim();
+    if (trimmed) out.push(JSON.parse(trimmed) as CalibrationSample);
+  }
+  return out;
+}
+
+function readFileSyncUtf8(path: string): string {
+  // Bun's readFileSync via node:fs is available in this runtime.
+  // (Kept local so the rest of the module is import-light.)
+  const { readFileSync } = require("node:fs") as typeof import("node:fs");
+  return readFileSync(path, "utf8");
+}
+
+/** Deterministic Fisher–Yates shuffle (mulberry32 PRNG seeded by `seed`).
+ *  NOTE: this does NOT reproduce numpy's RandomState shuffle order — it only
+ *  needs to be deterministic for our own reproducibility; the exact sample
+ *  ordering does not affect the allocator's 1:1 parity (that test feeds
+ *  synthetic SensitivityResults directly). */
+function shuffleInPlace<T>(arr: T[], seed: number): void {
+  const rand = mulberry32(seed >>> 0);
+  for (let i = arr.length - 1; i > 0; i--) {
+    const j = Math.floor(rand() * (i + 1));
+    const tmp = arr[i]!;
+    arr[i] = arr[j]!;
+    arr[j] = tmp;
+  }
+}
+
+function mulberry32(a: number): () => number {
+  return () => {
+    a |= 0;
+    a = (a + 0x6d2b79f5) | 0;
+    let t = Math.imul(a ^ (a >>> 15), 1 | a);
+    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}

package/src/quantize/config-writer.ts

@@ -0,0 +1,132 @@
+// Config + sidecar writer for a quantized model directory.
+//
+// Produces the `config.json` block the loader (src/config.ts → parseQuantization)
+// reads back: a top-level `quantization` object AND a mirror `quantization_config`
+// (HF compatibility — config.ts reads `raw.quantization ?? raw.quantization_config`).
+// The block carries default {group_size, bits, mode} plus per-module overrides
+// keyed by module path (e.g. "model.embed_tokens") — exactly the OptiQ layout.
+//
+// Aux files (tokenizer, chat template, generation config, …) are copied through
+// verbatim so the output directory is a complete, loadable model snapshot.
+
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+
+/** A per-module quantization override. `false` means "left unquantized". */
+export type PerLayerEntry = { bits: number; groupSize: number } | false;
+
+/** The default quantization definition for the model. */
+export interface QuantDef {
+  bits: number;
+  groupSize: number;
+  mode?: string;
+}
+
+/** The serialized quantization block, ready to drop into config.json. */
+export interface QuantizationBlock {
+  group_size: number;
+  bits: number;
+  mode: string;
+  // plus per-module entries: { [modulePath]: {bits, group_size} | false }
+  [modulePath: string]: number | string | { bits: number; group_size: number } | false;
+}
+
+/**
+ * Build the `quantization` config object: a `{group_size, bits, mode}` default
+ * plus one entry per module in `perLayer`. A module mapped to `false` is
+ * serialized as `false` (mlx's "not quantized" convention); otherwise it
+ * becomes `{bits, group_size}`.
+ */
+export function buildQuantizationBlock(
+  def: QuantDef,
+  perLayer: Map<string, PerLayerEntry>,
+): QuantizationBlock {
+  const block: QuantizationBlock = {
+    group_size: def.groupSize,
+    bits: def.bits,
+    mode: def.mode ?? "affine",
+  };
+  for (const [path, entry] of perLayer) {
+    block[path] = entry === false ? false : { bits: entry.bits, group_size: entry.groupSize };
+  }
+  return block;
+}
+
+/** Aux files copied through verbatim from the source dir when present. */
+const AUX_FILES = [
+  "tokenizer.json",
+  "tokenizer_config.json",
+  "tokenizer.model",
+  "spiece.model",
+  "chat_template.jinja",
+  "generation_config.json",
+  "special_tokens_map.json",
+  "added_tokens.json",
+  "vocab.json",
+  "merges.txt",
+  "README.md",
+  "kv_config.json",
+];
+
+/** Metadata for the OptiQ-style sidecar describing the achieved quantization. */
+export interface OptiqMetadata {
+  method: string;
+  base_model: string;
+  bits: number;
+  group_size: number;
+  achieved_bpw: number;
+  per_layer_count: number;
+}
+
+/**
+ * Write `config.json` (deep copy of `srcConfigRaw` with both `quantization`
+ * and `quantization_config` set to `block`) into `outDir`, copy through any
+ * present aux files, and write `optiq_metadata.json` describing the run.
+ *
+ * `*.model` tokenizer files beyond the known names are also swept in so SPM /
+ * tiktoken sidecars survive the round-trip.
+ */
+export async function writeQuantizedConfig(
+  srcConfigRaw: Record<string, unknown>,
+  outDir: string,
+  block: QuantizationBlock,
+  opts: { srcDir?: string; optiq?: OptiqMetadata } = {},
+): Promise<void> {
+  // Deep copy so we never mutate the caller's config object.
+  const config = JSON.parse(JSON.stringify(srcConfigRaw)) as Record<string, unknown>;
+  config.quantization = block;
+  config.quantization_config = block;
+  await Bun.write(join(outDir, "config.json"), JSON.stringify(config, null, 2));
+
+  if (opts.srcDir) await copyAuxFiles(opts.srcDir, outDir);
+
+  if (opts.optiq) {
+    await Bun.write(
+      join(outDir, "optiq_metadata.json"),
+      JSON.stringify(opts.optiq, null, 2),
+    );
+  }
+}
+
+/** Copy known aux files (and any `*.model`) from src → out when present. */
+async function copyAuxFiles(srcDir: string, outDir: string): Promise<void> {
+  const seen = new Set<string>();
+  for (const f of AUX_FILES) {
+    const src = join(srcDir, f);
+    if (existsSync(src)) {
+      await Bun.write(join(outDir, f), Bun.file(src));
+      seen.add(f);
+    }
+  }
+  // Sweep any additional *.model tokenizer sidecars not in the known list.
+  try {
+    const { readdirSync } = await import("node:fs");
+    for (const f of readdirSync(srcDir)) {
+      if (f.endsWith(".model") && !seen.has(f)) {
+        await Bun.write(join(outDir, f), Bun.file(join(srcDir, f)));
+      }
+    }
+  } catch {
+    // src dir unreadable — aux copy is best-effort.
+  }
+}

package/src/quantize/index.ts

@@ -0,0 +1,57 @@
+// Public API for native model quantization. The server (and the job system)
+// import from here: quantize a model directory, write sharded safetensors,
+// build the config block, and resolve/inspect models for the wizard UI.
+
+export { quantizeModelDir, isQuantizable } from "./quantizer";
+export type {
+  QuantizeOptions,
+  QuantizeResult,
+  ProgressEvent,
+} from "./quantizer";
+
+export {
+  writeShardedSafetensors,
+  DEFAULT_SHARD_BYTES,
+} from "./safetensors-writer";
+export type {
+  NamedTensor,
+  ShardInfo,
+  SafetensorsIndex,
+  WriteResult,
+  WriteOpts,
+} from "./safetensors-writer";
+
+export {
+  buildQuantizationBlock,
+  writeQuantizedConfig,
+} from "./config-writer";
+export type {
+  QuantDef,
+  PerLayerEntry,
+  QuantizationBlock,
+  OptiqMetadata,
+} from "./config-writer";
+
+// Mixed-precision (OptiQ sensitivity + knapsack) surface.
+export {
+  optimizeMixedPrecision,
+  computeBpw,
+  klReduction,
+  UpgradeHeap,
+} from "./allocator";
+export type {
+  LayerQuantConfig,
+  OptimizationResult,
+  OptimizeMixedPrecisionOptions,
+  UpgradeEntry,
+} from "./allocator";
+
+export { analyzeSensitivityExact, klFromRef } from "./sensitivity";
+export type {
+  SensitivityResult,
+  AnalyzeSensitivityOptions,
+} from "./sensitivity";
+
+export { loadLlmCalibration } from "./calibration";
+
+export { quantizeRunner, inspectModel } from "./job";

package/src/quantize/job.ts

@@ -0,0 +1,136 @@
+// The quantize JobRunner: wires a submitted quantize job to quantizeModelDir().
+// Speaks the shared src/jobs/types.ts contract. The submit config carries the
+// source model (a registry query, a model_id, or an explicit src_dir), the
+// output dir, and the target bits / group size.
+//
+// Registers itself with the jobs runner registry on import so both the
+// in-process submit path and the subprocess job-entry resolver (which
+// dynamic-imports this module via KIND_MODULES) find the "quantize" runner.
+
+import { statSync } from "node:fs";
+import type { Emit, JobRunner } from "../jobs/types";
+import { registerRunner } from "../jobs/runner";
+import { Registry } from "../registry";
+import { loadModelConfig } from "../config";
+import { quantizeModelDir, type QuantizeOptions } from "./quantizer";
+
+/** Resolve the source model directory from a job config. Accepts an explicit
+ *  filesystem path (src_dir) or a registry query / model id (model_id). */
+function resolveSrcDir(config: Record<string, unknown>): string {
+  const srcDir = config.src_dir as string | undefined;
+  if (srcDir) return srcDir;
+  const modelId = config.model_id as string | undefined;
+  if (!modelId) throw new Error("quantize job: missing src_dir or model_id");
+  // A path-looking model_id is taken literally; otherwise resolve via registry.
+  if (modelId.includes("/") && existsSyncSafe(modelId)) return modelId;
+  const reg = new Registry();
+  try {
+    return reg.resolve(modelId).path;
+  } finally {
+    reg.close();
+  }
+}
+
+function existsSyncSafe(p: string): boolean {
+  try {
+    statSync(p);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export const quantizeRunner: JobRunner = async (emit: Emit, config) => {
+  const outDir = String(config.out_dir ?? "");
+  if (!outDir) throw new Error("quantize job: missing out_dir");
+
+  const bits = Number(config.bits ?? 4) as 4 | 8;
+  const groupSize = Number(config.group_size ?? 64) as 32 | 64;
+  if (bits !== 4 && bits !== 8) throw new Error(`quantize job: bits must be 4 or 8 (got ${bits})`);
+  if (groupSize !== 32 && groupSize !== 64)
+    throw new Error(`quantize job: group_size must be 32 or 64 (got ${groupSize})`);
+
+  const srcDir = resolveSrcDir(config);
+
+  emit({
+    type: "stage",
+    stage: "starting",
+    progress: 0.01,
+    message: `Quantizing ${srcDir} → ${bits}-bit (g${groupSize})`,
+  });
+
+  const opts: QuantizeOptions = { bits, groupSize, mode: String(config.mode ?? "affine") };
+
+  const r = await quantizeModelDir(srcDir, outDir, opts, (e) =>
+    emit({ type: "stage", stage: e.stage, progress: e.progress, message: e.message }),
+  );
+
+  emit({
+    type: "stage",
+    stage: "done",
+    progress: 1,
+    message: `Quantized ${r.nQuantized} modules (${r.achievedBpw.toFixed(2)} bpw)`,
+    output_dir: r.outDir,
+  });
+
+  return { outputPath: r.outDir };
+};
+
+/** Lightweight model inspection for the `/api/quantize/inspect` route: report
+ *  whether a model is quantizable and its on-disk size, from the registry +
+ *  config (never touches tensor bytes). */
+export async function inspectModel(model_id: string): Promise<{
+  ok: boolean;
+  model_id: string;
+  arch: string | null;
+  support: boolean;
+  size_gb: number;
+  error?: string;
+}> {
+  try {
+    let path = model_id;
+    let arch: string | null = null;
+    let sizeBytes = 0;
+
+    if (!(model_id.includes("/") && existsSyncSafe(model_id))) {
+      const reg = new Registry();
+      try {
+        const rec = reg.resolve(model_id);
+        path = rec.path;
+        arch = rec.modelType;
+        sizeBytes = rec.sizeBytes;
+      } finally {
+        reg.close();
+      }
+    }
+
+    const config = await loadModelConfig(path);
+    arch = arch ?? config.modelType ?? (config.architectures[0] ?? null);
+    // Supported = a text architecture this quantizer can walk. v1 quantizes
+    // any model whose weights are plain 2D Linear/embedding tensors; we treat
+    // every loadable text config as supported and let eligibility filter
+    // per-tensor at quantize time.
+    const support = config.text.numHiddenLayers > 0;
+
+    return {
+      ok: true,
+      model_id,
+      arch,
+      support,
+      size_gb: sizeBytes / (1 << 30),
+    };
+  } catch (e) {
+    return {
+      ok: false,
+      model_id,
+      arch: null,
+      support: false,
+      size_gb: 0,
+      error: e instanceof Error ? e.message : String(e),
+    };
+  }
+}
+
+// Self-register on import: `kind: "quantize"` resolves both for in-process
+// submit and for the subprocess job-entry (KIND_MODULES → ../quantize/job.ts).
+registerRunner("quantize", quantizeRunner);