@mlx-bun/quantize 0.0.0-stage → 1.0.0

package/src/quantizer.ts

@@ -0,0 +1,691 @@
+import type { QuantizedLinear } from "@mlx-bun/inference/layers";
+// Native model quantizer: walk a model directory's weights, quantize every
+// eligible Linear / embedding / lm_head weight to uniform affine N-bit, pass
+// everything else through unchanged, and write a loadable quantized snapshot.
+//
+// No Python. The whole pipeline (load → dequantize-if-needed → cast bf16 →
+// mx.quantize → save) runs through the mlx-c FFI on the CPU stream where the
+// safetensors loader lives.
+//
+// Memory discipline: one module is materialized, quantized, written, and
+// disposed at a time; the allocator cache is cleared every CACHE_CLEAR_EVERY
+// modules so a large model never holds more than a few resident tensors.
+//
+// Uniform affine (every eligible weight gets the same bits/group_size) is the
+// default. When `opts.targetBpw` is set, the OptiQ sensitivity-driven
+// mixed-precision path runs: calibration → per-layer KL sensitivity → greedy
+// knapsack → a heterogeneous per-module bit allocation fed to the same writer.
+
+import { join } from "node:path";
+import { mkdirSync } from "node:fs";
+import { clearCache, Dtype } from "@mlx-bun/mlx/ffi";
+import { cpuStream, type MlxArray } from "@mlx-bun/mlx/array";
+import { quantize, dequantize } from "@mlx-bun/mlx/ops";
+import { Weights } from "@mlx-bun/inference";
+import { loadModelConfig, quantFor, type QuantizationConfig } from "@mlx-bun/inference/artifacts";
+import { createModel } from "@mlx-bun/inference";
+import { loadTokenizer } from "@mlx-bun/inference/input";
+import {
+  writeShardedSafetensors,
+  type NamedTensor,
+  type WriteResult,
+} from "@mlx-bun/inference/artifacts";
+import {
+  buildQuantizationBlock,
+  writeQuantizedConfig,
+  type PerLayerEntry,
+} from "./config-writer";
+import { loadLlmCalibration } from "./calibration";
+import { analyzeSensitivityExact, type SensitivityResult } from "./sensitivity";
+import { optimizeMixedPrecision } from "./allocator";
+import type {
+  WeightTransform,
+  WeightTransformContext,
+  WeightTransformPlan,
+} from "./weight-transform";
+import { writeAtomicDirectory } from "./atomic-output";
+
+/** Quantize the allocator cache this often (in modules processed). */
+const CACHE_CLEAR_EVERY = 16;
+
+/** Bits used per group for affine scale + bias (one bf16 each = 32 bits). */
+const AFFINE_GROUP_OVERHEAD_BITS = 32;
+
+/** The floating dtypes `mlx_lm.convert --dtype` accepts. */
+export type ConvertDtype = "float16" | "bfloat16" | "float32";
+export const CONVERT_DTYPES: readonly ConvertDtype[] = ["float16", "bfloat16", "float32"];
+const DTYPE_OF: Record<ConvertDtype, Dtype> = { float16: Dtype.float16, bfloat16: Dtype.bfloat16, float32: Dtype.float32 };
+const FLOATING = new Set<Dtype>([Dtype.float16, Dtype.bfloat16, Dtype.float32, Dtype.float64]);
+
+/** mlx-lm's per-model `cast_predicate`s, as one name rule: router/expert bias
+ * and state-space decay parameters stay in their source dtype. */
+const KEEP_DTYPE = /e_score_correction_bias|expert_bias|A_log$|dt_bias$/;
+
+/** Cast one floating tensor to `dtype` (lazy); anything else passes through. */
+export function castTensor(name: string, array: MlxArray, dtype: ConvertDtype | undefined): MlxArray {
+  if (!dtype || KEEP_DTYPE.test(name) || !FLOATING.has(array.dtype) || array.dtype === DTYPE_OF[dtype]) return array;
+  return array.astype(DTYPE_OF[dtype], cpuStream);
+}
+
+export interface QuantizeOptions {
+  bits: 4 | 8;
+  groupSize: 32 | 64;
+  /** `mlx_lm.convert --dtype`: dtype of the non-quantized tensors and of the
+   *  quantization scales/biases. Absent keeps bf16 scales/biases and the
+   *  source dtype for everything else. */
+  dtype?: ConvertDtype;
+  /** Quantization scheme; v1 only exercises "affine". */
+  mode?: string;
+  /** Extra per-module gate on top of the shape-eligibility check: return
+   *  false to keep a module full-precision even though its shape qualifies
+   *  (e.g. the DeepSpec drafter's confidence_head — threshold-sensitive and
+   *  tiny). Excluded modules are recorded as `false` in the quantization
+   *  block (mlx's "not quantized" convention). Shape-ineligible modules
+   *  never reach this. */
+  quantizePredicate?: (base: string, fullShape: number[]) => boolean;
+  /** Optional offline basis transform applied before quantization. */
+  weightTransform?: WeightTransform;
+
+  // --- mixed-precision (OptiQ sensitivity + knapsack) -------------------
+  // Setting targetBpw switches quantizeModelDir to the mixed path: it loads the
+  // model, runs calibration + per-layer KL sensitivity, knapsacks a per-module
+  // bit allocation, and feeds that heterogeneous map to the same writer.
+  /** Target bits-per-weight for sensitivity+knapsack mixed precision. */
+  targetBpw?: number;
+  /** Reference dtype for sensitivity scoring (e.g. "bf16"). Informational
+   *  only — the native path dequantizes the running model's own weights as
+   *  the self-contained bf16 source. */
+  reference?: string;
+  /** Candidate bit-widths the knapsack may choose from (e.g. [4, 8]). */
+  candidateBits?: number[];
+  /** Calibration data mixture spec ("optiq" or a JSONL path). */
+  calibrationMix?: string;
+  /** Number of calibration samples (forward passes per layer×bit probe). */
+  nCalibration?: number;
+  /** Dependency used to provide the loadable model for sensitivity sweeps. */
+  probeSource?: ProbeSource;
+}
+
+export interface PreparedProbe {
+  modelDir: string;
+  dispose(): void | Promise<void>;
+}
+
+/** Supplies a loadable sensitivity model without re-entering the top-level
+ *  quantization orchestrator. Tests and callers may inject an existing probe. */
+export interface ProbeSource {
+  prepare(
+    srcDir: string,
+    opts: QuantizeOptions,
+    onProgress?: (e: ProgressEvent) => void,
+  ): Promise<PreparedProbe>;
+}
+
+export interface ProgressEvent {
+  stage: string;
+  message: string;
+  /** 0..1 fraction of modules processed. */
+  progress: number;
+}
+
+export interface QuantizeResult {
+  outDir: string;
+  /** Achieved bits-per-weight over the quantized parameters (incl overhead). */
+  achievedBpw: number;
+  /** Number of modules quantized. */
+  nQuantized: number;
+  /** Sharding / size details from the writer. */
+  write: WriteResult;
+}
+
+type Suffix = "weight" | "scales" | "biases";
+
+/** Split a tensor name into (modulePath, suffix) for `.weight|.scales|.biases`. */
+function splitName(name: string): { base: string; suffix: Suffix } | null {
+  const m = /^(.*)\.(weight|scales|biases)$/.exec(name);
+  if (!m) return null;
+  return { base: m[1]!, suffix: m[2] as Suffix };
+}
+
+/**
+ * Eligibility predicate (port of mlx-lm's default class_predicate): a module
+ * is quantizable iff it has a 2D `.weight` whose last dim is divisible by the
+ * group size. This captures Linear projections, the token embedding, and the
+ * (untied) lm_head, and excludes norms/biases/1D tensors.
+ */
+export function isQuantizable(weightShape: number[], groupSize: number): boolean {
+  return weightShape.length === 2 && weightShape[weightShape.length - 1]! % groupSize === 0;
+}
+
+/**
+ * Quantize a model directory and write a loadable snapshot to `outDir`.
+ *
+ * Uniform affine N-bit by default. When `opts.targetBpw` is set, runs the
+ * OptiQ mixed-precision path (calibration → sensitivity → knapsack) to choose
+ * per-module bit-widths before writing.
+ */
+export async function quantizeModelDir(
+  srcDir: string,
+  outDir: string,
+  opts: QuantizeOptions,
+  onProgress?: (e: ProgressEvent) => void,
+): Promise<QuantizeResult> {
+  // Mixed-precision path: derive a per-module bit allocation first, then write.
+  let perLayerBits: Map<string, number> | undefined;
+  let mixedMeta: MixedMeta | undefined;
+  if (opts.targetBpw !== undefined) {
+    const alloc = await computeMixedAllocation(srcDir, opts, onProgress);
+    perLayerBits = alloc.perLayerBits;
+    mixedMeta = alloc.meta;
+  }
+
+  return writeAtomicDirectory(outDir, (stagingDir) =>
+    writeQuantizedModelDir(
+      srcDir,
+      stagingDir,
+      opts,
+      perLayerBits,
+      mixedMeta,
+      onProgress,
+      outDir,
+    ),
+  );
+}
+
+/** Lower-level writer used by the public orchestrator and its probe source.
+ *  Deliberately does not perform sensitivity/allocation or call back into
+ *  quantizeModelDir. */
+async function writeQuantizedModelDir(
+  srcDir: string,
+  outDir: string,
+  opts: QuantizeOptions,
+  perLayerBits?: Map<string, number>,
+  mixedMeta?: MixedMeta,
+  onProgress?: (e: ProgressEvent) => void,
+  publishedOutDir = outDir,
+): Promise<QuantizeResult> {
+  const groupSize = opts.groupSize;
+  const mode = opts.mode ?? "affine";
+
+  const bits = opts.bits;
+  // Dtype of the quantization input (so of its scales/biases): bf16 unless --dtype names one.
+  const castDtype = opts.dtype ? DTYPE_OF[opts.dtype] : Dtype.bfloat16;
+  const cast = (tensor: NamedTensor): NamedTensor => ({ name: tensor.name, array: castTensor(tensor.name, tensor.array, opts.dtype) });
+
+  mkdirSync(outDir, { recursive: true });
+
+  const progress = (stage: string, message: string, frac: number) =>
+    onProgress?.({ stage, message, progress: frac });
+
+  progress("loading", `Reading ${srcDir}`, 0);
+
+  const config = await loadModelConfig(srcDir);
+  const srcQuant: QuantizationConfig | null = config.quantization;
+  const weights = await Weights.open(srcDir);
+  let transformPlan: WeightTransformPlan | undefined;
+  let transformContext: WeightTransformContext | undefined;
+  try {
+    if (opts.weightTransform && srcQuant)
+      throw new Error(
+        `weight transform ${opts.weightTransform.id} requires full-precision source weights`,
+      );
+    transformPlan = opts.weightTransform?.plan(weights.tensorNames, config);
+    if (transformPlan) {
+      validateTransformPlan(transformPlan, weights.tensorNames);
+      transformContext = opts.weightTransform!.createContext(weights, transformPlan);
+      progress("transforming", `Planned ${transformPlan.id}`, 0);
+    }
+  } catch (error) {
+    transformContext?.dispose();
+    weights.dispose();
+    throw error;
+  }
+
+  // Preserve the source tensor order so the output diff is mlx-lm-shaped.
+  // tensorNames is sorted; we walk modules in that order. For each module's
+  // `.weight` we decide quantize / pass-through, then emit its tensors. We
+  // skip the source `.scales`/`.biases` of already-quantized modules (they're
+  // regenerated) and pass through everything else verbatim.
+  const names = transformPlan?.outputNames ?? weights.tensorNames;
+
+  // Group names by module base so we can detect "already quantized" (has a
+  // sibling `.scales`) and process a module atomically.
+  const moduleSuffixes = new Map<string, Set<Suffix>>();
+  const passthroughNames: string[] = []; // non-(weight/scales/biases) tensors
+  for (const name of names) {
+    const sp = splitName(name);
+    if (!sp) {
+      passthroughNames.push(name);
+      continue;
+    }
+    let set = moduleSuffixes.get(sp.base);
+    if (!set) {
+      set = new Set();
+      moduleSuffixes.set(sp.base, set);
+    }
+    set.add(sp.suffix);
+  }
+
+  // Decide the per-module plan in source order.
+  const moduleBases = [...moduleSuffixes.keys()];
+  const totalModules = moduleBases.length;
+
+  const out: NamedTensor[] = [];
+  const perLayer = new Map<string, PerLayerEntry>();
+  let nQuantized = 0;
+  let quantizedParams = 0;
+  let quantizedBits = 0;
+  let processed = 0;
+
+  try {
+    for (const base of moduleBases) {
+      const suffixes = moduleSuffixes.get(base)!;
+      const alreadyQuant = suffixes.has("scales");
+      const hasWeight = suffixes.has("weight");
+
+      if (!hasWeight) {
+        // Module with only scales/biases and no weight — shouldn't happen, but
+        // pass any present tensors through unchanged.
+        for (const suf of suffixes)
+          out.push(cast(outputTensor(weights, transformPlan, transformContext, `${base}.${suf}`)));
+        bumpProgress();
+        continue;
+      }
+
+      // Materialize the full-precision weight for this module.
+      const weightName = `${base}.weight`;
+      let fullWeight = outputArray(weights, transformPlan, transformContext, weightName);
+      const shape = fullWeight.shape;
+
+      // Determine the *original* (full-precision) shape: if already quantized,
+      // the packed weight's last dim is compressed — recover the real width
+      // from the scales (groups × srcGroupSize) when dequantizing.
+      let materialized: MlxArray | null = null;
+
+      if (alreadyQuant) {
+        const srcSpec = quantFor(srcQuant, base);
+        if (!srcSpec) {
+          // Has scales but config says unquantized — fall back to passthrough.
+          for (const suf of suffixes)
+            out.push(cast(outputTensor(weights, transformPlan, transformContext, `${base}.${suf}`)));
+          bumpProgress();
+          continue;
+        }
+        const scales = weights.tensor(`${base}.scales`);
+        const biases = suffixes.has("biases") ? weights.tensor(`${base}.biases`) : null;
+        materialized = dequantize(
+          fullWeight, scales, biases,
+          { bits: srcSpec.bits, groupSize: srcSpec.groupSize, mode: srcSpec.mode },
+          cpuStream,
+        );
+        fullWeight = materialized;
+      }
+
+      // Recover the full-precision 2D shape for eligibility + param counting.
+      // For a freshly-loaded bf16 weight this is just `shape`; for a
+      // re-quantized weight it is the dequantized array's shape.
+      const fullShape = alreadyQuant ? materialized!.shape : shape;
+
+      const predicateExcluded =
+        isQuantizable(fullShape, groupSize) &&
+        opts.quantizePredicate !== undefined &&
+        !opts.quantizePredicate(base, fullShape);
+      if (!isQuantizable(fullShape, groupSize) || predicateExcluded) {
+        // Not eligible: pass the full-precision weight through unchanged
+        // (re-materialized to bf16 if it had been quantized).
+        if (alreadyQuant) {
+          const dense = materialized!.astype(castDtype, cpuStream);
+          out.push({ name: weightName, array: dense });
+          materialized!.dispose();
+        } else {
+          out.push(cast({ name: weightName, array: fullWeight }));
+        }
+        // Policy-excluded (shape-eligible) modules get an explicit `false`
+        // in the quantization block — the loader must not apply the block's
+        // default spec to them (mlx's "not quantized" convention).
+        if (predicateExcluded) perLayer.set(base, false);
+        bumpProgress();
+        continue;
+      }
+
+      // Resolve this module's bit-width: the allocator's per-layer choice in
+      // mixed mode, else the uniform `bits`. The allocation map is keyed by
+      // module base path (e.g. "model.layers.0.self_attn.q_proj").
+      const moduleBits = perLayerBits?.get(base) ?? bits;
+
+      // Cast before quantizing (mx.quantize expects a float weight; the input
+      // dtype is the scales/biases dtype on disk: bf16, or --dtype).
+      const bf16 = fullWeight.astype(castDtype, cpuStream);
+      if (materialized) materialized.dispose();
+      else if (transformContext) fullWeight.dispose();
+
+      const q = quantize(bf16, groupSize, moduleBits, mode, cpuStream);
+      bf16.dispose();
+
+      out.push({ name: `${base}.weight`, array: q.packed });
+      out.push({ name: `${base}.scales`, array: q.scales });
+      out.push({ name: `${base}.biases`, array: q.biases });
+
+      // Account: every element of the original 2D weight is now `moduleBits`
+      // bits, plus 32 bits (bf16 scale + bf16 bias) per group of `groupSize`.
+      const params = fullShape.reduce((a, b) => a * b, 1);
+      quantizedParams += params;
+      quantizedBits += params * moduleBits + (params / groupSize) * AFFINE_GROUP_OVERHEAD_BITS;
+      perLayer.set(base, { bits: moduleBits, groupSize });
+      nQuantized++;
+
+      bumpProgress();
+    }
+
+    // Pass through non-weight/scale/bias tensors (e.g. anything unusual).
+    for (const name of passthroughNames)
+      out.push(cast(outputTensor(weights, transformPlan, transformContext, name)));
+
+    progress("writing", `Writing ${out.length} tensors to ${publishedOutDir}`, 1);
+    const write = writeShardedSafetensors(outDir, out);
+
+    const achievedBpw = quantizedParams > 0 ? quantizedBits / quantizedParams : 0;
+
+    // Default block bits: the uniform `bits` normally; in mixed mode use the
+    // lowest allocated bit-width so the loader's per-module overrides describe
+    // every deviation above the floor.
+    let defaultBits: number = bits;
+    if (perLayerBits && perLayerBits.size > 0) {
+      defaultBits = Math.min(...perLayerBits.values());
+    }
+
+    // Build + write the config block. For uniform mode, let the config writer
+    // emit its standard optiq_metadata.json sidecar. For mixed mode we write a
+    // richer OptiQ-style sidecar ourselves below (method "mixed_precision").
+    const block = buildQuantizationBlock({ bits: defaultBits, groupSize, mode }, perLayer);
+    const outputConfig = transformedConfig(config.raw, transformPlan);
+    const weightTransforms = transformPlan ? [transformPlan.metadata] : undefined;
+    await writeQuantizedConfig(outputConfig, outDir, block, {
+      srcDir,
+      optiq: mixedMeta
+        ? undefined
+        : {
+            method: "uniform_affine",
+            base_model: srcDir,
+            bits,
+            group_size: groupSize,
+            achieved_bpw: achievedBpw,
+            per_layer_count: nQuantized,
+            weight_transforms: weightTransforms,
+          },
+    });
+
+    if (mixedMeta) {
+      await Bun.write(
+        join(outDir, "optiq_metadata.json"),
+        JSON.stringify(
+          {
+            method: "mixed_precision",
+            base_model: srcDir,
+            bits: defaultBits,
+            group_size: groupSize,
+            target_bpw: mixedMeta.targetBpw,
+            achieved_bpw: achievedBpw,
+            candidate_bits: mixedMeta.candidateBits,
+            n_high: mixedMeta.nHigh,
+            n_low: mixedMeta.nLow,
+            per_layer_count: nQuantized,
+            weight_transforms: weightTransforms,
+            per_layer: Object.fromEntries(
+              [...perLayer].map(([p, e]) => [
+                p,
+                e === false ? false : { bits: e.bits, group_size: e.groupSize },
+              ]),
+            ),
+          },
+          null,
+          2,
+        ),
+      );
+    }
+
+    progress("done", `Quantized ${nQuantized} modules → ${publishedOutDir}`, 1);
+    return { outDir: publishedOutDir, achievedBpw, nQuantized, write };
+  } finally {
+    // Dispose every emitted array (their bytes are now on disk) and the source.
+    for (const t of out) {
+      try { t.array.dispose(); } catch {}
+    }
+    transformContext?.dispose();
+    weights.dispose();
+    clearCache();
+  }
+
+  function bumpProgress(): void {
+    processed++;
+    if (processed % CACHE_CLEAR_EVERY === 0) clearCache();
+    progress(
+      "quantizing",
+      `Module ${processed}/${totalModules}`,
+      totalModules > 0 ? processed / totalModules : 1,
+    );
+  }
+}
+
+/** Pass-through tensor: the lazy mlx array straight from the source weights. */
+function outputTensor(
+  weights: Weights,
+  plan: WeightTransformPlan | undefined,
+  context: WeightTransformContext | undefined,
+  name: string,
+): NamedTensor {
+  return { name, array: outputArray(weights, plan, context, name) };
+}
+
+function outputArray(
+  weights: Weights,
+  plan: WeightTransformPlan | undefined,
+  context: WeightTransformContext | undefined,
+  name: string,
+): MlxArray {
+  if (!plan || !context) return weights.tensor(name);
+  const sourceName = plan.sourceByOutput.get(name);
+  if (!sourceName) throw new Error(`weight transform ${plan.id}: no source for ${name}`);
+  return context.apply(name, weights.tensor(sourceName));
+}
+
+function validateTransformPlan(
+  plan: WeightTransformPlan,
+  sourceNames: readonly string[],
+): void {
+  const seen = new Set<string>();
+  const sources = new Set(sourceNames);
+  for (const name of plan.outputNames) {
+    if (seen.has(name)) throw new Error(`weight transform ${plan.id}: duplicate output ${name}`);
+    seen.add(name);
+    const sourceName = plan.sourceByOutput.get(name);
+    if (!sourceName)
+      throw new Error(`weight transform ${plan.id}: no source for ${name}`);
+    if (!sources.has(sourceName))
+      throw new Error(`weight transform ${plan.id}: missing source tensor ${sourceName}`);
+  }
+}
+
+function transformedConfig(
+  raw: Record<string, unknown>,
+  plan: WeightTransformPlan | undefined,
+): Record<string, unknown> {
+  if (!plan?.untieWordEmbeddings) return raw;
+  const out: Record<string, unknown> = { ...raw, tie_word_embeddings: false };
+  if (raw.text_config && typeof raw.text_config === "object") {
+    out.text_config = {
+      ...(raw.text_config as Record<string, unknown>),
+      tie_word_embeddings: false,
+    };
+  }
+  return out;
+}
+
+/** Summary of a mixed-precision allocation, for the sidecar + return value. */
+interface MixedMeta {
+  targetBpw: number;
+  candidateBits: number[];
+  nHigh: number;
+  nLow: number;
+  /** Knapsack's own BPW estimate over the analyzed layers (params-weighted,
+   *  no group overhead) — distinct from the writer's achievedBpw. */
+  allocBpw: number;
+}
+
+const defaultProbeSource: ProbeSource = {
+  async prepare(srcDir, opts, onProgress) {
+    const config = await loadModelConfig(srcDir);
+    if (config.quantization) return { modelDir: srcDir, dispose() {} };
+
+    const { mkdtempSync, rmSync } = await import("node:fs");
+    const { tmpdir } = await import("node:os");
+    const tempDir = mkdtempSync(join(tmpdir(), "mlx-bun-qbaseline-"));
+    onProgress?.({
+      stage: "sensitivity",
+      message: "Unquantized source — building an 8-bit baseline for the sweep",
+      progress: 0,
+    });
+    try {
+      await writeQuantizedModelDir(
+        srcDir,
+        tempDir,
+        {
+          ...opts,
+          bits: 8,
+          targetBpw: undefined,
+          probeSource: undefined,
+        },
+      );
+    } catch (error) {
+      rmSync(tempDir, { recursive: true, force: true });
+      throw error;
+    }
+    return {
+      modelDir: tempDir,
+      dispose: () => rmSync(tempDir, { recursive: true, force: true }),
+    };
+  },
+};
+
+/** Acquire/use/release helper that pins the injected ProbeSource lifecycle. */
+export async function withPreparedProbe<T>(
+  source: ProbeSource,
+  srcDir: string,
+  opts: QuantizeOptions,
+  run: (modelDir: string) => Promise<T>,
+  onProgress?: (e: ProgressEvent) => void,
+): Promise<T> {
+  const prepared = await source.prepare(srcDir, opts, onProgress);
+  try {
+    return await run(prepared.modelDir);
+  } finally {
+    await prepared.dispose();
+  }
+}
+
+/**
+ * Run the OptiQ mixed-precision pipeline for `srcDir`:
+ *   load model + tokenizer → calibration → per-layer KL sensitivity →
+ *   greedy knapsack → per-module bit allocation.
+ *
+ * Returns a `Map<modulePath, bits>` keyed by weight-file module base path
+ * (the same key the writer loop uses) plus a metadata summary. Modules not in
+ * the map fall back to the uniform `opts.bits` at write time.
+ */
+async function computeMixedAllocation(
+  srcDir: string,
+  opts: QuantizeOptions,
+  onProgress?: (e: ProgressEvent) => void,
+): Promise<{ perLayerBits: Map<string, number>; meta: MixedMeta }> {
+  const targetBpw = opts.targetBpw!;
+  const groupSize = opts.groupSize;
+  const candidateBits = (opts.candidateBits ?? [4, 8]).slice().sort((a, b) => a - b);
+  const nCalibration = opts.nCalibration ?? 2;
+  const calibrationMix = opts.calibrationMix ?? "optiq";
+
+  const progress = (stage: string, message: string, frac: number) =>
+    onProgress?.({ stage, message, progress: frac });
+
+  progress("sensitivity", "Loading model for sensitivity analysis", 0.0);
+
+  // The source dependency owns preparation and cleanup. The default builds a
+  // temporary uniform 8-bit baseline from bf16 through the lower-level writer;
+  // an injected source can point at a prebuilt probe without orchestrator
+  // recursion.
+  return withPreparedProbe(
+    opts.probeSource ?? defaultProbeSource,
+    srcDir,
+    opts,
+    async (probeDir) => {
+      const config = await loadModelConfig(probeDir);
+      const weights = await Weights.open(probeDir);
+      try {
+        const model = createModel(weights, config);
+        const tokenizer = await loadTokenizer(srcDir);
+
+        // Calibration: short sequences keep per-probe forward memory modest.
+        progress("sensitivity", "Building calibration set", 0.02);
+        const calibrationFn = loadLlmCalibration(tokenizer, {
+          nSamples: nCalibration,
+          seqLen: 128,
+          mix: calibrationMix,
+        });
+        const calIds = calibrationFn();
+
+        // Quantizable linears to probe: the model's LoRA-target linears
+        // (attention + MLP projections), keyed by full module path.
+        const layers = new Map([...model.loraTargets()].filter((entry): entry is [string, QuantizedLinear] =>
+          "exchangeQuantizedState" in entry[1]));
+
+        progress(
+          "sensitivity",
+          `Probing ${layers.size} layers × ${candidateBits.length} bit-widths`,
+          0.05,
+        );
+        const sensitivity: SensitivityResult[] = analyzeSensitivityExact(
+          model,
+          layers,
+          calIds,
+          { candidateBits, groupSize },
+          (done, total, layerName) =>
+            progress(
+              "sensitivity",
+              `Sensitivity ${done}/${total}: ${layerName}`,
+              0.05 + 0.85 * (done / Math.max(total, 1)),
+            ),
+        );
+
+        // Greedy knapsack → per-layer bit allocation.
+        progress("allocating", "Optimizing per-layer bit allocation", 0.92);
+        const opt = optimizeMixedPrecision(sensitivity, {
+          targetBpw,
+          candidateBits,
+          groupSize,
+        });
+
+        const perLayerBits = new Map<string, number>();
+        for (const c of opt.configs) perLayerBits.set(c.layerName, c.bits);
+
+        const meta: MixedMeta = {
+          targetBpw,
+          candidateBits,
+          nHigh: opt.nHighBits,
+          nLow: opt.nLowBits,
+          allocBpw: opt.achievedBpw,
+        };
+
+        progress(
+          "allocating",
+          `Allocated ${opt.nHighBits} high / ${opt.nLowBits} low (alloc bpw ${opt.achievedBpw.toFixed(2)})`,
+          0.95,
+        );
+
+        return { perLayerBits, meta };
+      } finally {
+        weights.dispose();
+        clearCache();
+      }
+    },
+    onProgress,
+  );
+}