MTP support in per-layer fused export

[Fridah-nv]

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:

• Exclusions are applied by the pre-quantize loop in hf_ptq.py, which appends
  {"quantizer_name": "*<prefix>*", "enable": False} to quant_cfg. feat(export): export each decoder layer as layerwise calibration finishes it #2136 already derives
  those prefixes from the checkpoint index before calibration, so MTP modules are simply
  never quantized.
• The exported quant config already lists them: finalize() calls _add_mtp_exclusions.
• Inlined MTP layers (model.layers.{N}, GLM-5.1 / DeepSeek-V3) are returned by
  get_homogeneous_hf_decoder_layers like any other decoder layer, so they already get
  their own shard — unquantized, because they are excluded.

That leaves orphans: MTP tensors with no slot in model.state_dict(), which the
separate-file conventions produce. load_mtp_weights() only fills existing slots and hands
the rest back, so it is safe to run before mtq.quantize. Per-layer export now does that
and stashes the leftovers on the model under MTP_EXTRA_STATE_ATTR; finalize() feeds them
into the extra_state_dict path it already had.

The stash exists because calibration owns the finalize() call, so hf_ptq cannot pass
them as an argument. _mtp_layer_prefixes is already carried across that same boundary the
same 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_dir set now exports instead of
raising NotImplementedError.

Testing

tests/gpu/torch/export/test_layerwise_export.py — 25 passed. One new test pins that
stashed orphans reach the tail shard and the index; it fails with
mtp.layers.0.weight missing from the exported checkpoint when the stash pickup is stubbed
out, so it is not vacuous.

tests/unit/recipe 282 · tests/unit/torch/export 172 · tests/examples/hf_ptq 36 ·
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_weights supports — inlined (GLM-5.1, DeepSeek-V3), standalone mtp.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:

1. Running load_mtp_weights before quantization may split in-slot vs orphan differently
   than 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.
2. An inlined MTP layer surviving as a decoder layer should produce a correctly named shard;
   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"

• Is this change backward compatible?: ✅ — only reachable with layerwise.export_dir set,
  which is opt-in. Non-layerwise MTP export is untouched.
• 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?: ⚠️ — one, covering the mechanism; the conventions
  are unverified (see Testing).
• Did you update Changelog?: ❌ — feat(export): export each decoder layer as layerwise calibration finishes it #2136 carries the layerwise.export_dir entry and lists
  MTP as refused; it needs amending once merge order is settled.
• Did you get Claude approval on this PR?: ❌ — not yet run.

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.

[copy-pr-bot]

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.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}