Don't rewrite packed modules as float weights when resuming

[Shubham-Padkonde]

Fixes #2350

Problem

When a run is resumed through AR_RESUME_DIR, the blocks a previous run already quantized are skipped, so their modules are still plain nn.Linear in the model tree while their packed tensors sit in the shard files the crashed run flushed.

ShardWriter.finalize()'s remaining-weights pass deduplicates by exact tensor name, so <layer>.weight never matches the saved <layer>.qweight: the stale floating-point weight is written on top of the packed one. Every module of every skipped block ends up in the checkpoint twice, with no warning. The issue reports a Qwen3.8-27B W4A16 artifact coming out at 47.4 GB instead of ~17.7 GB, with 318/407 modules duplicated.

There is a second half to it: discovery of the previous run's shards (_discover_existing_shards) is deferred to the first _flush_shard() call, which finalize() only reaches after the remaining-weights pass. So that pass ran with an empty _all_saved for anything the crashed run wrote — even tensors whose names match exactly would have been written twice.

Fix

• Adopt the previous run's shards at the start of finalize(), before deciding which weights are still missing. output_dir is already the final path there, which is the reason the call was deferred out of __init__ in the first place.
• Skip a module's unpacked weight when packed tensors for that module are already saved (qweight / weight_packed), and log how many were skipped so this is no longer silent.

Modules that were never packed are unaffected, and non-resuming runs never reach the new discovery call because it stays gated on AR_RESUME_DIR.

Review follow-up

Resuming an export can leave plain model weights in memory while their packed replacements already exist on disk. Adopt prior temporary or finalized shards before the remaining-weight pass, compare names in the serialized checkpoint namespace, and skip unpacked weights whose packed replacements are already saved.

The review follow-up also includes recovered tensors in the index's total_parameters and total_size. Safetensors counts come from validated header metadata without materializing payloads; PyTorch shards are inspected on the meta device. Ordinary non-resuming behavior is unchanged.

Original validation

The results in this section cover the initial implementation; current statistics validation is recorded below.

test_finalize_skips_unpacked_weight_of_resumed_packed_module writes a shard in the crashed run's temp layout containing linear.qweight / linear.scales / linear.bias, points AR_RESUME_DIR at it, and finalizes a writer whose model still holds the plain nn.Linear.

Without the change:

AssertionError: the packed module must not also be saved as a floating-point weight
1 failed, 7 passed

With it, all 8 pass and proj_out.weight — a module that was never packed — is still saved.

$ pytest test/unit/common/utils test/unit/common/export/test_export_utils.py
5 failed, 923 passed, 6 skipped

The same 5 failures (test_alg_ext, test_layer_config_resolution, test_weight_handler, two TestIsImmediateSavingMode cases) reproduce on an unmodified main with the same selection — 922 passed there, so this change adds a test and breaks nothing. black, isort and ruff are clean.

Current statistics validation

• All 12 resumed-shard cases fail their new index-total assertion before the statistics fix.
• All 65 shard-writer and export-utility tests pass afterward on Linux with CPU PyTorch 2.14.0 and safetensors 0.8.0.
• Coverage includes temporary, single-file and finalized multi-shard layouts, both serialization formats, transformed names, mixed dtypes, scalar and empty tensors.
• All applicable repository pre-commit checks and the staged whitespace check pass.
• The full model/hardware suite was not run for this follow-up; hosted CI must validate published commit 007955ef5ee1ba8c702e378d0aaffd32838bdf0c.

The published commit is the contributor-signed commit and matches the exact two-file patch tested above.

---

Disclosure: this change was written by Claude Code (Claude Opus 5) working as my agent, at my direction. The test, lint and reproduction output quoted above comes from real runs in my local environment; I am accountable for what is submitted here and will follow up on review feedback.

🤖 Generated with Claude Code

Codex prepared and tested the checkpoint-recovery and statistics review follow-ups.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Copilot]

Copilot review overview

🟡 Changes recommended

Unresolved critical resume-detection issues and incomplete shard metadata remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 2

Open (2)

• Finalized shards are not adopted during resumable checkpoint discovery · New
• Transformed layer names are missed when detecting resumed packed modules · New

What changed in this PR

This PR fixes duplicated floating-point weights when resuming packed shard exports.

Changes:

• Discovers prior shards before finalization.
• Skips stale unpacked weights for packed modules.
• Adds regression coverage.

File	Summary	Review findings
auto_round/​compressors/​shard_writer.py	Adds resume shard discovery and packed-weight filtering.	Critical (3 votes): Compare transformed layer names when detecting resumed packed modules. Moderate (1 vote): Include adopted shards in index size and element statistics. Critical (1 vote): Safely adopt finalized resumable shards to prevent overwriting prior output.
test/​unit/​common/​utils/​test_shard_writer.py	Adds regression coverage for resumed packed-module handling.	No findings.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[xin3he]

Nice catch! Thanks

[xin3he]

@Shubham-Padkonde Could you please address these high-priority comments?

[Shubham-Padkonde]

Published the signed checkpoint-recovery follow-up in f31a1b2. It handles finalized shards and transformed resume names, with regression coverage. The prepared patch passed 65 tests and pre-commit checks before sign-off; hosted CI still needs to validate the new commit.

Prepared with Codex assistance.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Shubham-Padkonde]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).