Stream AutoScheme's per-layer sensitivity scoring block-by-block#2063
Open
aquilarubra wants to merge 9 commits into
Open
Stream AutoScheme's per-layer sensitivity scoring block-by-block#2063aquilarubra wants to merge 9 commits into
aquilarubra wants to merge 9 commits into
Conversation
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
2 times, most recently
from
July 19, 2026 14:27
445f953 to
5721fe1
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>
AutoScheme's mixed-bit search scores every (scheme, block) pair to build its DP knapsack input -- a second, independent place (besides the compressor's own tuning loop) that needs real weights one block at a time, and one that runs *before* the compressor's OffloadManager cycle even starts, so it can't reuse that cycle directly. gen_auto_scheme.py: removed GenScheme.__init__'s unconditional `low_cpu_mem_usage = False` override (upstream commit 0c9c5b1), originally added to work around an acknowledged bug for mixed INT4/INT8 schemes in an old OffloadManager-based streaming path. This patch set doesn't use that path -- see below -- so the bug that motivated disabling it doesn't apply. delta_loss.py: threads a disk_index parameter through AutoScheme's own scoring call sites (gen_layer_config, _gen_layer_config, get_score_for_scheme, prepare_model_low_gpu, model_forward_low_gpu). When set, each decoder block is materialized (materialize_module, from utils/disk_stream_util.py) right before scoring and freed (free_module) right after, instead of being assumed already CPU-resident (old code) or moved with .to("cpu") (which never actually frees RAM once a block started resident either way). In the "object passed in" branch of gen_layer_config, reads disk_index = getattr(model, "_disk_stream_index", None) -- set by ModelContext when it builds a meta skeleton (context/model.py, earlier commit) -- so streaming works through the standard AutoRound(model=path, scheme=AutoScheme(...)) API without callers needing to pass anything new themselves. Also skips two unconditional model.to("cpu") calls (one right after scoring setup, one right after the DP knapsack selection finishes) when disk_index is set: both assume the model is, or should become, fully resident on cpu, which raises "Cannot copy out of meta tensor" for a meta skeleton with intentionally-still-meta decoder blocks. The second call site's attribute-cleanup loop (dropping scoring-only scale_dtype/imatrix/tuning_device markers) still runs regardless, since the tuning phase reuses the same model object afterward. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
from
July 19, 2026 14:31
bb4d4b2 to
499daf1
Compare
for more information, see https://pre-commit.ci
Contributor
|
Really appreciate this PR! It's exactly what we needed. We're currently in a code freeze while we refactor the architecture, so we aren't merging new feature PRs at the moment. Once the refactoring is complete (ETA 1 week), we'll merge your PR |
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
2 participants
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
Depends on #2061 (disk-streaming core) — stacked on top of it, so the diff includes those commits until it merges; only the last commit is new here.
AutoScheme's mixed-bit search scores every(scheme, block)pair to build its DP knapsack input — a second, independent place (besides the compressor's own tuning loop) that needs real weights one block at a time, and one that runs before the compressor'sOffloadManagercycle even starts.What's in this PR
auto_round/auto_scheme/gen_auto_scheme.py: removesGenScheme.__init__'s unconditionallow_cpu_mem_usage = Falseoverride (added in0c9c5b1dto work around a bug in an oldOffloadManager-based streaming path this patch set doesn't use).auto_round/auto_scheme/delta_loss.py: threads adisk_indexparameter through AutoScheme's scoring call sites, materializing/freeing each decoder block around scoring instead of assuming it's already resident or moving it with.to("cpu")(which never actually frees RAM once a block started resident). Readsdisk_indexoff the model object itself (set byModelContextin Stream large checkpoints from disk during quantization #2061), so streaming works through the standardAutoRound(model=path, scheme=AutoScheme(...))API with no caller-visible changes. Also skips two unconditionalmodel.to("cpu")calls that assume full CPU-residency and otherwise crash with "Cannot copy out of meta tensor" for a still-partially-meta model.Validation
Verified end-to-end (
AutoScheme(avg_bits=4.5, options=["W4A16","W8A16"])+ tuning) on synthetic fixtures up to 27.7GB: peak RAM stays flat through both scoring and tuning phases (vs. proportional-to-checkpoint-size on the baseline), and the resulting mixed-bitlayer_configis byte-identical to the unstreamed baseline. Also verified at real 207GB scale.