feat(export): export each decoder layer as layerwise calibration finishes it - #2136
Draft
Fridah-nv wants to merge 4 commits into
Draft
feat(export): export each decoder layer as layerwise calibration finishes it#2136Fridah-nv wants to merge 4 commits into
Fridah-nv wants to merge 4 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. |
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
Contributor
Author
|
/claude review |
Contributor
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2136 +/- ##
==========================================
- Coverage 78.38% 77.91% -0.47%
==========================================
Files 522 523 +1
Lines 60342 60705 +363
==========================================
Hits 47297 47297
- Misses 13045 13408 +363
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:
|
Fridah-nv
force-pushed
the
fridah/layerwise-fused-export
branch
from
August 10, 2026 22:08
e62644b to
c904e6c
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>
Fridah-nv
force-pushed
the
fridah/layerwise-fused-export
branch
from
August 11, 2026 00:11
c904e6c to
aa52dbd
Compare
…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>
Fridah-nv
force-pushed
the
fridah/layerwise-fused-export
branch
from
August 11, 2026 00:17
aa52dbd to
5f88cf2
Compare
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>
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.
What does this PR do?
Type of change: new feature
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.
This adds
layerwise.export_dir: each decoder layer's final quantized tensors are flushed to their own shard the moment layerwise calibration finishes with it. When the last layer lands the directory is already a complete, loadable checkpoint, soexport_hf_checkpoint()is skipped. Combined withlayerwise.checkpoint_dirthe shards are the resume artifact — a restarted run skips layers already on disk instead of recalibrating and re-exporting them, and the per-layerweights.pt/quantizer_buffers.ptfiles are no longer written.Setting the config field is the entire switch; there is no CLI flag.
hf_ptq.pyrewrites its value to--export_path, exactly as it already rewriteslayerwise.checkpoint_dir.Two commits, fix first:
dc5a47cda2fix(hf_ptq):_is_layerwiseprobed the recipe's algorithm withgetattr, but an algorithm loaded from YAML is a plain dict — so it answeredFalsefor every layerwise recipe in the repo, and the auto batch-size probe it gates was never actually skipped. Pre-existing onmain; surfaced because the new detection shares the same accessor. Included because the feature depends on it, split out so it can be judged on its own. Behaviour change: with--batch_size 0, layerwise recipes now usebatch_size=1instead of probing, which is what the existing comment intends.aa52dbd346One layer per shard is what makes resume safe. 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 duplicate keys.
Resident models needed
transient_module_state. Export is destructive — it swapsweightfor packed bytes, registers scale buffers, grafts on per-expert submodules. An offloaded model discards that when its materialization window closes; a resident one has no such window, and calibration still has every later layer to run through it. So the layer subtree's_parameters/_buffers/_modulesare snapshotted and restored around each export.NVFP4 works by rediscovering fusion groups per layer. The groups
_fuse_shared_input_modulesoperates on — q/k/v behindinput_layernorm, gate/up behindpost_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 synthetictorch.ones([1, 2])the whole-model pass builds.What is refused, and why. These raise
NotImplementedErrorbefore calibration starts, because each would otherwise produce a silently different checkpoint rather than fail:requantize_resmooth_fused_llm_layers' pre-quant-scale steps, which are still whole-modelsync_tied_input_amaxmerges amaxes across a tie partner that may be uncalibrated or already writtenexport_hf_checkpoint()already streams these layer by layer--vllm_fakequant_export, non-dense sparsity--export_pathBecause a resumed run never recalibrates the layers it skipped, the in-memory model is not valid for inference afterwards; enabling the field implies
--skip_generate.Two pieces move out of this path rather than being duplicated:
save_non_weight_artifacts()is extracted from the streaming exporter's tail (both now share one copy), andFUSION_FREE_FORMATSmoves tomodel_config.py, where_fuse_shared_input_moduleshad held the same set inline as a literal list.Usage
python examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path <model> \ --export_path <out> \ --recipe modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export.yamlInterrupt it and rerun the same command: calibration resumes from the last committed layer and the finished shards are reused as-is.
Testing
Real model, byte-identical. Qwen3.6-35B-A3B (text-only view, FP8, 40 layers, 256 experts/layer, 21,746 quantizers), against a baseline of the same layerwise calibration followed by
export_hf_checkpoint():Key sets match exactly and
hf_quant_config.jsonagrees (FP8, 172 excluded modules). Fused run 174s / 41 shards / 34G; baseline 187s / 4 shards / 34G. Shard layout differs by design; contents do not. Re-verified after the tied-cache change, since that is the path the MoE dedup exercises.7 new GPU tests (
tests/gpu/torch/export/test_layerwise_export.py): equivalence vs whole-model export incl. on-disk artifacts; equivalence after a simulated mid-model resume; resume fails fast on mismatched dirs; per-layer resume files not written; KV-cache-quantized equivalence; NVFP4 equivalence; AWQ refusal.tests/gpu/torch/export/130 passed / 2 skipped (pre-existing env skips) · config + layerwise-calibration units 92 passed · pre-commit clean.Also verified directly: the model is bit-identical after export (logits and
state_dictunchanged, amax buffers intact), and the exported directory reloads throughAutoModelForCausalLMand runs a forward.Before your PR is "Ready for review"
export_dirdefaults toNone; every existing path is unchanged when unset. The one exception is thedc5a47cda2batch-size behaviour change, called out above.CONTRIBUTING.md: N/A — no new dependencies.Additional Information
A 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. Bisected:get_qdq_activations_from_prev_layer=Trueavoids it entirely, which pins the cause to the pre-calib_funccache_outputs_for_next_layer_calibpass — and explains why only the last layer, the one that skips it, is correct. It stayed hidden because both shipped NVFP4 layerwise recipes are experts-only and never quantizeo_proj. The NVFP4 test here excludeso_projfor the same reason.Two things for reviewers to weigh in on:
layerwise_calibrate(quantization) constructs aLayerwiseExporter(export), whilemodelopt/torch/export/__init__.pyalready imports quantization. A function-local import avoids the cycle, but the cleaner shape is formode.pyto inject the exporter as a callback so calibration never names the export layer. Left as follow-up; happy to do it here._export_transformers_checkpoint,_export_transformers_checkpoint_streaming,LayerwiseExporter).save_non_weight_artifacts()converges one piece; the per-tensor postprocess step (_collectvs_stream_tensor— same logic, different sink) is the next honest candidate. The sharding policies genuinely cannot converge: streaming packs by size and names shards at finalize, layerwise needs one shard per layer named by index, which is exactly what makes a crash resumable.