Stream targeted block re-quantization's calibration pass - #2064
Open
aquilarubra wants to merge 34 commits into
Open
Stream targeted block re-quantization's calibration pass#2064aquilarubra wants to merge 34 commits into
aquilarubra wants to merge 34 commits into
Conversation
aquilarubra
force-pushed
the
pr/targeted-block-calib
branch
2 times, most recently
from
July 19, 2026 14:31
11ba30e to
3be3fae
Compare
chensuyue
requested review from
xin3he and
yiliu30
and removed request for
yiliu30
July 20, 2026 02:12
aquilarubra
force-pushed
the
pr/targeted-block-calib
branch
5 times, most recently
from
July 27, 2026 09:25
6642d22 to
c6e097e
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>
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>
Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
to_quant_block_names lets a caller restrict tuning to a subset of decoder blocks -- useful for cheaply re-quantizing just a couple of blocks in an already-produced checkpoint at higher precision instead of redoing a full multi-hour run. Combined with AR_DISK_STREAM_MODEL, this crashed: the "cache block inputs" forward pass (calibrate_on_cpu branch) needs real weights in every block leading up to (and, when there's only one target block, all the way through) the target block(s), but nothing materializes blocks outside quant_block_list for this specific forward pass -- they stay meta forever, and the forward silently propagates meta-ness through them until it collides with a genuinely-materialized module. Reproduced against a tiny hybrid-MoE fixture with to_quant_block_names restricted to one block: "Tensor on device meta is not on the expected device cpu!" inside the final norm. Full (unrestricted) runs never hit this, since quant_block_list already covers every block in that case. Fix: when disk streaming is active and any decoder block still has meta parameters at this point, wrap the calibration forward with the existing stream_block_forward primitive (already used by _streaming_eval_model() in bin/el_quantize_autoround_mixed.py for the analogous held-out-loss-eval case), scoped to just the still-meta blocks -- materializing each on demand and freeing it right after. Only activates when there's something left meta to fix, so it's a no-op for the normal full-quantization path. Also adds an AR_CALIB_STREAM_DEVICE env-gated fast path: the default keeps this forward entirely on cpu (mixing a GPU-streamed block with cpu-resident hidden states crashed with a device mismatch the first time this was tried at full 397B scale), but a full cpu forward through every pre-target block of a 100B+ model is unusably slow for this specific targeted-requant use case. When set, every already-real (non-meta) param/buffer is moved to that device for the duration of the pass (including stray non-parameter buffers like RoPE inv_freq), blocks stream-materialize there too, and everything is moved back afterward so the tuning phase sees the exact layout it would have without this. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
Remove a "Local addition (not upstream)" comment prefix -- local-only tooling metadata that doesn't apply upstream. No logic change. Signed-off-by: Fabrizio del Tin <devotedmystic@gmail.com>
test_targeted_block_with_disk_streaming_does_not_crash reproduces the exact crash this PR fixes: restricting to_quant_block_names to the last of 3 blocks leaves the earlier blocks meta-only (never in quant_block_list), and the calibration forward pass used to propagate that meta-ness until it hit "Tensor on device meta is not on the expected device cpu!". Verified this test fails with that error against the pre-fix commit (15aa885) and passes with the fix. test_targeted_block_without_disk_streaming_still_works is the baseline: the same targeted re-quantization without streaming never hit this bug, so it must keep working unchanged. 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>
aquilarubra
force-pushed
the
pr/targeted-block-calib
branch
from
July 29, 2026 07:34
c6e097e to
3841bb3
Compare
for more information, see https://pre-commit.ci
This was referenced Jul 29, 2026
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
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.
to_quant_block_nameslets a caller restrict tuning to a subset of decoder blocks — useful for cheaply re-quantizing just a couple of blocks in an already-produced checkpoint at higher precision instead of redoing a full multi-hour run. Combined withAR_DISK_STREAM_MODEL, this crashes: the "cache block inputs" calibration forward pass needs real weights in every block leading up to (and sometimes through) the target block(s), but nothing materializes blocks outsidequant_block_listfor this specific pass — they stay meta forever, and the forward silently propagates meta-ness until it collides with a genuinely-materialized module. Full (unrestricted) runs never hit this, sincequant_block_listalready covers every block in that case.What's in this PR
auto_round/calibration/llm.py: when disk streaming is active and any decoder block still has meta parameters at the point this calibration forward runs, wraps it with the existingstream_block_forwardprimitive (from #2061), scoped to just the still-meta blocks. Only activates when there's something left meta to fix, so it's a no-op for the normal full-quantization path.Also adds an
AR_CALIB_STREAM_DEVICEenv-gated fast path: the default keeps this forward entirely on CPU (mixing a GPU-streamed block with CPU-resident hidden states crashes with a device mismatch), but a full CPU forward through every pre-target block of a 100B+ model is unusably slow for this specific targeted-requant use case. When set, every already-real tensor is moved to that device for the pass's duration and moved back afterward.Validation
Reproduced and fixed against a tiny hybrid-MoE fixture with an MTP head, with both a single restricted block and two adjacent ones, plus a control run of the full (unrestricted) path confirming no regression.