feat: checkpointing and resuming quantization process from last checkpoint

[okdshin]

Summary

Adds an opt-in, crash-survivable resume path for GPTQModel.quantize() (env GPTQMODEL_RESUME=1), for sequential GPTQ jobs on very large models that can run tens of hours and previously had to restart from scratch after any crash.

Follows up on the discussion in #3015, where @Qubitium asked whether this was going to be upstreamed rather than reimplemented independently.

Design

• GPTQMODEL_RESUME=1 on the original run makes it resumable (forces synchronous per-layer finalize draining, writes quant_resume_state.json once a layer's modules are confirmed durable on disk). The same flag on a restart triggers the actual resume. offload_to_disk users who never set it get the old async-drain behavior unchanged, at no added cost.
• On restart, completed transformer layers are replayed forward-only (quantized modules loaded back from the offload directory swap in for the original meta modules) to regenerate the next layer's calibration inputs, then quantization resumes at the first unfinished layer. The most-recently-completed layer's replay itself can be skipped via a small activation cache.
• A resume fingerprint (model identity, calibration content, quant config, checkpoint weight-file metadata) guards against silently resuming against a different checkpoint or calibration run, which would otherwise produce a hybrid model with no error.

Split into 3 commits for review:

1. fix: make disk-offload writes crash-atomic — prerequisite fix, unrelated to resume itself (the old rmtree-then-write offload scheme had a crash window that could destroy the only durable copy of an already-finalized module).
2. feat: crash-survivable mid-quantization resume — the core mechanism.
3. fix: harden resume fingerprint against calibration/checkpoint drift — closes several ways two different runs could compare as "unchanged" and resume against stale state.

Validation

• Unit tests (tests/test_resume.py, 47 tests) covering the resume-target/fast-forward decision logic, fingerprint edge cases, and activation-cache branching in isolation.
• A dummy-model end-to-end driver (reference run vs. kill-mid-run-and-resume vs. bitwise safetensors comparison) — bitwise-identical output, re-verified after every change to the core replay/hashing logic.
• Real production use across several multi-hour quantization jobs on this branch's predecessor state (42-layer/~12h DeepSeek-V4-Flash run, two 64-layer Qwen runs, one deliberately interrupted and resumed mid-run) — zero resume-correctness issues, only unrelated infra bugs (JIT-cache/device-isolation issues under concurrent processes, tracked separately in Concurrent quantize() processes on the same machine can hang indefinitely (shared pack_block_cpu JIT extension build directory is not process-isolated) #3015).

Open questions for maintainers

As noted in the original design: marker file format (quant_resume_state.json), the GPTQMODEL_RESUME env var naming/semantics (single flag gating both "make resumable" and "resume"), and the activation-cache replay-skip semantics are all judgment calls I'd like feedback on before considering this final.

[Qubitium]

Reviewed head 86c6f9744edf2fdd06684b93c524fccc999e87b3. I found five issues that should be addressed before merging.

1. [P1] Cache-miss replay uses the wrong weights. stage_layer.py:425–441

   The original-weight replay premise is incorrect for GPTQ: GPTQProcessor.process() assigns module.weight.data = wq (gptq_processor.py:591), and the normal layer stage then performs its post-process forward before packing. Consequently, the next layer receives activations computed with the quantized dense reconstructions. Resume instead materializes original checkpoint weights, forwards them, and only afterward restores packed modules. A missing/unusable activation cache—including a crash between marker publication and cache publication—therefore changes the calibration inputs and subsequent quantization. Reconstruct the same post-process dense weights for replay, or require a matching durable activation checkpoint. Add an uninterrupted-versus-resumed comparison that explicitly removes the activation cache and exercises the real GPTQ processor.

2. [P1] Activation payload and metadata can come from different layers after interruption. resume.py:273–289

   The two independent replacements are not an atomic cache commit. Start with layer 0's cache, publish layer 1's marker, then interrupt after replacing the tensor file but before replacing its JSON. Layer 1's cache lookup correctly misses, but fallback replay subsequently calls load_activation_cache(..., 0, ...), which accepts layer 0's old metadata together with layer 1's new tensors. With matching batch structure, this silently injects an extra layer's activations. I reproduced this using the actual save/load helpers with filesystem fault injection and tensor serialization stubbed. Use immutable generation-specific payload files and atomically publish metadata referencing the exact generation, or bind the layer/fingerprint into the tensor file and validate both.

3. [P1] The fingerprint accepts incompatible quantization settings. resume.py:121–138

   Output-affecting settings such as damp_percent, act_group_aware, static_groups, and mse are absent. I wrote a marker, changed these settings, and confirmed read_resume_target() still accepts it. A restart then mixes layers produced under the old settings with newly quantized layers under the new settings. Fingerprint a canonical representation of all algorithm/output-affecting configuration, including relevant processor-specific options, while excluding runtime-only settings. Add mismatch tests for these fields.

4. [P2] Fully dynamically excluded intermediate layers make resume fail. stage_layer.py:529–535

   If all tracked modules in an intermediate layer are excluded via dynamic, the normal path replays that untouched layer and creates no finalized bundles. Once a later layer writes a marker, resume attempts to restore every earlier transformer layer and raises when this legitimate excluded layer has no bundles. The same unconditional empty-restore error exists in _resume_replay_layer(). The model-level should_quantize_layer check does not handle ordinary per-module dynamic exclusions. Preserve untouched layers without requiring bundles; replay them when upstream activations are needed. Test an excluded layer between two quantized layers, both with and without an activation-cache hit.

5. [P2] Class-only restore loses the recovered .old bundle location. resume.py:577–578

   _offloaded_layer_modules() recovers a lone <module>.old/module.safetensors, but class-only restore discards that path and moves the module to meta. Later, get_state_dict_for_save() searches the canonical module/ancestor directories, not .old; the recovered index also references the missing canonical bundle path. Thus restore reports success but save cannot resolve the weights. I verified the mismatch between the recovery helper and save's search paths. Promote the recovered directory back to its canonical location before class-only restoration, or explicitly preserve a corrected offload mapping. Cover restore-through-save after interruption between the two offload renames.

Validation: inspected all five changed files and the surrounding processing, replay, and save paths. Ran isolated checks against extracted PR helper functions for findings 2, 3, and 5; filesystem operations were real, while tensor/model collaborators were stubbed. PyTorch is unavailable in this environment, so I could not run tests/test_resume.py or a real GPU end-to-end quantization comparison.

[Qubitium]

@okdshin Please resolve the ai reviews before I begin manual review. Thank you for thr pr.

[Qubitium]

@okdshin Going to refractor this PR with a slighly different design for better usability and future expansion.

[Qubitium]

Proposed redesign: explicit looper boundaries with checkpointing as an extension

The current implementation publishes progress separately from continuation activations and falls back to replaying original weights. That replay does not match GPTQ's normal forward after module.weight.data = wq. The redesign should eliminate recovery-specific layer replay rather than add more branches to it.

The looper needs four first-class concepts:

1. A stable execution plan: uniquely identified input-embedding, transformer-layer, and output steps, including explicit skipped steps.
2. A cursor identifying the next step in that plan.
3. A boundary whose finalization futures can be drained and whose continuation can be captured before the next step mutates it.
4. A continuation containing the complete InputCache, supported processor/model shared state, and RNG state where needed. Serialization uses explicit versioned adapters, not arbitrary Python object serialization.

An extension receives boundaries on the orchestration thread. A checkpoint extension decides when to wait for finalization, capture state, and commit. Existing quantization/subset/device scheduling stays in the normal execution path. The looper does not implement filesystem layout, fingerprints, retention, or recovery policies.

Conceptual flow:

prepare execution plan
extension.prepare(context) -> optional ResumePoint
initialize inputs OR restore ResumePoint
for step from cursor:
    result = execute_step(step)
    extension.on_boundary(boundary)
drain finalizers
extension.on_complete(context)

Finalization should return module names, quantization specifications, and artifact references. Committed artifacts must be immutable; ordinary disposable offload paths are not sufficient. A restore adapter builds the model shell, registers completed modules lazily against those artifacts, installs the continuation, and starts directly at the next step. Skipped modules stay backed by the source checkpoint.

A checkpoint store writes artifacts and continuation first, flushes them according to its durability contract, then atomically publishes one manifest referencing both. A CURRENT pointer selects a complete manifest. A killed writer leaves the previous complete checkpoint available. Retention keeps at least two generations and removes only unreferenced artifacts. There is no original-weight replay fallback: recovery from a damaged checkpoint uses an older complete continuation or fails clearly.

Concurrency: workers own finalization tasks; one coordinator publishes checkpoints after the required futures complete. A run-level lease rejects a second writer to the same checkpoint directory. Independent module artifacts and independent runs can still be written in parallel. SIGINT/SIGTERM set a stop request handled at a boundary; SIGKILL needs no cleanup handler.

Proposed public configuration:

model.quantize(calibration, checkpoint=CheckpointConfig(
    path="/checkpoints/my-run", resume="auto", every_layers=1, keep_last=2,
))

required rejects missing/incompatible checkpoints; never creates a new run without overwriting an existing one. Inspection reports the next step and field-level compatibility differences. Canonical run identity includes the execution plan, source checkpoint identity, prepared calibration identity, and supported algorithm-affecting options.

Implementation sequence:

1. Introduce and exercise the boundary/extension contract in the existing looper, preserving its default scheduling.
2. Add stable plan/cursor and versioned continuation adapters, initially for sequential GPTQ.
3. Implement immutable artifacts, transactional store, and lazy restoration; connect the public configuration.
4. Replace the legacy resume branches and validate dense/MoE equivalence across graceful stops, SIGKILL at publication boundaries, write failures, corrupt generations, skipped layers, repeated resumes, and concurrent-writer rejection with the GIL disabled.

Validation correction: the earlier local Llama driver lacked a __main__ guard and was re-executed by the forkserver. Its duplicate writers invalidate the earlier attribution of the directory race to no-GIL thread safety. Checkpoint creation alone also did not establish uninterrupted/resumed equivalence. Those runs must be replaced by guarded subprocess tests before claiming recovery correctness.

Starting with the boundary extension foundation; this is a staged implementation, not a claim that the existing recovery path is already corrected.

[Qubitium]

Implemented the checkpoint redesign in 02e0f08e on this PR's existing feat-quant-resume head branch.

Implementation

• Public quantize(..., checkpoint=CheckpointConfig(path=..., resume="auto", every_layers=1, keep_last=2)) API and CheckpointStopped exception.
• Generic looper startup plan/cursor and boundary extension contracts. The layer stage has 343 lines of legacy logic removed, with no persistence or model restoration logic added to it.
• Versioned complete InputCache/shared-state/log/RNG continuation; no pickle and no original-weight replay.
• Streamed content-addressed artifacts, manifests and atomic CURRENT publication; file/directory fsync, exclusive run lease, retained-generation fallback, corruption repair, and orphan cleanup.
• Quantized modules restored on meta; disposable save indexes reference immutable artifacts directly.
• SIGINT/SIGTERM request a safe-boundary commit/stop. SIGKILL resumes from the latest complete published generation.
• Removed the old marker/replay implementation and its obsolete tests; added contract, codec, transaction, concurrency and real quantization recovery coverage. Removed the global offload flock: finalizer threads retain parent-module locks, while checkpoint runs use an exclusive lease and private attempt directories.

Validation

All runs below used free-threaded Python 3.14.7 with PYTHON_GIL=0. Subprocess drivers assert that the GIL remains disabled after model-stack imports.

• Core/offload/looper/concurrency suites: 95 passed, 1 skipped.
• Dense/MoE subprocess recovery plus stage, weight-only and offload-config regression suites: 56 passed.
• Additional completed-checkpoint restart check: 2 passed (overlaps the dense/MoE tests above).
• Storage termination tests include partial object writes and death immediately before/after CURRENT publication, for both SIGTERM and SIGKILL.
• Two-layer Llama and Qwen3 MoE: every saved tensor matches an uninterrupted run after SIGINT/SIGTERM, SIGKILL before/after commits, and injected publication failure. Completed layers are not re-executed.
• Local /monster/data/model/Llama-3.2-1B on CUDA: uninterrupted baseline vs SIGKILL/restart and a separate SIGTERM/restart both matched all 482 saved tensors exactly. These use four synthetic 32-token calibration batches to test recovery, not model quality.
• Ruff on the new implementation/tests and git diff --check: passed.

Environment caveat: the available free-threaded environment has Transformers 5.5.4, older than this repository's declared requirement, and lacks torchvision. Validation used temporary, uncommitted import-compatibility stubs for torchvision and create_recurrent_attention_mask. These are not part of the PR. This is not a claim of validation against the full supported dependency matrix or of exhaustive thread-safety proof.

Explicit initial scope

The first model adapter supports sequential GPTQ for Llama and Qwen3 MoE with disk offload. Dynamic exclusions, rotation, endpoint quantization, GPTAQ/FOEM, adapters and other model families are rejected explicitly pending their own continuation coverage. Initial calibration capture is reconstructed for identity validation on restart. Checkpoint retention bounds generations; private attempt directories remain available for live model save references. Durability assumes a local filesystem with working atomic rename, locking and directory fsync, not arbitrary network/object storage.

Architecture, usage and operational limits are documented in docs/looper-checkpoint-redesign.md.

[Qubitium]

Hessian recovery and strict GPU topology follow-up

Implemented in f72d94be, on this PR's existing branch.

Recovery guarantees

Uncommitted layer work is discarded. Resume starts fresh Hessian collectors from the last committed layer boundary, replaying that layer's calibration batches and quantization steps. Partial Hessians and sample counts are not restored.

The subprocess tests now interrupt both early calibration and late calibration after earlier subsets have already been quantized, using actual SIGKILL. They check fresh zero-count collectors, calibration input hashes, final Hessian hashes/counts, committed cursor position, and byte-exact saved tensors. The dense and small MoE recovery matrix also covers graceful signals and injected errors.

On /monster/data/model/Llama-3.2-1B, CUDA, with PYTHON_GIL=0: killed layer 1's down projection after 2 forwards / 64 samples. Resume from cursor 1 reproduced all 105 finalized Hessians and 240 audited batch events in layers 1–15. All 482 saved tensors matched the uninterrupted baseline byte-for-byte, including dtype and shape. This uses small deterministic calibration data to test recovery, not quantization quality.

GPU-aware snapshots

• Record GPU count, visible GPU UUID/index mapping, capabilities, and ordered execution device pools in checkpoint identity.
• Reject topology mismatches before decoding continuation tensors, including multi-GPU to single-GPU, reordered/replaced GPUs, or changed pools.
• Preserve each continuation tensor's original device and device-valued metadata.
• Restore the looper's round-robin placement cursor and module/device map under its scheduler lock.
• Schema version 2 deliberately rejects older checkpoints without these guarantees.

Review also found cross-device Hessian summation followed worker arrival order. A three-device simulated regression reproduced different floating-point bits in 4 of 6 arrival permutations. Reduction now follows stable device-index order; all 6 permutations pass. Embedding count reduction uses the same ordering.

Validation and limits

Targeted suites: 204 passed, 4 skipped (66 checkpoint/extension/staging; 73 passed + 1 skipped store/concurrency/offload/helpers; 12 passed + 3 skipped topology/device; 53 Hessian/deduplication). Additional overlapping byte-exact recovery rechecks passed. Scoped Ruff and git diff --check passed.

This host has one physical GPU. Actual two-GPU tensor placement and dense/MoE E2E tests are included but explicitly skipped here. Mocked multi-GPU topology rejection, scheduler continuation, and reduction ordering tests passed; these are not a substitute for real multi-GPU execution. Free-threaded tests establish observed behavior, not a universal proof of thread safety or bitwise determinism across all kernels/hardware.

Environment caveat: installed Transformers 5.5.4 predates the repository's >=5.14 requirement and torchvision is absent. Runs used temporary, uncommitted import compatibility stubs under /tmp/gptqmodel_test_stubs (torchvision and the missing recurrent attention-mask symbol). These are not shipped in the PR. The full single-GPU Llama run preceded the final device-order sorting change; that change does not alter single-GPU reduction order, and subsequent tiny-model recovery and Hessian regression tests passed.

[Qubitium]

@okdshin Putting this to draft mode as major refractor and full quant method/format compat is currently happening so this branch will be unstable.

[Qubitium]

@okdshin Ready for testing

[Qubitium]

@okdshin Going to merge now. Passing local tests.. If you find bugs, please open new Pr.