Stream AutoScheme's per-layer sensitivity scoring block-by-block - #2063
Merged
Conversation
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
3 times, most recently
from
July 19, 2026 14:31
bb4d4b2 to
499daf1
Compare
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 |
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
from
July 22, 2026 12:25
58e3ccf to
21a28b4
Compare
Contributor
|
Hi @aquilarubra We have observed that using the default It would be very helpful if your change could reduce the memory usage based on that. |
Contributor
Author
|
Thanks for the heads-up @xin3he — yes, I see the overlap in On the memory question: this PR's |
wenhuach21
approved these changes
Jul 27, 2026
Contributor
|
@aquilarubra recently we have refactored the arch to better support more algorithms, could you kindly help fix the conflict? then we can merge your prs. |
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
from
July 27, 2026 07:17
21a28b4 to
2c6a023
Compare
Contributor
Author
|
/azp run Performance-Test-AutoRound |
|
Commenter does not have sufficient privileges for PR 2063 in repo intel/auto-round |
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
from
July 27, 2026 09:25
cb53811 to
b71b8b9
Compare
Contributor
|
Hi @aquilarubra FYI. #2083 is merged now and you can resolve the conflicts. Looking forward to getting this PR merged soon. 😀 I plan to set the |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends AutoRound’s disk-streaming support so AutoScheme’s per-layer sensitivity scoring can materialize/free one decoder block at a time (instead of assuming full CPU-residency), reducing peak RAM for very large checkpoints while preserving scoring outcomes.
Changes:
- Add disk-streaming primitives (
SafetensorsIndex, meta↔real materialize/free, and forward wrapping) and thread disk streaming through AutoScheme scoring paths. - Improve offload/reload behavior for meta-skeleton models (including fused-MoE key handling and resumability-oriented reload behavior).
- Add CPU tests for disk-stream utilities and AutoScheme streaming parity, plus EN/CN docs for
AR_DISK_STREAM_MODEL.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| test/test_cpu/utils/test_disk_stream_util.py | Unit tests for materialize/free round-trip and dtype behavior of disk-stream utilities. |
| test/test_cpu/schemes/test_auto_scheme_disk_stream.py | Verifies AutoScheme streaming produces identical mixed-bit layer_config to baseline and exercises primitives. |
| docs/environments.md | Documents AR_DISK_STREAM_MODEL (and includes AR_RESUME_DIR section). |
| docs/environments_CN.md | Chinese translation for the AR_DISK_STREAM_MODEL documentation update. |
| auto_round/utils/offload.py | Adds fused-expert key splitting, meta-aware load behavior, and resumability-oriented reload/dir logic. |
| auto_round/utils/disk_stream_util.py | New disk-streaming utilities: safetensors indexing, per-module materialize/free, and streaming forward wrapper. |
| auto_round/envs.py | Adds AR_DISK_STREAM_MODEL, AR_RESUME_DIR, and AR_CALIB_STREAM_DEVICE env handling. |
| auto_round/context/model.py | Adds meta-skeleton load path and non-block parameter materialization for disk streaming. |
| auto_round/compressors/orchestrator.py | Ensures streamed blocks are reloaded/materialized when needed even if low_cpu_mem_usage is disabled. |
| auto_round/compressors/base.py | Propagates checkpoint dir to OffloadManager for streamed models; defers resume-state clearing until save succeeds. |
| auto_round/auto_scheme/gen_auto_scheme.py | Stops forcibly disabling low_cpu_mem_usage in AutoScheme generator init. |
| auto_round/auto_scheme/delta_loss.py | Threads disk_index through scoring and implements per-block materialize/free during forward/backward replay. |
Contributor
Great! How about the runtime? If the overhead is low, we could make this the default. |
xin3he
approved these changes
Jul 29, 2026
xin3he
left a comment
xin3he
left a comment
Contributor
There was a problem hiding this comment.
Nice work!
@aquilarubra
Please help resolve the conflicts, Please feel free to ask help if you have any issue.
Contributor
|
@copilot resolve the merge conflicts in this pull request |
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>
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>
for more information, see https://pre-commit.ci
- Remove "Local addition"/LOCAL_PATCHES.md/vendor/reap references from every touched file -- local-only tooling metadata that doesn't apply upstream. - materialize_module() in disk_stream_util.py always forced the checkpoint's raw on-disk dtype onto rematerialized decoder blocks, fighting the compute dtype ModelContext._set_amp_dtype() had already promoted the meta skeleton to. This crashes with a dtype mismatch (e.g. BFloat16 vs Half) the moment a checkpoint's native dtype differs from the chosen amp dtype. Now prefers the meta parameter's already-declared (promoted) dtype, only falling back to the checkpoint's dtype for the one case that motivated the original behavior: an untyped meta context defaulting to float32. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
- Materialize/free round-trip for a real decoder block, and a re-materialize-is-a-no-op check for already-real (e.g. tied) params. - Regression coverage for the meta-materialize dtype bug: materialize_module must prefer the meta parameter's already-declared dtype (reflecting whatever compute dtype the caller promoted the model to) over the checkpoint's raw on-disk dtype, except when the declared dtype is an untyped-context float32 default. Verified this test fails against the pre-fix code with the exact reported "BFloat16 vs Half"-style mismatch. - build_meta_model + materialize_non_block_params: non-block params (embeddings) become real while decoder blocks stay meta. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
…t flake) Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
…t flake) Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
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>
Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
When AR_RESUME_DIR is set, DataDrivenCompressor.quantize() now builds one ResumeState per block group (auto_round/utils/resume.py, added in the next commit), keyed by a signature over model path + scheme + dataset + nsamples/seqlen + block list. On a partial resume, the group's first not-yet-done block substitutes its cached input_others from the pre-existing all_inputs cache, but the chained input_ids/ q_input come from the ResumeState's cached tensors, not that cache -- the pre-cache pass and the in-loop reference forward aren't numerically identical, so reusing the wrong one produced a 20x larger tuning loss on the first resumed block in testing. _quantize_blocks() starts its loop at resume_state.resume_index instead of 0 (nblocks=1 only), forces shard_writer._flush_shard() after each block when resuming is active (write() alone only buffers until the shard-size budget is hit -- a lie about durability that a real crash-and-resume test exposed as zero files on disk), and calls resume_state.mark_block_done(...) only after that write, so a crash before it correctly re-does the block rather than skipping it with incomplete output. Clearing the resume manifest is deferred to quantize_and_save() (see compressors/base.py, previous commit) rather than done right after the tuning loop, for the shard-export path specifically: quantize() returning successfully isn't the end of the pipeline there, and a crash during the packing/config-write step that follows would otherwise wipe resumability for no reason. The final "reload everything before returning" call now passes the full flattened block list explicitly when AR_RESUME_DIR is set (skipped entirely under shard-export/is_immediate_saving): reload(names=None) only reloads names already in the offloader's own _saved dict, which never includes a block a resumed process skipped entirely via ResumeState. Under shard export, reloading those blocks back to real memory is actively harmful, not just unnecessary -- the shard_writer's subsequent is_finalize=True write would re-emit their raw, unpacked weights alongside the already-correct packed ones already flushed by a prior process, producing duplicate/inconsistent tensors for the same layer (confirmed by diffing tensor names against an uninterrupted control run). Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
New auto_round/utils/resume.py, no upstream equivalent. Tracks completed blocks plus cached chain tensors (resume_q_input.pt/ resume_input_ids.pt) for a tuning run, keyed by a signature hash over model path + scheme + dataset + nsamples/seqlen + block list, so a resume directory reused for a different run is detected and ignored rather than silently misapplied. Both chain tensors (q_input and input_ids) are cached, not just q_input: the FP reference chain (input_ids) is not numerically identical between AutoRound's pre-tuning cache pass and the in-loop reference forward, so reconstructing it from the pre-cache instead of persisting the live value produced a 20x larger tuning loss on the first resumed block in testing. Also adds layer_config_fingerprint(), folded into the run signature by this file's callers (auto_round/compressors/data_driven.py, previous commits): str(self.scheme) (or the literal "rtn_with_imatrix") alone is bits-blind for AutoScheme runs -- two runs against the same model/dataset/nsamples/seqlen but different avg_bits targets produced identical signatures, so the second run silently resumed the first's already-complete manifest and saved an output containing no layer tensors at all. Folding the resolved per-layer bit allocation into the signature fixes this. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
A fresh process's ShardWriter has no memory of shards a previous, crashed process already flushed to output_dir -- it would restart shard_counter at 0, collide with existing shard filenames, and finalize()'s index would only cover this process's tensors, producing a corrupt/incomplete checkpoint. Adds _discover_existing_shards(): when AR_RESUME_DIR is set, on the first real _flush_shard() call (not __init__ -- see below), scans output_dir for leftover pre-rename model-shard-NNNNN.<ext> files, reads each one's tensor names straight from its safetensors/torch header (no data materialization needed), and seeds shard_counter/shard_meta/ _all_saved from them so numbering doesn't collide and finalize()'s index covers both processes' shards. Discovery has to be deferred past __init__: ShardWriter.__init__ runs during post_init(), before quantize_and_save()'s _get_export_dir() appends the final subfolder (e.g. <model>-w4g128/) to output_dir -- discovering at construction time silently looked in the wrong directory and found nothing, confirmed by a real crash-and-resume test where blocks 0-2 resumed correctly through tuning but still lost their output. Fixed by running discovery lazily inside _flush_shard() itself, guarded by a self._existing_shards_discovered flag, by which point output_dir is always the final path. Gated on AR_RESUME_DIR throughout, so normal non-resuming runs never change behavior even if output_dir happens to be reused. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
A resumed disk-streamed run only materializes/quantizes the blocks it didn't already finish in a prior (crashed) process; blocks it skipped are untouched in this process and stay on the meta device, while their packed weights already live in shard files the previous process flushed to disk (see ShardWriter._discover_existing_shards, earlier commit). The global post-tuning packing pass otherwise crashed trying to read .scale off such a layer. Early-return when the layer's weight is still on meta: there is nothing to pack, and the on-disk export for it is already complete. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
for more information, see https://pre-commit.ci
Remove "Local addition"/LOCAL_PATCHES.md references from every touched file -- local-only tooling metadata that doesn't apply upstream. No logic change. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
- test/test_cpu/utils/test_resume.py: unit tests for ResumeState (mark_block_done ordering, q_input/input_ids round-trip, signature mismatch and non-prefix manifest handling both correctly discard stale state, clear()), compute_run_signature, and layer_config_fingerprint. - test/test_cpu/core/test_resume_integration.py: end-to-end test that simulates a crash after the first block (injected via a ResumeState.mark_block_done wrapper that raises right after persisting state) and verifies a fresh AutoRound run against the same AR_RESUME_DIR resumes from the second block only, producing a complete layer_config for both blocks. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
…t flake) Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
…_usage, docs, tests) Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Reconstructed against the post-intel#2083 caching/parallel-scoring rewrite of delta_loss.py: threads disk_index through the serial scoring path (prepare_model_low_gpu, model_forward_low_gpu, get_score_for_scheme, gen_layer_config/_gen_layer_config) the same way as before, but now explicitly excludes streaming from the parallel multi-process scoring path added by intel#2083 -- each parallel worker fully loads its own copy of the model in a separate process, which defeats disk streaming's entire purpose. Per-scheme score caching is unaffected either way. The two model.to("cpu") calls that used to need an explicit disk_index-aware skip are now handled for free by safe_to_cpu_() (added upstream independently), which already checks for meta tensors before moving -- no manual guard needed at either call site anymore. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
aquilarubra
force-pushed
the
pr/autoscheme-streaming
branch
from
July 29, 2026 07:32
b71b8b9 to
b9eb98f
Compare
for more information, see https://pre-commit.ci
Contributor
Author
|
@xin3he Thanks for reviewing — glad the memory reduction checks out. Quick note left on #2061 for context: this branch is stacked on it, so merging #2063 brings #2061's commits along too, and #2062/#2064 (also stacked on #2061) will need a rebase once this lands. Also flagging: I'll be on vacation for the next 3 weeks starting now, so I won't be able to respond to review comments or handle the rebase during that window — will pick it back up once I'm back. #2065 is independent and fine to merge separately anytime. |
Contributor
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
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
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.