Add crash/resume support for tuning runs (AR_RESUME_DIR)#2062
Open
aquilarubra wants to merge 13 commits into
Open
Add crash/resume support for tuning runs (AR_RESUME_DIR)#2062aquilarubra wants to merge 13 commits into
aquilarubra wants to merge 13 commits into
Conversation
aquilarubra
force-pushed
the
pr/resumability
branch
from
July 19, 2026 14:20
32d501d to
4d5ab67
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/resumability
branch
from
July 19, 2026 14:31
3f881e3 to
124fab4
Compare
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>
CalibratedRTNCompressor._quantize_via_rtn_blockwise()'s per-block loop
never reloaded blocks from the disk-stream index the way the standard
tuning loop does -- materialize_model_() only rebuilds MoE
ReplacementModuleBase wrappers, it has no knowledge of the disk-stream
index, so under AR_DISK_STREAM_MODEL every block was still meta by the
time block.to("cpu") ran. Added the same
self._offloader.reload(model, block_name) call used by the standard
loop, gated on low_cpu_mem_usage.
This path also never wrote anything to shards *during* its loop --
unlike the standard tuning loop, it only called
shard_writer.write(is_finalize=True) once, at the very end of the
entire run. Naively layering ResumeState bookkeeping on top of that
would have been actively dangerous: a resumed run would trust the
manifest, skip already-"done" blocks, and produce a checkpoint
silently missing their weights, since nothing had ever been flushed
for them. Fixed in two layers: first, added the missing per-block
immediate_pack + shard_writer.write(block, is_finalize=False) call
(mirroring the standard loop's equivalent block), making this path's
output incremental on its own, independent of resumability; only then
added ResumeState wiring on top (one ResumeState per block group,
signature includes "rtn_with_imatrix" so it can never collide with a
standard-tuning-loop resume manifest), including the same forced
_flush_shard() and mark_block_done() pattern as the standard loop.
q_input is always passed as None to mark_block_done() -- this path
doesn't track a separate quantized-input chain value.
The final "reload everything before returning" call gets the same
fix as the standard loop's equivalent call (previous commit): pass the
full flattened block list explicitly when AR_RESUME_DIR is set, so a
resumed process's skipped blocks are still reloaded; skipped entirely
under shard export for the same duplicate-tensor reason.
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>
aquilarubra
force-pushed
the
pr/resumability
branch
from
July 19, 2026 14:47
408691d to
95b7762
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
Depends on #2061 (disk-streaming core) — this branch is stacked on top of it, so the diff below includes those commits until #2061 merges; only the last 5 commits are new here.
Adds an opt-in
AR_RESUME_DIRenv var: when set, the tuning loop (both the standard SignRound loop and the--iters 0RTN-with-imatrix path) checkpoints its progress to disk after each completed block, so a crash or kill mid-run can resume from the first not-yet-done block instead of restarting from block 0. Useful independently of disk streaming, but especially so for very large models where a full run spans hours.What's in this PR
auto_round/utils/resume.py(new):ResumeState, tracking completed blocks plus cached chain tensors (q_input/input_ids) keyed by a signature hash over model/scheme/dataset/block-list, so a reused directory from an incompatible run is detected and ignored. Alsolayer_config_fingerprint(), folded into the signature so twoAutoSchemeruns with differentavg_bitstargets (which otherwise produce identical signatures) don't silently resume each other's state.auto_round/compressors/data_driven.py: wiresResumeStateinto both the standard tuning loop and the RTN-with-imatrix path (the latter needed a prerequisite fix first — it never wrote per-block output during its loop at all, only in one bulk pass at the end, which would have made resumability actively unsafe).auto_round/compressors/shard_writer.py: makes shard-based export resume-aware across process restarts (discovers a crashed prior process's already-flushed shards and continues numbering from there).auto_round/export/export_to_autoround/export.py: skips packing for a resumed, still-meta layer (its packed weights already live in a previously-flushed shard).auto_round/compressors/base.py(already in Stream large checkpoints from disk during quantization #2061): defers clearing the resume manifest untilsave_quantized()actually succeeds, not right after tuning finishes.Validation
Verified via real crash-and-resume tests (
kill -9mid-run, fresh process resumes) across dense/MoE architectures, bothiters>0anditers=0, both bare.quantize()and shard-export paths: resumed output is bit-identical to an uninterrupted control run in every case, confirmed by diffing full tensor sets. Also verified at real 207GB-checkpoint scale.Happy to squash/reorganize commits if a different granularity is preferred for review.