MTP support in per-layer fused export - #2259
Closed
Fridah-nv wants to merge 1 commit into
Closed
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 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
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>
Fridah-nv
force-pushed
the
fridah/layerwise-mtp-support
branch
from
August 30, 2026 22:56
53b3865 to
c69ff4f
Compare
Contributor
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2259 +/- ##
==========================================
- Coverage 78.64% 78.64% -0.01%
==========================================
Files 525 525
Lines 61104 61108 +4
==========================================
Hits 48057 48057
- Misses 13047 13051 +4
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
added a commit
that referenced
this pull request
Sep 8, 2026
…2303) ### What does this PR do? Type of change: New feature **Layerwise export now supports multimodal and MTP models.** Both were refused outright, and both were refused for the same reason: `finalize()` was called from inside `layerwise_calibrate`, which is the wrong scope for it. **1. Calibration does not know which model the checkpoint describes.** It only sees the module it was handed. A VLM calibrates its *language model*, so the shards, the exclusions and `config.json` all came out describing that submodel rather than the whole VLM. Moving the call out lets the caller root the exporter at the parent — and without the key prefixing, tower collection or ambient parent handle an earlier attempt needed, because the decoder layers are the same objects from either root. **2. Calibration runs before things the export needs exist.** Orphaned MTP weights are loaded *after* calibration, by which point every shard had already been written, so they could not be passed at all. After the move they are an ordinary argument to `finalize()`, with no staging attribute stashed on the model. ### How it works The exporter is created by whoever owns the export and **announced on the model** that `mtq.quantize` is given. Calibration picks it up, binds it, and drives it per layer; the export that follows reads it back and finishes the checkpoint: ```python LayerwiseExporter(full_model, export_path).announce(language_model) mtq.quantize(language_model, quant_cfg, forward_loop=loop) ... getattr(full_model, LAYERWISE_EXPORTER_ATTR).finalize(extra_state_dict=mtp_state_dict) ``` Calibration and export are handed *different* models, so `announce()` publishes the exporter on each end separately: the caller announces on the model being calibrated, and `bind()` announces on the export root. Neither side has to know where the other looked, and the lookup stays an O(1) `getattr` rather than a `named_modules()` scan — worth avoiding at roughly 1.65 µs/module, or ~500 ms on a Kimi-K3-sized model. For a non-VLM both roots are the same object and the second announcement is a no-op. `finalize()` clears every attachment it recorded, so the module graph does not retain a live exporter afterwards. `mtq.quantize` and `mtq.calibrate` are **unchanged** — a layerwise-only feature does not belong in the public quantization API. The attribute follows `_mtp_layer_prefixes`, which crosses the same calibration→export boundary the same way (`hf_ptq.py:538` sets it, `unified_export_hf.py:870` reads it back). Construction is inert: `__init__` records only the export root and the directory, because the caller builds it before `mtq.quantize`, when there are no quantizers yet to validate or read a config from. `bind()` does that, called from calibration after quantizer insertion and before any layer is converted — the only window where both hold, and the same instant the exporter used to be constructed, so unsupported models still fail in seconds rather than hours. Only the calibration pass that sets `export_dir` drives the exporter: a list-form algorithm runs one pass per entry, and an earlier one must not convert layers a later one still has to calibrate. ### Usage Nothing changes for a plain layerwise-export recipe: `layerwise.export_dir` still drives it. Pre-attaching an exporter is the opt-in for the two cases that need it — a checkpoint whose root is wider than the calibrated model, and orphaned tensors to merge at the end. The one behaviour change for a config-only caller is that `mtq.quantize` now writes the layer shards but no longer finishes the checkpoint. Both exit paths warn with what is still owed, and `LayerwiseConfig.export_dir`'s description has been corrected — it previously promised "a complete, loadable checkpoint when the last layer lands" and still listed multimodal and MTP as raising `NotImplementedError`. ### Testing `tests/gpu/torch/export/test_layerwise_export.py` — **29 passed**. Beyond the 24 inherited from #2136, five new ones, each with a negative control confirming it fails without its fix: - orphaned MTP tensors reach the tail shard *and* the index - an exporter rooted at the parent widens the checkpoint's namespace - the config-only path announces an exporter that can be finished, and finalize clears it - only the pass that sets `export_dir` drives the exporter - an exporter whose root holds a different number of layers is refused at `bind()` Full suites: `tests/gpu/torch/export` + `tests/gpu/torch/quantization` **1012 passed / 55 skipped**, `tests/unit` **3318 passed / 15 skipped**, pre-commit clean. Both suites also report failures in `test_implicit_gemm.py` (FP4 conv kernels), `test_triton_fa_p_qdq.py`, `test_autocast_quantize_int8` and `test_engine_builder.py` collection; all reproduce unchanged on `main` and none touch the paths in this diff. Measured against the whole-model exporter on a tiny Gemma3-VL, towers prepared exactly as `hf_ptq` does: ``` keys: baseline=80 layerwise=80 only-baseline=[] only-layerwise=[] differing values: 0 vision tower present: True VLM namespace: True config.json is the VLM: True hf_quant_config match: True exclude_modules: ['language_model.lm_head', 'vision_tower.vision_model*'] (both sides) ``` #### End-to-end through `hf_ptq.py` Same FP8 recipe both sides; the baseline drops `layerwise.export_dir` and is exported by `main`, so the diff isolates this PR. Every tensor matches in key, dtype, shape and value, and `config.json` / `hf_quant_config.json` match too. | Model | Covers | Keys | Differing | |---|---|---|---| | Qwen3-VL-8B-Instruct | multimodal | 1254 = 1254 | 0 | | GLM-4.7-Flash | MoE + MTP | 28119 = 28119 | 0 | The VLM checkpoint keeps the vision tower unquantized (351 `model.visual.*` keys, no `weight_scale` among them) while the language model is FP8. The MTP run reports 212 orphaned tensors; all 212 land in `model-tail.safetensors` and in the index, with `model.layers.47*` in `exclude_modules`. **Not yet validated:** an accelerate-offloaded run, and a serving canary on the exported checkpoints. ### Refusals `export_dir` without `enable`, and an exporting algorithm entry with no calibration method, are both refused before calibration starts — neither reaches the per-layer pass, so both would otherwise export nothing. The early gate is a heuristic on the recipe, so `hf_ptq` also raises a plain `RuntimeError` at export time if calibration turned out not to have run; that backstop, not the gate, is what makes the failure legible on paths the recipe check cannot predict. `bind()` requires the layers calibration will drive and refuses a root that discovers a different number of them. Only the count is checked here: `export_layer` already rejects a reordering or a substituted module on its first call, and a length difference is the one mismatch it structurally cannot catch — every call would pass and `_write_index` would then open a shard that was never written, at the very end of the run. Orphan tensors are merged into the tail with no collision check, matching the whole-model path (`unified_export_hf.py:1623`). `load_mtp_weights` returns exactly the keys absent from `model.state_dict()`, so a collision with an exported tensor is not reachable through the only producer, and a guard would only make the two export paths diverge. ### Why not reuse `export_hf_checkpoint` It was the first idea and it is the most expensive one. Its transformers path is whole-model at every step — `_prepare_moe_inputs`, `requantize_resmooth_fused_llm_layers` (which runs a dummy forward that would fail on already-converted layers), `_process_quantized_modules`, a full `model.state_dict()` in host RAM, then `save_pretrained` rewriting shards already on disk — and it raises outright under `has_accelerate_offload`. `save_pretrained(state_dict={})` is not an escape either: safetensors' shared-storage check fires on MoE even with an empty dict. The natural consolidation target is the **streaming** exporter, which is already most of `finalize()`: 122 lines vs 74, sharing `decoder_owned_ids`, `enable_weight_access_and_writeback`, `_dispatch_export_handler`, `_reconstruct_fused_moe_linear`, `_add_mtp_exclusions`, `_postprocess_single_tensor`, `requires_weight_materialization` and `save_non_weight_artifacts`. Folding them together needs roughly four knobs: skip the whole-model prep, skip layers already written, seed the index with the existing shards, and inject the quant config. That is a separate change and deliberately not in this one. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — `mtq.quantize`/`mtq.calibrate` signatures are unchanged, and a recipe that only sets `layerwise.export_dir` behaves as before. The one behaviour change is that `mtq.quantize` no longer finishes the checkpoint on its own: callers must now call `finalize()` on the exporter, which calibration leaves on the model. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ❌ — pending. - Did you get Claude approval on this PR?: ❌ — the last review's findings are all addressed; needs a re-run. ### Additional Information Follow-ups this enables: #2259 (MTP) reduces to close to nothing, and the multimodal work in #2218 no longer needs `export_parent`, the key prefixing, or the tower collection. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Layerwise export now supports clearer control over export locations and calibrated layer handling. - Export workflows provide improved support for resuming, sharded checkpoints, mixture-of-experts models, and nested model namespaces. - **Bug Fixes** - Improved handling of exported checkpoint shards and extra tensors. - Added clearer warnings when exports require completion before loading. - **Documentation** - Clarified that layerwise exports write shards during calibration and require an explicit finalization step. - Documented that the in-memory model is not suitable for inference after layerwise export. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.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
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.