MTP support in per-layer fused export - #2259
Draft
Fridah-nv wants to merge 26 commits into
Draft
Conversation
_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>
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>
transient_module_state restored which tensor each name points at, not tensor contents, and its docstring claimed export only ever rebinds. Scale fusion disproves that: preprocess_linear_fusion unifies a group's amax through _amax.data.copy_(), an in-place write the dict restore cannot undo. Measured on NVFP4, five of the model's amax buffers came back carrying the fused group maxima. Buffers are now clone-restored; parameters are not, since export replaces weights by rebinding and cloning them would copy the whole layer. The earlier "model unchanged" check used FP8, which never fuses, so it could not see this. detect_resume_point returns None once the manifest is complete, so start_layer fell back to 0 and a rerun recalibrated and overwrote every finished shard -- whether the previous run crashed between the last ckpt.save and finalize, or simply finished. In export mode a complete manifest now means finalize-only: the shards are already on disk and only the tail, index and configs can be missing. The layer-order guard becomes a raise. As an assert, python -O strips it and layer N's tensors land in layer M's shard under a well-formed index -- silent corruption, which is the failure mode the rest of this file refuses. set_layerwise_export_dir now retargets only entries that already declare export_dir, rather than switching per-layer export on for any entry that happens to carry a layerwise block. Document the new recipe in modelopt_recipes/ptq.md: its absence was failing test_every_general_ptq_recipe_is_documented and the recipe-count check. Its own metadata also still said resident-only, which offload support had made stale. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…he checkpoint Four more review findings. Nothing tied the exported shards to the run that produced them: the resume manifest records no model or quantization identity, and assert_shards_present only checks that files exist. Two runs pointed at one export directory with the same layer count could pair one run's manifest with the other's shards and finalize a structurally valid, wrong checkpoint. The exporter now writes .layerwise_export.json (model class, layer count, formats, kv-cache format) on first export and refuses when an existing one disagrees, naming the fields that differ. The CLI's config-hash checkpoint subdirectory made this hard to hit there, but direct library use had nothing. set_layerwise_export_dir returned the config untouched when it found nothing to retarget. Since the caller decides to skip the real export from a separately parsed recipe object, a mismatch would send shards to the recipe's placeholder path and leave --export_path empty on a run that reports success. It raises now. The co-located resume directory was a child of --export_path, so its cached activations shipped inside the deliverable; nothing deletes it on success, and deleting it would make re-running a finished command trip assert_no_orphan_shards, which cannot tell "already done" from "resume record lost". It is a sibling now, and named .layerwise_resume: with per-layer export on the shards are the weights, so that directory holds no checkpoint at all -- only the resume point, the next layer's cached activations, and skip-mode shape metadata. The changelog advertised a --layerwise_export flag removed earlier in this PR in favour of the recipe field, so it pointed users at an argparse error. Rewritten, and cut to two sentences per CONTRIBUTING.md. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The shipped recipe was the only one in the repo to set layerwise.checkpoint_dir, and it set it to /tmp -- container-local, so a run that outlasted its GPU session came back to a wiped manifest, restarted at layer 0 and overwrote every finished shard. colocate_layerwise_checkpoint_dir existed to survive that by rewriting the path unconditionally, which also discarded any path the user had chosen deliberately. Removing the default removes the reason for the override. The recipe now ships no checkpoint_dir, matching every other layerwise recipe, and the helper becomes default_layerwise_resume_dir: it fills the field only when unset, so an explicit choice is honoured, and resume still works out of the box because leaving it unset would otherwise mean no resume at all -- the one thing this recipe is for. Three review findings alongside it. The int8 refusal tested "int8_sq" against --qformat, but that is the export format constant; the qformat preset is int8_smoothquant, and "int8_sq" is not a substring of it, so the check could never fire and an int8 run would have reached the TensorRT-LLM exporter and overwritten the shards. Matches both tokens now. The same dead test exists in export_quantized's own branch, left alone here. The export identity recorded only model class, layer count, format names and kv-cache format, so two FP8 runs differing in which modules were quantized shared one identity. It now carries a digest of the resolved quant config, which covers module selection and block settings, plus the source path. Source weights are not digested -- that would mean reading the whole model -- so differently-trained weights at the same path still compare equal; the code says so. The changelog entry was filed under Megatron Framework rather than Quantization. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
State why rather than what, and drop the restated-signature prose. No behavior change. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Export retargeting already walked a list-form algorithm, but the checkpoint-dir helpers only looked at a dict, so such a recipe calibrated with no resume manifest and lost its progress on interrupt. One list-aware accessor now backs all three helpers. Also correct the recipe's equivalence claim: layerwise export writes one shard per layer, so it matches whole-model export tensor for tensor, not byte for byte. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Two failures found in review, both on the per-layer export path. cache_outputs_for_next_layer_calib leaves the layer it just calibrated in 'run' mode with a drained deque. The qdq_from_prev=False branch already reset it; the True branch did not, because until this PR nothing ran the layer afterwards. The fusion probe does, so every fusing format asserted on layer 0. Reset it there too, and cover the combination -- it is the one layerwise axis the tests missed. The run identity is written before layer 0, so a run that died early left a directory holding only .layerwise_export.json, and re-running into it after a config change was refused as if shards were at stake. Bind it only once shards exist, and rewrite a stale one instead of keeping it forever. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
A sequential algorithm list made export ownership ambiguous. Detection read the first layerwise block, so an exporting block behind a non-exporting one missed the opt-in and exported to the recipe's placeholder; deriving the resume dir had the mirror bug, letting another pass's explicit checkpoint_dir stand in for the exporting one's, and giving every block one shared directory when none was set. Export finalizes shards as calibration walks the layers, so only the final pass can own them: layerwise_export_block() enforces exactly one export_dir on the last entry and every export-driven helper reads through it. resolve_checkpoint_dir now resolves each block against its own base, so two passes cannot collide on one manifest. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
AWQ and SVDQuant are only distinguishable after calibration -- their discriminators are registered by the calibrator -- so the constructor's gate saw plain NVFP4 and passed. An NVFP4-AWQ recipe then exported unfused pre_quant_scale plus the layernorm the whole-model path folds it into: 24 extra keys and 24 mismatched tensors against export_hf_checkpoint on a tiny Llama. Re-checking on the first exported layer catches any format that only becomes visible after calibration, not just AWQ by name. Also state what per-layer export leaves behind. The tail is converted in place, because its tensors are read from model.state_dict() after dispatch and restoring first would collect unpacked weights -- so the model is unusable for inference on every run, not only a resumed one. Ungate that warning and correct the docstrings that promised otherwise, including transient_module_state, which does not restore attributes outside _parameters/_buffers/_modules. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Five tests built a whole-model baseline and compared it to a per-layer export, differing only in the quant config; the two AWQ refusals and the two run-identity refusals were paired the same way. Parameterize each family so the axis under test reads as a table, and name the config variants once instead of inlining them. Same 17 cases collected. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…view gaps The auto-derived resume dir was never opted into, and kept one full activation set per layer -- on a 40-layer model that dwarfs the checkpoint it sits beside. Prune every boundary but the committed one, after the manifest commits it so a crash mid-write still resumes. Scoped to the per-layer export path, leaving the explicit checkpoint_dir semantics untouched. Fuse the sibling experts the probe never routed to, mirroring the replay requantize_resmooth_fused_llm_layers does for the same reason: sync_moe_gate_up_amax covers weight_quantizer.amax but not a static quantizer's global_amax, so unrouted gate/up pairs would keep unmerged scales. Drop the format re-check latch -- it inspected only one layer, so a recipe applying a pre-quant-scale format to a subset still shipped unfused pre_quant_scale. Refuse MTP checkpoints before calibration on the normal load path: only the FSDP2 loader flags the prefixes early, so the run used to write a complete-looking checkpoint and fail afterwards. Also: no MoE gate/up warning at finalize (the sync happens inside transient_module_state, so it fired on every MoE run and reported a miss that did not happen); delete out-of-range shards from a longer previous run; null-safe manifest read plus a num_layers check; drop __all__, which advertised names the package never exported. The resume test now interrupts for real rather than rewinding a finished run's manifest -- a state a crash cannot produce. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
… config The format gate lived in two places with a duplicated message, because AWQ and SVDQuant only become visible once the calibrator registers their discriminators. assert_formats_supported() is now that check, called before calibration to fail early and on each exported layer as the authority; assert_layerwise_export_supported keeps the structural cases. Record the parity with requantize_resmooth_fused_llm_layers as a table next to _fuse_shared_inputs: every step of that pass is intra-layer, so each maps to a per-layer equivalent or an explicit refusal, and a gap is visible rather than discovered. Resume is an axis of the equivalence oracle, not one example: every config now runs fresh and after a real interruption, since a resumed run re-enters export part-way through the model. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The reset exists so the fusion probe hits a real forward. Without an exporter nothing runs the layer before _set_layer_states does, and layerwise calibration with get_qdq_activations_from_prev_layer=True has no test coverage, so leave that path exactly as it was. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…g it Tensor equality against a whole-model export never resolves the index, so a weight_map entry naming the wrong shard compares equal and still fails to load. That mapping is the one artifact per-layer export builds differently. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
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>
The refusal said MTP exclusions and orphaned weights are applied after calibration, when every shard is already written. The first half stopped being true once the prefixes were derived from the checkpoint index before calibration: the pre-quantize exclusion loop already leaves MTP modules unquantized, and finalize() already calls _add_mtp_exclusions. That leaves the orphans -- MTP tensors with no slot in state_dict(), which the separate-file conventions produce. load_mtp_weights only fills existing slots and returns the rest, so it is safe to run before quantize; per-layer export now does that and stashes the leftovers on the model under MTP_EXTRA_STATE_ATTR, which finalize() feeds into its existing extra_state_dict path. The stash exists because calibration owns the finalize() call, so hf_ptq cannot pass them as an argument -- the same reason _mtp_layer_prefixes is already carried that way. The blanket refusal is replaced by a narrow one: if the post-calibration load finds tensors that were not staged, the run still fails, because the shards are written by then and they cannot be added. Prototype: covered by an inlined-convention test that fails without the stash pickup. The separate-file conventions (GLM-4.7 standalone mtp.safetensors, Qwen3-Next indexed tail shard) have no local fixture and are unverified. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
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 Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
Fridah-nv
force-pushed
the
fridah/layerwise-fused-export
branch
from
August 30, 2026 04:56
0cb804b to
5b32946
Compare
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
Stacked on #2136, which refuses MTP models outright. This lifts that refusal.
Why the refusal existed, and what changed. It said MTP exclusions and orphaned MTP
weights are applied after calibration, by which point per-layer export has already written
every shard and the quant config. Two thirds of that stopped being true:
hf_ptq.py, which appends{"quantizer_name": "*<prefix>*", "enable": False}toquant_cfg. feat(export): export each decoder layer as layerwise calibration finishes it #2136 already derivesthose prefixes from the checkpoint index before calibration, so MTP modules are simply
never quantized.
finalize()calls_add_mtp_exclusions.model.layers.{N}, GLM-5.1 / DeepSeek-V3) are returned byget_homogeneous_hf_decoder_layerslike any other decoder layer, so they already gettheir own shard — unquantized, because they are excluded.
That leaves orphans: MTP tensors with no slot in
model.state_dict(), which theseparate-file conventions produce.
load_mtp_weights()only fills existing slots and handsthe rest back, so it is safe to run before
mtq.quantize. Per-layer export now does thatand stashes the leftovers on the model under
MTP_EXTRA_STATE_ATTR;finalize()feeds theminto the
extra_state_dictpath it already had.The stash exists because calibration owns the
finalize()call, sohf_ptqcannot passthem as an argument.
_mtp_layer_prefixesis already carried across that same boundary thesame way, so this follows an existing convention rather than inventing one.
The blanket refusal becomes a narrow one: if the post-calibration load finds keys that were
not staged, the run still fails, because the shards are written by then and nothing can be
added to them.
Usage
No new flags. An MTP checkpoint with
layerwise.export_dirset now exports instead ofraising
NotImplementedError.Testing
tests/gpu/torch/export/test_layerwise_export.py— 25 passed. One new test pins thatstashed orphans reach the tail shard and the index; it fails with
mtp.layers.0.weight missing from the exported checkpointwhen the stash pickup is stubbedout, so it is not vacuous.
tests/unit/recipe282 ·tests/unit/torch/export172 ·tests/examples/hf_ptq36 ·pre-commit clean.
Draft, because the coverage is narrower than the feature. The test exercises orphan
delivery, which is the mechanism this PR changes. It does not exercise the four conventions
load_mtp_weightssupports — inlined (GLM-5.1, DeepSeek-V3), standalonemtp.safetensors(GLM-4.7), indexed
mtp.*tail shard (Qwen3-Next) — none of which has a local fixture.Two risks I have not been able to close:
load_mtp_weightsbefore quantization may split in-slot vs orphan differentlythan running it after, since
model.state_dict()changes once quantizers are inserted.The late guard compares key sets precisely because of this, but a real MTP model should
confirm the split is what we expect.
that path is reasoned about, not observed.
A full export on a real MTP checkpoint is the missing step before this leaves draft.
Before your PR is "Ready for review"
layerwise.export_dirset,which is opt-in. Non-layerwise MTP export is untouched.
guidance in
CONTRIBUTING.md: N/Aare unverified (see Testing).
layerwise.export_direntry and listsMTP as refused; it needs amending once merge order is settled.
Additional Information
Depends on #2136 and must not merge before it. #2136's refusal test and its real-checkpoint
verification both assert that MTP is refused, so whichever lands second needs them
reconciled.