Stream large checkpoints from disk during quantization#2061
Open
aquilarubra wants to merge 8 commits into
Open
Stream large checkpoints from disk during quantization#2061aquilarubra wants to merge 8 commits into
aquilarubra wants to merge 8 commits into
Conversation
This was referenced Jul 19, 2026
aquilarubra
force-pushed
the
pr/disk-streaming-core
branch
from
July 19, 2026 14:20
1b6bfec to
b8674e9
Compare
Two opt-in switches used by the disk-streaming and resumability work that follows in later commits. Default off (unset) preserves upstream behavior exactly. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
New auto_round/utils/disk_stream_util.py, no upstream equivalent.
Lazy, mmap-backed reads of individual tensors by name straight from a
checkpoint's safetensors shards, plus meta<->real materialize/free for
a whole module. Also provides build_meta_model() (a meta skeleton +
tokenizer + SafetensorsIndex, narrower than llm_load_model -- no
bagel/glm/mxfp4/HPU special-casing) and materialize_non_block_params()
(real-loads everything outside the decoder blocks: embeddings/
lm_head/final norm).
Both materialize functions pass dtype=values[full_name].dtype
explicitly to accelerate's set_module_tensor_to_device(): without it,
accelerate casts real checkpoint data to whatever dtype the meta
skeleton's parameter happened to declare, not the checkpoint's real
dtype -- silently wrong for any module built without a matching dtype
context (e.g. an unfused-MoE replacement module's per-expert
nn.Linears, built under torch.device("meta") alone with no dtype,
which default to float32 regardless of the checkpoint's actual dtype).
This is the streaming primitive; it isn't wired into AutoRound's own
model loading or tuning loop yet -- that follows in later commits.
Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
ModelContext._load_model() unconditionally called llm_load_model(..., device="cpu") (or mllm_load_model for multimodal checkpoints) whenever model was a string, fully materializing the checkpoint on CPU RAM before any block-wise AutoScheme/tuning logic ever ran. This is the fix for the initial-load problem: nothing downstream can be memory-safe if the model is already 100%+ resident before either ever runs. When AR_DISK_STREAM_MODEL=1 and model is a string (not diffusion), build an all-meta skeleton via the new disk_stream_util.build_meta_model() instead, for both the plain-text and multimodal (mllm_load_model) paths. Sets model.path = model_name (satisfies the existing but previously-dead unsupported_meta_device() escape hatch, which only allows an all-meta model) and stashes the SafetensorsIndex both on self._disk_stream_index and on model._disk_stream_index, so code that only has the model object (e.g. AutoScheme's gen_layer_config, which runs after ModelContext has already turned a string into an object) can still find it. Materializes non-block params (embeddings/lm_head/ final norm) for real right after the meta-device guard passes, leaving the (typically 100+GB combined) decoder blocks meta for later per-block materialize/free. Falls back to the original full CPU load on any exception. Also re-ties output embeddings via model.tie_weights() right after materializing: a tied lm_head.weight has no entry of its own in the checkpoint's safetensors index (relies on the model re-establishing the tie at load time, which a normal from_pretrained() does automatically but per-tensor materialization does not), so without this the tied module is silently left on meta. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
OffloadManager's existing per-block offload/reload cycle assumed every block started CPU-resident; starting from a meta skeleton (AR_DISK_STREAM_MODEL=1) broke it in three places: - _load_state_dict_into_module() copied a freshly-read real tensor onto the target parameter's existing device -- but for first-time materialization from meta, that existing device IS meta, so the copy silently discarded the real data instead of landing it on cpu. Now targets "cpu" specifically when the existing parameter is meta. - _save_to_disk() unconditionally recorded a block as saved even when its state_dict was empty (an all-meta block that hasn't been materialized yet has nothing real to persist). A later reload() then trusted that record and loaded an empty file, leaving the block meta. Now skips recording in that case. - _reload(), in "offload" mode, silently did nothing for a block not in self._saved (true for a still-meta block, or one _save_to_disk just started correctly skipping). Now falls back to load_block_from_model_files(self.model_dir, name, module) -- an existing upstream function, previously only used by "clean" mode -- reading the block directly from the original checkpoint. Requires compressors/base.py to propagate model_dir onto the offloader (next commit). Also fixes a real-scale bug found against qwen3.5-397b-base: when a checkpoint's on-disk MoE layout uses fused 3D expert tensors (experts.gate_up_proj/down_proj) but the in-memory module tree has already been replaced by unfused per-expert nn.Linears, assigning the fused key resolves to nothing and the experts stay meta. Added _maybe_split_fused_expert_keys(), which detects that mismatch and splits the fused tensor into per-expert keys via the existing missing_tensors.split_fused_expert_tensors() helper before assignment. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
OffloadManager._ensure_dir() always used tempfile.mkdtemp() -- a fresh, uniquely-named directory every process, impossible for a resumed process to ever find again. Whenever AR_RESUME_DIR is set, use a stable path (<AR_WORK_SPACE>/offload/<prefix>_resume/) instead, so a resumed process's OffloadManager can find and reuse whatever a prior crashed process already offloaded there (see the companion discovery check in _reload(), previous commit). Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Two small additions supporting the disk-streaming/resumability work in adjacent commits: - After constructing self.model_context, if it was built as a disk- streamed meta skeleton, propagate its checkpoint path onto self._offloader.model_dir, so OffloadManager can materialize never-yet-offloaded blocks directly from disk (see the reload fix in utils/offload.py). - New self._resume_states, cleared by quantize_and_save() only after save_quantized() actually returns successfully -- not right after the tuning loop finishes, since a crash during the export/packing step that follows would otherwise wipe resumability for no reason. Populated by DataDrivenCompressor.quantize() in the next commit. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Three fixes needed for the tuning/RTN loops to work correctly when a model was built as a meta skeleton (AR_DISK_STREAM_MODEL=1), unrelated to resumability: - The standard tuning loop's per-block reload only fired when low_cpu_mem_usage was true. GGUF export forces low_cpu_mem_usage False for reasons of its own (unrelated to disk streaming), so a streamed block was never materialized before GGUF's tuning loop touched it. Now also reloads when AR_DISK_STREAM_MODEL is set, regardless of low_cpu_mem_usage. - configure_layer_config() disables low_cpu_mem_usage for any non-MoE-patched (dense) model on the assumption that the whole model is already CPU-resident, so per-block offload/reload buys nothing. False once the initial load is itself no longer full-residency: keep it enabled when self.model_context._disk_stream_index is not None. - CalibratedRTNCompressor (--iters 0 path)'s safe_to_cpu_() call tries to consolidate the whole model onto CPU, including decoder blocks intentionally still on meta -- crashing with "Cannot copy out of meta tensor". Skipped in both the normal and OOM-fallback branches when AR_DISK_STREAM_MODEL is set. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
aquilarubra
force-pushed
the
pr/disk-streaming-core
branch
from
July 19, 2026 14:31
502d108 to
15aa885
Compare
for more information, see https://pre-commit.ci
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an opt-in
AR_DISK_STREAM_MODELenv var that lets AutoRound quantize a checkpoint far larger than available RAM+VRAM, by never fully materializing it in CPU memory.Today,
ModelContext._load_model()unconditionally callsllm_load_model(..., device="cpu")(ormllm_load_modelfor multimodal checkpoints) whenevermodelis a string, fully loading the entire checkpoint onto CPU RAM before any block-wise tuning logic ever runs. Even the "regular" per-block offload/reload machinery (OffloadManager) only actually engaged for fused-MoE-patched architectures — a plain dense model gotlow_cpu_mem_usagesilently forced back toFalse, since upstream's own assumption is "the whole model is already CPU-resident anyway, so per-block offload/reload buys nothing." That assumption breaks once the initial load itself is no longer full-residency.When
AR_DISK_STREAM_MODEL=1,ModelContextinstead builds an all-meta skeleton (accelerate.init_empty_weights(), ~0 RAM) and materializes/frees one decoder block's real tensors — read directly from the checkpoint's safetensors shards — right before/after each place that block is actually used. Default (unset) behavior is byte-for-byte identical to upstream.What's in this PR
auto_round/envs.py: newAR_DISK_STREAM_MODEL/AR_RESUME_DIRflags (the latter used by a follow-up resumability PR; unused by this PR alone).auto_round/utils/disk_stream_util.py(new): the streaming primitive — lazy, mmap-backed per-tensor reads from safetensors shards (never caches file handles, so touched pages don't stay RSS-resident), plus meta↔real materialize/free for a whole module.auto_round/context/model.py: builds the meta skeleton (text and multimodal paths) and materializes non-block params (embeddings/lm_head/final norm) right after the existingunsupported_meta_device()guard passes; re-ties tied output embeddings viatie_weights().auto_round/utils/offload.py: makes the existing per-block offload/reload cycle actually work starting from a meta model (first-time materialization target-device fix, meta-aware save/reload, fused-3D-expert-tensor key splitting for on-disk layouts that don't match an already-unfused module tree).auto_round/compressors/base.py: propagates the checkpoint path toOffloadManagerso it can materialize never-yet-offloaded blocks directly from disk.auto_round/compressors/data_driven.py: keepslow_cpu_mem_usageenabled for streamed dense models, reloads streamed blocks for GGUF export, and skips amodel.to("cpu")call that can't work on an intentionally-still-meta model.Validation
Verified on synthetic Llama fixtures (8-layer/2.8GB, 20-layer/27.7GB): peak RAM stays flat regardless of checkpoint size (vs. proportional-to-size on the baseline), per-block tuning losses and final quantized weights are bit-identical to the unstreamed baseline. Also verified end-to-end at real scale against a 207GB checkpoint and against synthetic MoE/hybrid-attention/MTP-head fixtures — happy to share more detail if useful for review.
This is the first of a small stack of PRs splitting up a larger local patch set (disk streaming → resumability → AutoScheme streaming → calibration fixes); happy to adjust scope/splitting based on maintainer preference.