feat(export): export each decoder layer as layerwise calibration finishes it - #2136
feat(export): export each decoder layer as layerwise calibration finishes it#2136Fridah-nv wants to merge 34 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesThe PR adds configurable, resumable layerwise Hugging Face checkpoint export. It writes per-layer safetensors shards, validates resume state and unsupported configurations, integrates PTQ calibration, centralizes artifact handling, adds a PTQ recipe, and expands GPU coverage. Layerwise HF export
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🟡 Moderate · up to This PR adds per-layer export and resume, but current behavior can omit resume artifacts for list-form configurations, silently produce an empty export when export_dir is set without layerwise mode, and potentially combine stale shards with shards from different weights. These gaps can yield incomplete or mixed checkpoints, so merge should wait for fixes or explicit owner acceptance. Sequence Diagram(s)sequenceDiagram
participant HFPTQ
participant LayerwiseCalibration
participant LayerwiseExporter
participant HFArtifacts
HFPTQ->>LayerwiseCalibration: configure export_dir and resume paths
LayerwiseCalibration->>LayerwiseExporter: export calibrated decoder layers
LayerwiseExporter->>HFArtifacts: write indexed shards and non-weight artifacts
HFPTQ-->>HFArtifacts: report the layerwise checkpoint
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2136 +/- ##
==========================================
- Coverage 78.95% 76.14% -2.82%
==========================================
Files 522 524 +2
Lines 60550 63715 +3165
==========================================
+ Hits 47810 48514 +704
- Misses 12740 15201 +2461
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
aa52dbd to
5f88cf2
Compare
_is_layerwise probed the recipe's algorithm with getattr(obj, "layerwise", None), but an algorithm loaded from YAML is a plain dict, where getattr always returns None. It therefore answered False for every layerwise recipe in the repo. The one thing it gates is whether --batch_size 0 skips auto batch-size probing, which its own comment says must be skipped because the probe "runs a full-model forward which defeats the point and can OOM on very large models". That protection has never engaged for the recipes it was written for. Replace it with an accessor that handles both shapes: dicts from YAML, and the config objects the deprecated --auto_quantize_* path still builds. Behaviour change: with --batch_size 0, layerwise recipes now use batch_size=1 instead of probing, which is what the existing comment intends. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…shes it A PTQ run that outlasts its GPU session loses every calibrated layer, and a completed one still pays for a second whole-model export pass over a full-precision intermediate checkpoint. Add layerwise.export_dir: each decoder layer's final quantized tensors are flushed to their own shard the moment calibration finishes with it, so the shards double as the resume artifact. Combined with layerwise.checkpoint_dir a restarted run skips layers already on disk instead of recalibrating them, and the per-layer weights.pt / quantizer_buffers.pt files are no longer written. When the last layer lands the directory is already a complete, loadable checkpoint, so export_hf_checkpoint() is skipped. Setting the field is the whole switch; hf_ptq.py rewrites it to --export_path, as it already does for checkpoint_dir. One layer per shard is what makes that work: a shard is only ever written whole, its name derives from the layer index, and the index is rebuilt at the end from the shards on disk. A crash can lose the layer in flight but never corrupt an earlier one, and a re-run overwrites in place rather than appending duplicates. Resident modules have no materialization window to discard export's damage, so transient_module_state snapshots the layer subtree's _parameters/_buffers/_modules and restores them, leaving calibration free to run every later layer through it. Tied-weight dedup is off here for the reason registry.py already turns it off for offload: data_ptr() cannot identify a tensor across an export that keeps rolling packed weights back. Weight-tied quantized modules are refused up front anyway. NVFP4 works because export_layer rediscovers the scale-fusion groups itself: the groups _fuse_shared_input_modules operates on -- q/k/v behind input_layernorm, gate/up behind post_attention_layernorm -- never cross a layer boundary, so a probe forward over one layer finds them. The probe uses that layer's real cached activations rather than the synthetic input the whole-model pass builds. Scope is resident, single-process models. AWQ and SVDQuant additionally need requantize_resmooth_fused_llm_layers' pre-quant-scale steps, which are still whole-model, so they stay refused along with accelerate offload, multi-process jobs, weight-tied quantized modules, multimodal models, MTP models and speculative decoding are refused with NotImplementedError before calibration starts; each would otherwise produce a silently different checkpoint rather than fail. A resumed run never recalibrates the layers it skipped, so the in-memory model is not valid for inference and layerwise.export_dir implies --skip_generate. Two pieces move out of this path to avoid duplicating what already exists: save_non_weight_artifacts() is extracted from the streaming exporter's tail, and FUSION_FREE_FORMATS moves to model_config.py, where _fuse_shared_input_modules had held the same set inline as a literal list. Verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, FP8, 256 experts per layer): 93,273 tensors, 0 mismatched, and identical again after a simulated mid-model resume. NVFP4 equivalence is covered by a GPU test. The NVFP4 test leaves o_proj unquantized: layerwise calibration leaves its input amax at 0 on every layer but the last, which no export path can write. That is a pre-existing bug, unrelated to per-layer export. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Per-layer export exists so a run that outlasts its GPU session keeps its finished layers, and offloaded single-GPU runs on very large models are exactly the multi-hour runs that hit session limits -- but offload was the one case refused. Calibration already handled it; only export did not. finalize()'s tail walked model.state_dict() directly, and _collect drops meta tensors, so an offloaded model's embeddings, norms and lm_head were skipped with no error. Give them the same per-module materialization window the streaming exporter uses. Tail collection splits in two: modules needing a window, then everything already resident, which on a non-offloaded model is the whole tail. Tie detection cannot run under offload -- data_ptr() cannot group weights that are not resident -- so it now says so instead of reporting a clean bill of health. Resolving ties by name would fix it properly, the same way ExportContext.__post_init__ already has a TODO for. Also fix the fusion gate added with NVFP4 support: it asked get_quantization_format(layer), which returns the first format found, so a layer with FP8 attention and NVFP4 experts reported fp8 and silently skipped fusing its NVFP4 groups. Use the per-module scan the model-level gate already uses. The existing NVFP4 tests could not catch this -- both were single-format layers -- so this adds a mixed FP8/NVFP4 case, which fails without the fix on mlp.up_proj.weight_scale_2. Verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, experts-only NVFP4, accelerate disk offload, max_memory 20GiB): 123,493 tensors, 0 mismatched, with embed_tokens, lm_head and norm correctly captured in the tail shard. Per-layer export is ~63s slower than the streaming export for this model (271s vs 208s). The cause is the per-layer fusion probe forward, not offload; the argument for this path is durability, not speed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Seven fixes from a review of the layerwise export path. Each would fail late or silently rather than at the point of the mistake. hf_ptq only retargets export_dir and runs the compatibility refusals on the mono-quantize path, so an AutoQuantize recipe carrying export_dir exported to the recipe's placeholder directory and then skipped export_hf_checkpoint(), leaving --export_path with no weights and printing success. Refuse that combination. set_layerwise_export_dir indexed algorithm as a dict, but detection accepts a list of algorithms too, so a list-shaped recipe died on a str index before calibration. Handle both shapes. The refusal loop listed --sparsity_fmt as the only route to the TRT-LLM exporter; int8_sq and encoder-decoder model_type reach it as well, and a Whisper model has discoverable decoder layers, so shards were written and then overwritten by a second checkpoint. resolve_checkpoint_dir hashed the config while it still held the recipe's placeholder export_dir, so two runs to different --export_path values shared one checkpoint dir and the second resumed against the wrong shards. Retarget first. _fusion_probe replayed a cached batch without the past_key_values reset _layer_forward_loop performs for the same tuples, so the probe could see kv_len at twice the attention mask width. save_file rejects two keys backed by one storage, and _collect's .cpu() is a no-op when the tensor is already there; copy aliases before writing, as _StreamingShardWriter already does. Finally, finalize()'s resident tail loop could still reach a module holding meta tensors, where packing raises deep inside the export handler. Raise there instead with the module name -- skipping it would drop weights silently, which is the failure this path exists to prevent. Re-verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, experts-only NVFP4, accelerate disk offload): 123,493 tensors, 0 mismatched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ise export LayerwiseExporter kept model, dtype and is_modelopt_qlora both on self and inside the ExportContext it builds from them, leaving two sources of truth for the same facts across 9 read sites. Build the context first in __init__ and read everything through it. No behaviour change. This is what docs/design/export-feature-scoping.md asks for in phase 1 -- construct the context once and thread it through, rather than rebuilding it or shadowing it per phase. Doing it now means that phase is already satisfied for this branch when the CheckpointExporter base class lands, instead of being work the refactor has to carry. Re-verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, experts-only NVFP4, accelerate disk offload): 123,493 tensors, 0 mismatched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Comments and docstrings only, no code change. State the reason and stop; drop restatements of what the code already says. Also corrects a stale claim in assert_layerwise_export_supported: it still said restricting to FUSION_FREE_FORMATS is what makes skipping requantize_resmooth_fused_llm_layers safe, which stopped being true when NVFP4 gained per-layer fusion. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Per-layer export presents its shards as the resume artifact, but the resume *point* does not come from them. _CheckpointState.from_folder returns `start = info[0] if info else 0`, read from checkpoint_dir's manifest, and nothing derives it from the shards on disk. So the manifest and the shards have to share a lifetime. They did not. Both shipped layerwise-export recipes default checkpoint_dir to /tmp/modelopt_layerwise_ckpt, and /tmp is container-local on a Slurm/Pyxis node. A run that outlasts its GPU session -- precisely the case this path exists for -- comes back to a wiped manifest, restarts calibration at layer 0, and overwrites every finished shard. assert_shards_present(0) passes trivially on the way through, so nothing warns: the run just silently redoes hours of work. Found on the Kimi-K3 run, where the manifest landed in /tmp/modelopt_layerwise_ckpt/Kimi-K3-bf16_4abe5702/ while the shards were on shared storage. At ~6 min per layer over 93 layers that would have cost a full session on every restart. Co-locate the checkpoint dir under --export_path when per-layer export is enabled, so the invariant is structural instead of something the user has to know. Only applied in that mode; otherwise checkpoint_dir is an ordinary resume directory and its placement is the caller's business. Verified on midi-K3: the manifest now resolves under <export_path>/.layerwise_checkpoint/<model>_<hash>/, and a run whose manifest is rewound to layer 3 reports "resuming layerwise calibration from layer 4/8" and skips the finished layers rather than recalculating them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 9d0b846)
Name-based tied-weight dedup (#2194) removed tied_cache and moe_tied_cache from ExportContext, so constructing one with them now raises TypeError. Drop them. That change also supersedes this path's tie detection. Grouping quantized modules by weight.data_ptr() sees nothing once the weights are on meta, so under offload the check passed vacuously and only emitted a warning saying so. TiedWeightMap is keyed by name and survives offload, so group by that instead and delete the warning. data_ptr remains as a fallback for transformers <5.0, which publishes no map. The alias-key dedup in _collect deliberately stays as it is, matching the TODO(tied-map) note the streaming path carries: swapping that one needs offload-specific validation (meta tensors, per-tensor order, disk round-trip), and this path is per-tensor and offloaded too. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…is gone The shards are the resume artifact but not the resume point: restarting at layer K needs the cached activations feeding it and every skipped layer's output_meta, neither reconstructible from quantized weights. So the point comes from the checkpoint manifest, which makes one state dangerous -- shards on disk, manifest gone. start_layer is then 0, assert_shards_present checks an empty range and passes, and calibration recalculates and overwrites every finished layer without a word. A checkpoint_dir on ephemeral storage is the usual cause, and it is exactly the long run this feature exists for. Raise instead. A manifest that merely lags the shards stays allowed: that is a mid-window interrupt, and re-exporting those layers is idempotent. Adapted from 64ea3b8 on the stacked Kimi-K3 branch, without its multimodal machinery. Its guard also assumed checkpoint_dir is set, which raises TypeError in the export-without-resume mode; that case is now skipped, since with no checkpoint_dir there is no resume to lose and re-exporting is documented behaviour. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
06ace1e to
8c1673f
Compare
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The feature addresses a real durability problem: avoiding loss of already-calibrated decoder layers and the second whole-model export pass. Existing alternatives are (1) extending/reusing the existing streaming exporter and _StreamingShardWriter in modelopt/torch/export/unified_export_hf_streaming.py, (2) retaining the existing layerwise checkpoint artifacts and invoking that streaming exporter as the final/resume phase, or (3) factoring a shared per-module shard writer around the existing ExportContext/export-handler path. The PR body acknowledges the streaming exporter and its tail-pass duplication, but does not justify why a separate 549-line LayerwiseExporter is preferable or why the existing writer cannot be extended; this remains an architectural concern for a 1,286-line PR. More importantly, I found two correctness issues in the core path: the hf_ptq integration calls an undefined helper, and a completed manifest is treated as a fresh run, leaving a crash window after the last layer that defeats the durability claim. New-file license headers match LICENSE_HEADER, and the GPU equivalence/resume tests are useful, but they do not cover either failure below.
Additional comments (outside the PR diff):
examples/hf_ptq/example_utils.py:1193— > Bot comment.
_layerwise_checkpoint_dir_location is not defined or imported anywhere in this file/repository. Consequently the documented hf_ptq.py path reaches colocate_layerwise_checkpoint_dir() and raises NameError whenever the layerwise config has a checkpoint directory (including the new shipped recipe). Please implement/reuse the intended lookup and add an example-utils or hf_ptq integration test, since the direct mtq.quantize GPU tests bypass this code.
Two CI failures, both introduced here. build-docs: assert_layerwise_export_supported's docstring used a `.. todo::` directive. The convention was copied from _CheckpointState, which gets away with it by being private and never rendered; this function is in __all__, so autodoc renders it and sphinx.ext.todo is not enabled -- "Unknown directive type", fatal under --fail-on-warning. Plain prose instead. code-quality: colocate_layerwise_checkpoint_dir called _layerwise_checkpoint_dir_location, which does not exist on this base. The function came from the stacked Kimi-K3 branch, written when that helper returned a (shape, current) tuple; main has since replaced it with _layerwise_checkpoint_dir returning the directory. Rewritten against the current API, which also drops the dead legacy-flat branch. This was not only a type error: it would have raised NameError the first time a recipe set checkpoint_dir. Neither showed up locally because pre-commit was run over changed files while CI runs --all-files, and the docs were never built. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — 0 CRITICAL · 0 IMPORTANT · 1 SUGGESTION
Claude review passed — no blocking issues found. LGTM
Scope: full review (the trigger comment carried no scoping instructions). 16 files changed. Reviewed all 8 modelopt/ files, both examples/hf_ptq/ files, the new recipe YAML, modelopt_recipes/ptq.md, CHANGELOG.rst, and both new test files. For context I re-read the unchanged sides this touches: requantize_resmooth_fused_llm_layers, _fuse_shared_input_modules, collect_shared_input_modules, preprocess_linear_fusion, get_quantization_format, _StreamingShardWriter.add, _CheckpointState.{setup_resume,full_restore,save}, and export_quantized's is_tensorrt_llm_export predicate. Nothing deliberately skipped.
Prior rounds: both IMPORTANT findings resolved
| Prior finding | Status |
|---|---|
| IMPORTANT: derived resume dir accumulates one full activation cache per layer, unbounded | Fixed — _CheckpointState._prune_stale_next_inputs(keep=layer_idx) runs at each committed boundary under save_layer_state=False, leaving exactly one next_inputs.pt. test_layerwise_export_replaces_resume_artifacts pins that a completed run leaves none at all, which holds because the final save() gets next_layer_inputs=None and keep points at a layer that never had one. |
IMPORTANT: _warn_on_unsynced_moe_gate_up in finalize() guaranteed to false-alarm |
Fixed — the call is gone, with a comment at layerwise_export.py:433-434 recording why (the sync lands in the shard and is rolled back on the live model), so the next person doesn't re-add it. |
I verified the prune's commit ordering rather than taking the comment at its word: next_inputs.pt → manifest (atomic tmp + os.replace) → prune. Placing the prune after the manifest is what makes it safe — if the manifest never lands, the previous boundary's next_inputs.pt is still on disk and setup_resume finds it; if it does land, the new boundary's copy is the one keep retains. Checked the save_every > 1 and mid-export_layer crash cases too: a truncated shard always sits at an index >= start_layer, so it is re-exported and overwritten rather than trusted.
Checks that came back clean this round
Recording these so they aren't re-litigated. Two of them I opened intending to file a finding and closed after reading the baseline:
_fuse_unrouted_expertsis a faithful port, including its cost profile. I went in expecting to flag theO(routed_groups × num_experts)sibling replay — for each routed expert group it walks every expert id, so an unrouted expert is re-fused once per routed group.requantize_resmooth_fused_llm_layers(unified_export_hf.py:518-540) has the identical loop with the identical redundancy, so this is parity with the path it must match, not a regression.preprocess_linear_fusionis idempotent (max/mean over amaxes), so the repetition is wasted work rather than wrong scales. Also confirmed the two loops agree oncount=1for the group key and nocountfor the members.- Probe-strength difference does not change the exported scales. The whole-model path probes with
torch.ones([1, 2]);_fusion_probereplays one real 16-token calibration batch, so it routes more experts. That changes which groups land infused_linearsbut not the union of experts fused, because the sibling walk covers all ids0..N-1from any single present template. Equivalence holds as long as at least one expert routes, which is true for both probes. - Refusal list has no hole against the current code.
is_tensorrt_llm_export(hf_ptq.py:882-886) is exactlymodel_type in [t5, bart, whisper] or sparsity_fmt != "dense" or "int8_smoothquant" in args.qformat, and--export_fmtis deprecated tohfat line 1824.assert_layerwise_export_compatiblecovers all three plus VLM, MTP, spec-dec,--cast_mxfp4_to_nvfp4and--vllm_fakequant_export. Matching bothint8_sqandint8_smoothquantis over-broad by one token, which errs in the safe direction. - Format gate is sound and conservative.
QUANTIZATION_NONE is None, so it really is insideFUSION_FREE_FORMATSand a fully-disabled layer takes the no-probe path.SUPPORTED_FORMATSis{None, fp8, fp8_pb_real, nvfp4}— every other constant (fp8_pc_pt,fp8_pb_wo,w4a8_*, mxfp4, int8) is refused rather than silently mis-exported.assert_formats_supportedruns both pre-calibration and per exported layer, which is what closes the AWQ/SVDQuant hole from an earlier round. - Finalize-only path. Returns before
_patch_all_layers, so no patching is left dangling;assert_shards_present(num_layers)gates it; a complete manifest can't coexist with a truncated final shard because the shard write precedes the manifest write. - Tail collection has no duplicate-key or double-dispatch path. Loop 1 records
handled_idsfor every descendant of a materialized module andseen_keysfor every key it collected; loop 2 skips both, and themodel.state_dict()sweep skipsskip_prefixes(original names, matchingstate_dict()'s namespace — the_name_mapperrename happens downstream in_collect).requires_weight_materializationinspecting only a module's own_parameters/_buffersis still what prevents a container-level window pulling decoder tensors into the tail shard. transient_module_state+writeback=calib_mutates_weights=True. Still only a live-model concern, never a checkpoint one:finalize()skips_decoder_owned_idsand_write_indexreads shards from disk, so an in-place param mutation written back to the offload store cannot reach the exported checkpoint.save_layer_state=Falsemeansfull_restorereturns early, so it can't reach a resume either.save_layer_stateback-compat. Drift-checked withif ckpt_value is not None, so manifests written by earlier versions still resume.- Streaming-writer parity for storage aliasing.
_copy_storage_aliasesusestensor.data_ptr(), the same key_StreamingShardWriter.adduses (unified_export_hf_streaming.py:96-99), so the new path is no weaker than the one it mirrors. save_non_weight_artifactsextraction is behaviour-preserving line for line, and still runs before_write_hf_export_config, which reads theconfig.jsonit wrote.- New tests earn their keep.
test_index_resolves_every_key_to_the_shard_holding_itcloses a real gap — tensor-for-tensor equality againstexport_hf_checkpointnever exercises theweight_map, so a shard misrouting would have compared equal and still failed in vLLM. Thenvfp4_qdq_from_prev_layercase covers the axis that produced the previous round's IMPORTANT.
One note on the PR description
The second "pre-existing bug" under Additional Information — that export_quantized's TRT-LLM branch tests "int8_sq" in args.qformat and is therefore unreachable — does not match the current code: line 885 reads "int8_smoothquant" in args.qformat, and that is the only occurrence in the file. Worth dropping from the description so it doesn't get filed as an issue that no longer exists. The o_proj input-amax-0.0 finding does still check out and is worth its own issue.
Risk
Low. export_dir defaults to None and every new path is gated behind it; the only unconditional behaviour change remains the _is_layerwise dict-vs-object fix and its documented --batch_size 0 consequence, which is called out in both the PR body and the changelog. The resume/commit ordering, the shard-identity binding, and the fusion-equivalence argument all hold up on re-reading, and this round's delta (the prune plus the removed MoE warning) is correct and covered by a test. The single SUGGESTION is a silent flag override plus a half-stated comment — non-blocking.
The separate-exporter architectural question raised by the other bot review is still the author's and CODEOWNERS' call, not a correctness blocker; nothing in the code forecloses the shared-writer refactor later.
Added to mirror requantize_resmooth_fused_llm_layers, but sync_moe_gate_up_amax already walks every expert routed or not, and gate/up is the only input-sharing group among expert linears -- so the replay can only revisit pairs that sync has handled. Three runs on Qwen3-30B-A3B with static expert weights, replay on and off, were bitwise identical to the whole-model export. The hazard it was meant to cover is real but unreachable here: a static quantizer's weight_scale_2 reads global_amax, which sync does not write, but promote_static_block_weight_quantizers has no caller on this path, so global_amax stays unset and weight_scale_2 falls back to the amax sync does write. Fixing that belongs in the sync, not in a per-layer replay that no test can reach. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…reason The flag flips silently, so a user who expected the post-quantization sample just sees it missing; every other consequence of export_dir announces itself. The comment also predated the tail-mutation fix -- finalize() converts the non-decoder modules in place on every run, not only a resumed one, which is how model_calib already words it. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ule_state The snapshot restored a layer nothing reads again. After its shard is written the next layer's inputs have already been captured, resume skips finished layers in favour of their shards, full_restore is a no-op once the shards are the resume artifact, and ckpt.save reads no weights when save_layer_state is False. It was protecting the returned model, a guarantee finalize() then broke anyway by converting the tail in place. It cost a per-submodule snapshot of _parameters/_buffers/_modules plus a clone of every buffer, per layer -- on a 256-expert MoE that is the whole layer's buffer set cloned 40 times. Verified rather than argued: with the restore disabled, 23 of 24 GPU tests passed unchanged, and the one failure was the test asserting the guarantee being dropped. hf_quant_config.json and config.json are byte-identical either way, which the tensor-only test helper would not have caught. Qwen3.6-35B-A3B NVFP4 experts-only under accelerate disk offload -- the configuration where the surrounding persistent_materialization window can write back -- matches the whole-model export on all 123,513 tensors. That test is replaced by one pinning what actually matters: the checkpoint is unaffected, and the returned model is in export form. The wording elsewhere said export converts 'the non-decoder modules', true only while layers were restored; it now says the model. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…p where it is used The identity guarded against resuming one run's shards under another run's manifest. hf_ptq already separates those: resolve_checkpoint_dir appends a hash of the quant config to checkpoint_dir, so a changed config finds no manifest and assert_no_orphan_shards raises. The manifest itself drift-checks num_layers, save_every, calib_mutates_weights and save_layer_state. What was left is a library caller reusing one checkpoint_dir across differing configs at the same layer count -- their own responsibility, and not worth a sidecar file, a digest and two tests to state. _name_to_module moves into finalize(), its only consumer. It was built in __init__ and held for the whole run, which on a 256-expert MoE is ~100k entries alive across hours of calibration for a map used once at the end. get_quant_config stays in __init__ -- kv_cache_quant_algo still comes from it, and deriving that another way would give the exported tensors and the exported config two different sources for one value. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…reasons _fuse_shared_inputs said what it operated on rather than what it did; it unifies scales across modules that share an input, so it is now _unify_shared_input_scales. export_layer names its two phases once -- settle the scales, then change the representation -- with the ordering constraint that ties them, in place of prose scattered over each call. The rest is subtraction: the parity table became the sentence it was making, and comments that restated their own code are gone. What is left answers why -- why a set and not get_quantization_format, why not get_kv_cache_dtype, why the layer identity check raises instead of asserting. 550 -> 518 lines, 56 -> 42 comment lines. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Neither fired. _copy_storage_aliases cloned tensors sharing storage because save_file rejects them -- measured across five configs and 831 tensors, it cloned none. Silently copying an alias we cannot classify is the wrong default anyway: a declared tie is legitimate, two logically distinct tensors on one storage is a bug, and cloning writes duplicate data under both names either way. save_file raising is the louder answer. The tied-weight dedup went with it, because nothing reaches it. Conversion registers every nn.Linear and nn.Embedding, so both sides of a tie_word_embeddings model carry quantizers and _tied_quantized_modules refuses the model outright -- including under a base_disable_all config, since "enable: False" leaves the quantizer module in place. The refusal message now says so, so nobody tries excluding lm_head to get past it, and a test pins both shapes. The streaming exporter keeps its equivalent for now; that is a main-path change and belongs in its own PR. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
RB:
need pre-quant-scale steps that are still whole-model -> Is this true? I dont think so. pre-quant-scale fusion is entirely limited to one layer (the fusion is to a previous linear layer or pre_norm layer , but in the same decor layer) Can you check? |
🐝 Looking into this now — I’ll verify the actual fusion/resmoothing scope against the current head. |
Re: this RB question: you’re right. The operations are not inherently whole-model: The current refusal is therefore a limitation of this implementation, not a true cross-layer dependency. The whole-model helper ( I’d change the wording from “steps ... are still whole-model” to “not yet implemented in the per-layer exporter”; the unsupported-format guard remains appropriate. |
| @model_validator(mode="after") | ||
| def validate_layerwise_checkpoint_dir(self): | ||
| """Raise if layerwise.checkpoint_dir is set but layerwise.enable is False.""" | ||
| if self.layerwise.checkpoint_dir is not None and not self.layerwise.enable: | ||
| raise ValueError( | ||
| "layerwise.checkpoint_dir requires layerwise.enable=True. " | ||
| "Set layerwise.enable=True or remove layerwise.checkpoint_dir." | ||
| ) | ||
| """Raise if a layerwise directory is set but layerwise.enable is False.""" | ||
| for field in ("checkpoint_dir", "export_dir"): | ||
| if getattr(self.layerwise, field) is not None and not self.layerwise.enable: | ||
| raise ValueError( | ||
| f"layerwise.{field} requires layerwise.enable=True. " | ||
| f"Set layerwise.enable=True or remove layerwise.{field}." | ||
| ) | ||
| return self |
There was a problem hiding this comment.
Should we just get rid of this check?
It does not make sense to me. The behavior should be layerwise will be done if layerwise.enable=True
There was a problem hiding this comment.
This is checking is to refuse combinations like layerwise.enable=False + checkpoint_dir/export_dir specified. I think this is the validators' job. Do you agree?
| self._ctx = ExportContext( | ||
| model=model, | ||
| dtype=_resolve_export_dtype(model, dtype), | ||
| is_modelopt_qlora=is_modelopt_qlora, |
There was a problem hiding this comment.
Do we need to support QLoRA? It is not a heavily used feature. Will removing QloRA support simplify any thing? Wdyt?
There was a problem hiding this comment.
RB: What are your thoughts?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Acknowledged — I’m checking the QLoRA-specific export path and its real maintenance cost.
There was a problem hiding this comment.
🤖 Bot comment.
I’d remove layerwise QLoRA support for this PR. It is more than a flag: correct QLoRA export needs the PEFT base_layer skip plus the distinct base-model/adapter artifact layout. The layerwise caller never passed is_modelopt_qlora to LayerwiseExporter, so this path was not actually wired or covered as QLoRA support; the latest change simply removes that dangling parameter.
The simplification is worthwhile, but I’d ask for an explicit early NotImplementedError when the model is QLoRA (and a focused test). Otherwise a QLoRA model can enter layerwise export as ordinary export and produce an invalid/mislaid checkpoint rather than a clear unsupported-feature error. Normal whole-model QLoRA export can remain supported.
There was a problem hiding this comment.
I don't think ExportContext as a general structure should carry is_modelopt_qlora (a specific module check), and the flag is scattered in the export code. We can do a clean up with another PR.
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Acknowledged — checking the QLoRA/export implications.
There was a problem hiding this comment.
🤖 Bot comment.
I’d scope this feature to non-QLoRA rather than carry QLoRA support. QLoRA export has special semantics—skip PEFT wrappers/adapters and strip base_layer from checkpoint keys—so it is not just a plumbing boolean. The current tip has already removed LayerwiseExporter’s QLoRA argument, meaning layerwise export would otherwise proceed with False and can emit incompatible keys/adapters. Please add an early, explicit NotImplementedError for ModelOpt QLoRA when export_dir is used (and a test). This simplifies the new path; it does not require removing existing whole-model QLoRA export support.
…sure export_layer took a probe_forward callable that the driver built per layer, so the knowledge of how to replay a layer was split across both files. It now takes the cached inputs and builds the probe itself: one argument, and the KV-cache reset lives next to the code that needs it. The exported tensors are unchanged by construction -- same first input tuple, same cache reset, same call -- and measured: Qwen3.6-35B-A3B NVFP4 experts-only under accelerate disk offload matches the whole-model export on all 123,513 tensors, as it did before the change. Two more things move to where they are used. _decoder_owned_ids is recomputed in finalize() rather than snapshotted in __init__, which also fixes a staleness the snapshot had: calibration installs SharedQuantState modules inside the decoder layers afterwards, and 16 of them were missing from the set, leaving the tail pass to dispatch handlers on modules a layer shard already owned. Harmless today because no handler matches that type, but it would not stay harmless. extra_state_dict is gone from finalize(): it is the hand-off for orphaned MTP weights, both call sites pass nothing, and MTP is refused on this path. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
There was a problem hiding this comment.
Overall looks very nice!!
There was a problem hiding this comment.
RB: Can you review this file? Are there any redundancies or scope for any further simplifications? Review applying this repos coding guidelines.
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Acknowledged — reviewing this file against the repository coding standards.
There was a problem hiding this comment.
🤖 Bot comment.
I found one maintainability cleanup worth requesting; no independent correctness blocker in this file.
It has 17 function-local imports (for example lines 50, 73, 114, 155–161, 231–234, 320–331, 446, and 465). CONTRIBUTING asks for module-top imports unless a concrete circular-dependency, optional-dependency, or unusually-heavy-import reason is documented. Please consolidate the normal internal/safetensors imports at the top, retaining only genuinely necessary deferred imports with a short reason. This makes the module's dependencies visible and removes repeated local import plumbing.
I would not combine completed_layers, assert_shards_present, and assert_no_orphan_shards: they encode distinct resume invariants. A follow-up-only comment/docstring trim could help readability, but it is not worth broad refactoring in this PR.
assert_no_orphan_shards accepted any manifest that parsed, so a hand-edited or foreign manifest.json missing last_completed_layer let calibration restart at layer 0 and silently overwrite finished shards -- the exact outcome the guard exists to prevent. The reconciliation now runs in one helper that reads the manifest once and decides between finalize-only, resume, and refuse; the guard is only ever reached without a usable record, so its parameter is gone. Also from the same review round: - fold the gate/up amax sync into the renamed _unify_shared_quantization_params, so one call site owns every scale the fused kernels share - move the export-form warning into finalize(), which is what leaves the model in that state -- the finalize-only resume path never warned before - move the recipe layerwise traversal out of hf_ptq.py into example_utils.recipe_layerwise_blocks, normalizing both recipe shapes to dicts - drop LayerwiseExporter's is_modelopt_qlora parameter: no caller can set it, so QLoRA was never wired in and the parameter implied otherwise - point layerwise_calibrate's docstring at LayerwiseConfig instead of restating it Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…s its length The probe replayed the captured kwargs after calling Cache.reset() on the past_key_values they carry. That call does not clear anything: transformers' DynamicLayer.reset() zeroes the key/value tensors "while preserving the objects", so get_seq_length() stays at the captured length and the next update() appends to it. The probe therefore attended over twice the keys the captured attention mask was built for. Invisible in the tests because a tiny unpadded batch reaches the layer with attention_mask=None, where the extra keys only change values the probe discards. A padded batch materializes a 4D mask and the forward raises "The size of tensor a (64) must match the size of tensor b (32) at non-singleton dimension 3". A throwaway forward run only to fire hooks has no use for a cache at all, so pass none. The identical reset() call in the calibration replay predates this branch (#1223) and is PR #2248's to fix. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The probe now runs layer_inputs[0] verbatim instead of substituting a fresh past_key_values. A cache only reaches here when the calibration forward loop left caching on, and by then the layer's amax is already wrong -- every replay attended over the keys its own earlier replay wrote. Papering over that in export hides it behind a checkpoint that loads and runs. PR #2248 owns the real fix and warns at the calibration entry point, which is where a user can act on it. Letting the probe raise is the better failure until then: loud, and pointing at the run that produced it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
A pass over the files this branch touches, dropping comments that restate the code they sit above and docstring paragraphs a reader gets from the signature. What stays is what a reader cannot recover: rejected alternatives that look correct (get_kv_cache_dtype not recursing, get_quantization_format stopping at the first hit), ordering constraints between the export passes, and why decoder_owned_ids cannot be snapshotted in __init__. 36 lines net. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
4 participants
What does this PR do?
Type of change: new feature
Layerwise calibration can already resume, but only through a full-precision scratch checkpoint, and a completed run still pays for a second whole-model export pass over it.
layerwise.export_dirwrites each decoder layer to its own quantized shard as soon as calibration finishes with it, so the directory is a complete, loadable checkpoint when the last layer lands andexport_hf_checkpoint()is skipped. The shards are the resume artifact: a restarted run reuses layers already on disk instead of recalibrating and re-exporting them, so no full-precision copy of the model accumulates. The resume directory beside it holds only the current boundary's cached activations and the per-layer output shapes.Setting the config field is the whole switch — no CLI flag.
hf_ptq.pyrewrites its value to--export_path, and derives the resume directory (<export_path>.layerwise_resume) when you haven't chosen one.One shard per layer is what makes resume safe: shards are written whole and named from the layer index, so a crash can lose the layer in flight but never corrupt an earlier one, and a re-run overwrites in place.
Because a resumed run never recalibrates the layers it skipped, the in-memory model is not valid for inference afterwards; the field implies
--skip_generate.Includes a pre-existing
mainfix this depends on:_is_layerwiseusedgetattron an algorithm that YAML parses as a dict, so it answeredFalsefor every layerwise recipe in the repo and the batch-size probe it gates was never skipped. Behaviour change:--batch_size 0now yieldsbatch_size=1for layerwise recipes, as its comment intends. Detection also now scans every algorithm entry rather than the first, so a list-form recipe whoselayerwiseblock is not first is recognised as layerwise — same batch-size consequence. Nothing else on the non-fused paths changes:FUSION_FREE_FORMATSis the exact set the inline list held,save_non_weight_artifactsis a lift of the streaming exporter's own block, and the calibration-loop changes are gated on an exporter being present.Refused before calibration starts, since each would otherwise produce a silently different checkpoint rather than fail:
sync_tied_input_amaxmerges amaxes across a partner that may be uncalibrated or already writtenexport_dir--vllm_fakequant_export, non-dense sparsity,int8_smoothquant, encoder-decodermodel_type--export_pathexport_diron more than one algorithm entry, or on any but the lastShards are also bound to the run that produced them (
.layerwise_export.json: model class, layer count, formats, KV-cache format, and a digest of the resolved quant config), so one run's manifest cannot finalize another's shards. Source weights are not digested — that would mean reading the whole model — so differently-trained weights at the same path compare equal.Why a separate exporter
Three reuse paths were considered before adding one:
_StreamingShardWriter. It buffers bymax_shard_sizeinto__shard_part_*temp names and renames to canonical names only infinalize(), once the shard count is known. The resume invariant needs the opposite: a stablemodel-layer-00007.safetensorscommitted when layer 7 finishes, so "shard exists" means "layer done" across a restart. Forcing a per-layer flush still leaves temp names, finalize-time renaming, and an in-memory_key_to_part— every method would change.save_layer_stateis off under per-layer export, but withcalib_mutates_weights: false(the shipped recipe) the checkpoint holds just amax buffers either way.ExportContext. The right long-term shape, but it touches all three existing export paths; doing it here makes this change larger, not smaller.Happy to take a different call on this — flagging it for maintainer sign-off rather than assuming it.
Usage
Interrupt and rerun the same command: calibration resumes from the last committed layer, finished shards are reused, and a run that had already finished every layer only re-runs
finalize().Testing
Each row exports the same calibration two ways — per-layer, and whole-model via
export_hf_checkpoint()— and compares them tensor for tensor. All re-run on the current head.nvfp4_staticweights) +mse, offloadSIGKILLafter 25/48 layers, then resumedRefusals verified on real checkpoints, each writing zero shards and never reaching calibration — the "refused before calibration starts" claim above, demonstrated rather than asserted: multimodal and MTP (Qwen3.6-35B, the MTP case on a text-only view since the multimodal gate fires first), tied embeddings (Qwen3-0.6B), and multi-process (2-rank
torchrun --use_fsdp2, Llama-3.1-8B).Served, not just compared. Under vLLM 0.27.1 (Marlin NVFP4 kernels, SM 8.9): the 30B checkpoint exported three ways — whole-model, per-layer, per-layer-resumed-after-a-kill — and the 8B exported both ways all load and produce identical greedy generations, 4/4 prompts within each model.
Index integrity on every checkpoint above: each
weight_mapkey resolves to the shard actually holding it; 0 missing, 0 extra, 0 mis-routed. Tensor equality alone never exercises that, and it is the one artifact per-layer export builds differently.Resume state stays bounded: 332 KB beside 22 GB of shards on the 35B, 396 KB beside 19 GB on the 30B — the committed boundary's activations only, not one set per layer.
Not covered: the
trust_remote_code*.pycopy path. Nemotron-Nano-12B-v2-Base fails with a CUDA illegal memory access on these cards, on the whole-model baseline too, so it is an environment limit rather than a result.24 GPU tests in
tests/gpu/torch/export/test_layerwise_export.py. The equivalence oracle is a cross-product: {FP8, NVFP4, NVFP4 +get_qdq_activations_from_prev_layer, mixed FP8/NVFP4, KV-cache} × {fresh, resumed-after-interruption}, each compared tensor-for-tensor againstexport_hf_checkpoint. Plus MoE export; resume fail-fast; resume artifacts replaced and pruned; complete-manifest finalize-only; shards-without-manifest refusal; shards-from-a-different-run refusal (format and module selection); identity-without-shards does not block a rerun; export-does-not-mutate-the-model; index routes every key to the shard holding it; AWQ refusal (from config, and after calibration); export-without-checkpoint_dir.5 unit tests in
tests/examples/hf_ptq/test_example_utils.pycover the list-valuedalgorithmshapes: which entry owns export, whosecheckpoint_diris derived, per-entry resume bases, and both ambiguity refusals.tests/gpu/torch/export/149 passed / 3 skipped (pre-existing env skips) ·tests/unit/recipe282 ·tests/unit/torch/export172 ·test_layerwise_calibrate33 ·test_example_utils36 · pre-commit clean.Also verified: the exported directory reloads through
AutoModelForCausalLMand runs a forward.Not a speed win: per-layer export was slower than the streaming export in one offload pairing (271s vs 208s, the per-layer fusion probe), though those runs shared GPUs so the magnitude is not cleanly measured.
Before your PR is "Ready for review"
export_dirdefaults toNone; existing paths unchanged when unset, except the batch-size change noted above.CONTRIBUTING.md: N/AAdditional Information
Pre-existing bug found on the way, not fixed here. Layerwise calibration leaves
self_attn.o_proj's input amax at0.0on every layer but the last, so a full-NVFP4 layerwise model cannot be exported by any path.get_qdq_activations_from_prev_layer=Trueavoids it, pinning the cause to the pre-calib_funccapture pass — which also explains why only the last layer, the one that skips it, is correct. That combination now works with per-layer export (it asserted on layer 0 until review caught it). Hidden until now because the shipped NVFP4 layerwise recipes are experts-only; the NVFP4 tests here excludeo_projfor the same reason. Deserves its own issue.