Skip to content

RFC: Streaming conversion - quantize checkpoints larger than local disk in mlx_lm.convert #1842

Description

@azamamirza

Problem

mlx_lm.convert requires the full source checkpoint on disk before conversion starts. The snapshot download completes first, then load_model(lazy=True) -> quantize -> save_model. For frontier-size checkpoints, this makes conversion impossible on otherwise perfectly
capable machines.

A 532 GB BF16 MoE checkpoint quantizes to ~150 GB, but converting it today requires ~700 GB of free space, even though the machine only ever needs the output plus one source shard at a time. Lazy loading already avoids the memory problem (weights evaluate shard by
shard).

This RFC is about the disk problem.

Approach

Stream the conversion! Download one source shard at a time, sanitize and quantize its tensors, append them to mlx-community-style output shards (~5 GB), and then delete the source blob before fetching the next one. The full BF16 checkpoint never exists on disk.

Important Considerations

  1. Cross-shard fused groups must not force source coexistence. Models whose sanitize() fuses tensors (e.g. stacking per-expert weights, or fusing q/k/v projections) can have one fused group scattered across several source shards. Pending parts are mx.eval 'd (copied out of the source mmap) before the source blob is deleted, and spilled to a state-*-pending.safetensors file at each shard boundary. Buffering + spill is far cheaper than keeping source spans on disk (observed pending peak: ~1.2 GB against 17 GB source shards).
  2. A manifest write is the single commit point (atomic tmp+rename), recording per-source-shard status/sha256 and flushed output shards. Order per shard: flush full output shards -> write state spills -> write manifest -> delete old spills -> delete source blob. A kill at any point resumes byte-identically (verified with a mid-run simulated crash).
  3. Disk guard before each fetch: abort resume-ready if free < min_free + estimated shard size, so the tool degrades into "free some space and re-run" instead of corrupting anything.
  4. Dry-run plans the whole conversion through the real sanitize path on lazy zero tensors, reporting output size and the disk high-water mark without downloading anything.

Real-world Example

Implemented and validated end-to-end for a 532 GB BF16 multimodal MoE checkpoint (Inkling-Small), quantizing routed experts to 4-bit with BF16 passthrough elsewhere, on a machine with under 190 GB free:

  • mxfp4 recipe: 31 output shards, 148.8 GB output, 168.7 GB disk max disk usage (output + one ~17.4 GB source shard + spills).
  • affine 4-bit: 156.9 GB output, 175.7 GB max disk usage.
  • Kill/resume mid-run reproduces byte-identical output.
  • Output is standard mlx-community layout (sanitized names, top-level quantization block, tokenizer/processor files). Loads with stock
    mlx-lm conventions.

Proposed Addition to Upstream

An opt-in flag on the existing CLI: mlx_lm.convert --stream [--min-free-gb N], sharing the existingquantization/predicate machinery:

  • The generic pieces (shard-ordered fetch via the index, pending-group buffering around sanitize(), manifest/resume, disk guard, output-shard writer) are model-agnostic and would live in mlx_lm/convert.py (or a sibling module).
  • Per-model sanitize() stays as the only model-specific hook, exactly as today; the streaming path just calls it incrementally.
  • Non-streaming behavior is untouched; --stream trades wall-clock (serialized download/convert) for disk.

Happy to adapt the implementation to whatever shape maintainers prefer and submit it as a PR.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions