From fda0fde8d8aa27fbd2c8d209928088eb735a3762 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Mon, 31 Aug 2026 20:26:41 +0000
Subject: [PATCH 1/5] feat(export): Kimi-K3 on per-layer fused export, VLM
 included

Brings per-layer fused export up on moonshotai/Kimi-K3 -- 1.5 TB, 896 experts,
93 layers, a VLM -- quantized to NVFP4 on a single B200. Four pieces.

Decoder discovery through arbitrary nesting. The walk unwrapped `.model` then
`.language_model` once each, so it only found layers exactly two wrappers deep in
that order. K3 keeps its decoder at language_model.model.layers, so the walk
stopped on the intermediate wrapper and reported the architecture unsupported. It
now descends through the wrappers and reads `layers` at the bottom, which
generalises the fixed two-level unwrap rather than bounding a shallow-first
search. Not K3-specific.

Multimodal support, replacing the refusal. The refusal was right about the
failure: calibration runs on the extracted language model, so the shards and
config.json described that submodel rather than the whole VLM. The caller now
marks the submodel through an export_parent() context manager over a ContextVar,
and the exporter resolves the parent by identity and walks from there. The decoder
layers are the same objects either way, so parent-namespace tensor names, the full
VLM config and the unquantized towers all fall out of the passes that already
exist -- no key prefixing, no separate tower collection. The refusal narrows
rather than disappears: a VLM whose language model is not reachable from the full
model is still refused, since the parent would be undefined.

Hub names on transformers 4. The whole-model path writes through save_pretrained,
which is what reverses _checkpoint_conversion_mapping. Per-layer export writes
shards directly with save_file and never passes through it, so it emitted
in-memory names where the whole-model path emits published ones. Also not
K3-specific: any transformers-4 model with that mapping was affected.

Recipe. Experts-only NVFP4 + FP8 KV for offloaded models, scoped `*.experts.*`
rather than `*block_sparse_moe*` -- the broad glob also matches `shared_experts.*`
and `routed_expert_*_proj`, 552 modules the vendor left unquantized -- plus its
ptq.md row. checkpoint_dir is left unset so it derives
<export_path>.layerwise_resume beside the shards, instead of a container-local
/tmp path that a run outlasting its GPU session comes back to find wiped.

_force_attn_implementation is reduced to the model config and its direct
sub-configs, since K3's remote code only rewrites text_config. Tied-weight alias
handling is left out: K3 sets tie_word_embeddings=False and tied weights are
unsupported upstream, so it belongs in its own change.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 examples/hf_ptq/example_utils.py              | 21 +++++
 examples/hf_ptq/hf_ptq.py                     | 44 ++++++----
 modelopt/torch/export/layerwise_export.py     | 81 ++++++++++++++++++-
 .../torch/quantization/plugins/huggingface.py | 19 +++--
 ..._only-kv_fp8_layerwise_export_offload.yaml | 63 +++++++++++++++
 modelopt_recipes/ptq.md                       |  3 +-
 .../gpu/torch/export/test_layerwise_export.py | 64 ++++++++++++++-
 7 files changed, 265 insertions(+), 30 deletions(-)
 create mode 100644 modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index 862bc0f7fe5..6035fbba603 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -41,6 +41,7 @@
     AutoModelForCausalLM,
     AutoProcessor,
     AutoTokenizer,
+    PretrainedConfig,
     PreTrainedTokenizerBase,
     ProcessorMixin,
 )
@@ -703,6 +704,20 @@ def _resolve_init_config(hf_config, auto_model_module, ckpt_path, config_kwargs)
         return hf_config
 
 
+def _force_attn_implementation(model, attn_implementation: str) -> None:
+    """Re-apply the caller's ``attn_implementation`` after ``__init__``.
+
+    Kimi-K3's remote code rewrites ``text_config._attn_implementation`` to
+    ``flash_attention_2`` unconditionally, so an uninstalled backend fails later at the
+    export trace forward instead of at load. Drop this once flash-attn is installed
+    reliably here and the checkpoint is fixed upstream.
+    """
+    sub_configs = (v for v in vars(model.config).values() if isinstance(v, PretrainedConfig))
+    for cfg in (model.config, *sub_configs):
+        if getattr(cfg, "_attn_implementation", None) not in (None, attn_implementation):
+            cfg._attn_implementation = attn_implementation
+
+
 def _get_config_dtype(config):
     config_dtype = (
         getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16
@@ -967,6 +982,12 @@ def has_pack_quantized_config(config):
             **model_kwargs2,
         )
     model.eval()
+
+    # Honour the caller's explicit choice even when remote modeling code overwrote it
+    # during __init__ (see _force_attn_implementation).
+    if attn_implementation is not None:
+        _force_attn_implementation(model, attn_implementation)
+
     if has_pack_quantized_config(hf_config):
         _unpack_compressed_linear_weights(model, ckpt_path)
 
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 0e3c8035bf8..590ebee5193 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -78,6 +78,7 @@
     has_spec_opt,
     save_expert_token_count_table,
 )
+from modelopt.torch.export.layerwise_export import export_parent
 from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model
 from modelopt.torch.quantization.config import need_calibration
 from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights
@@ -758,7 +759,12 @@ def mono_quantize(
                 language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop
             )
         else:
-            language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop)
+            # A VLM calibrates its language model but must export the whole thing.
+            parent = full_model if args.layerwise_export else language_model
+            with export_parent(parent):
+                language_model = mtq.quantize(
+                    language_model, quant_cfg, forward_loop=calibrate_loop
+                )
 
         # For VL models, update full_model to use the quantized language model
         if is_nemotron_vl_model:
@@ -779,12 +785,16 @@ def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) ->
     calibration begins, the user has already paid for the whole run.
     """
     if is_multimodal_model(full_model):
-        raise NotImplementedError(
-            "layerwise.export_dir does not support multimodal models: calibration runs on the "
-            "extracted language model, so the shards and config.json would describe that "
-            "submodel rather than the full VLM, and the VLM export path would then "
-            "overwrite config.json with the unquantized source config."
-        )
+        # Calibration runs on the extracted language model, so the exporter is pointed at the
+        # full model below. That needs the submodel to be reachable from it by identity.
+        lineage = get_language_model_from_vl(full_model)
+        language_model = lineage[-1] if lineage else None
+        if language_model is None or all(m is not language_model for m in full_model.modules()):
+            raise NotImplementedError(
+                "layerwise.export_dir does not support this multimodal model: its language "
+                "model could not be located inside the full model, so the exported tensor "
+                "names cannot be resolved against the VLM."
+            )
 
     if mtp_layer_prefixes:
         raise NotImplementedError(
@@ -859,15 +869,17 @@ def export_quantized(
         is_vlm = is_multimodal_model(full_model)
 
         if is_vlm:
-            # Save original model config and the processor config to the export path for VLMs.
-            print(f"Saving original model config to {export_path}")
-
-            config_kwargs = {"trust_remote_code": args.trust_remote_code}
-            if args.attn_implementation is not None:
-                config_kwargs["attn_implementation"] = args.attn_implementation
-            AutoConfig.from_pretrained(args.pyt_ckpt_path, **config_kwargs).save_pretrained(
-                export_path
-            )
+            # Skipped under per-layer export: calibration already wrote a config carrying the
+            # quantization_config, and the source config is unquantized.
+            if not args.layerwise_export:
+                print(f"Saving original model config to {export_path}")
+
+                config_kwargs = {"trust_remote_code": args.trust_remote_code}
+                if args.attn_implementation is not None:
+                    config_kwargs["attn_implementation"] = args.attn_implementation
+                AutoConfig.from_pretrained(args.pyt_ckpt_path, **config_kwargs).save_pretrained(
+                    export_path
+                )
 
             # Try to save processor config if available
             try:
diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py
index a196c9a9e52..ea7306722a0 100644
--- a/modelopt/torch/export/layerwise_export.py
+++ b/modelopt/torch/export/layerwise_export.py
@@ -16,6 +16,7 @@
 """Write each decoder layer's quantized checkpoint shard as soon as it is calibrated."""
 
 import contextlib
+import contextvars
 import json
 import warnings
 from pathlib import Path
@@ -146,6 +147,73 @@ def assert_layerwise_export_supported(model: nn.Module) -> None:
         )
 
 
+_export_parent: contextvars.ContextVar[nn.Module | None] = contextvars.ContextVar(
+    "layerwise_export_parent", default=None
+)
+
+
+@contextlib.contextmanager
+def export_parent(parent: nn.Module):
+    """Export the checkpoint for ``parent`` while calibration runs on one of its submodules.
+
+    Multimodal pipelines calibrate the extracted language model, but the shards and config
+    have to describe the whole VLM. The decoder layers are the same objects either way, so
+    walking the parent yields parent-namespace tensor names, the full config and the
+    untouched towers with no prefixing.
+    """
+    token = _export_parent.set(parent)
+    try:
+        yield
+    finally:
+        _export_parent.reset(token)
+
+
+def _resolve_export_parent(model: nn.Module) -> nn.Module:
+    """Return the model the checkpoint should describe. Membership is by identity, not name."""
+    parent = _export_parent.get()
+    if parent is None or parent is model:
+        return model
+    if all(m is not model for m in parent.modules()):
+        raise ValueError(
+            f"export_parent() was given a {type(parent).__name__} that does not contain the "
+            "calibrated model."
+        )
+    return parent
+
+
+def build_legacy_name_mapper(model: nn.Module):
+    """Hub-name mapper for transformers < 5, or ``None``.
+
+    The whole-model path writes through ``save_pretrained``, and it is save_pretrained that
+    reverses ``_checkpoint_conversion_mapping`` (in-memory ``model.language_model.*`` ->
+    published ``language_model.model.*``). Per-layer export writes shards directly, so it
+    never passes through that. ``build_reverse_name_mapper`` does not cover it either: it
+    reads transformers 5's ``conversion_mapping`` module and raises on 4.x.
+
+    The mapping is stored hub-pattern -> in-memory-prefix, so it is inverted here. Longest
+    in-memory prefix first, or a short rule shadows a longer one (``lm_head`` inside
+    ``model.language_model...``).
+    """
+    import re
+
+    mapping = getattr(model, "_checkpoint_conversion_mapping", None)
+    if not mapping:
+        return None
+    rules = sorted(
+        ((re.compile("^" + re.escape(mem)), hub.lstrip("^")) for hub, mem in mapping.items()),
+        key=lambda r: -len(r[0].pattern),
+    )
+
+    def _map(name: str) -> str:
+        for pattern, replacement in rules:
+            new, n = pattern.subn(replacement, name, count=1)
+            if n:
+                return new
+        return name
+
+    return _map
+
+
 class LayerwiseExporter:
     """Writes one decoder layer's quantized shard per call, then the tail and index.
 
@@ -163,6 +231,7 @@ def __init__(
 
         Runs before calibration, so nothing amax-dependent exists yet.
         """
+        model = _resolve_export_parent(model)
         assert_layerwise_export_supported(model)
         # Splits regroup tensors across the whole state dict; no per-layer pass reverses that.
         _assert_no_split_rules(model)
@@ -208,10 +277,14 @@ def __init__(
         try:
             self._name_mapper = build_reverse_name_mapper(model)
         except Exception as exc:
-            warnings.warn(
-                f"Reverse name mapper unavailable ({exc}); exported tensor names may not "
-                "match the original HF hub checkpoint."
-            )
+            # transformers < 5 has no conversion_mapping module; fall back to the legacy
+            # mapping save_pretrained would have applied.
+            self._name_mapper = build_legacy_name_mapper(model)
+            if self._name_mapper is None:
+                warnings.warn(
+                    f"Reverse name mapper unavailable ({exc}); exported tensor names may not "
+                    "match the original HF hub checkpoint."
+                )
 
     def export_layer(
         self,
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 4acb4d30dfa..b06e70e6f5d 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1800,13 +1800,18 @@ def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None:
     if not _is_supported_hf_model(model):
         return None
 
-    decoder = model
-    if hasattr(decoder, "model"):
-        decoder = decoder.model
-    if hasattr(decoder, "language_model"):
-        decoder = decoder.language_model
-    if hasattr(decoder, "layers"):
-        return decoder.layers
+    # Descend through every wrapper before reading ``layers``: multimodal models nest them
+    # in either order, e.g. Kimi-K3 keeps its layers at ``language_model.model.layers``.
+    decoder, seen = model, set()
+    while id(decoder) not in seen:
+        seen.add(id(decoder))
+        for attr in ("model", "language_model"):
+            inner = getattr(decoder, attr, None)
+            if isinstance(inner, nn.Module):
+                decoder = inner
+                break
+        else:
+            return getattr(decoder, "layers", None)
 
     return None
 
diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
new file mode 100644
index 00000000000..b041123cfa7
--- /dev/null
+++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  nvfp4: configs/numerics/nvfp4
+  kv_fp8: configs/ptq/units/kv_fp8
+
+metadata:
+  recipe_type: ptq
+  description: >
+    NVFP4 W4A4 on routed experts only, FP8 KV cache, max layerwise calibration, with each
+    decoder layer exported to its own shard as soon as it is calibrated. Same intent as
+    nvfp4_experts_only-kv_fp8_layerwise_export, but scoped and paired for a model too large
+    to hold resident: use with --offload_folder and per-device memory budgets.
+
+    Expert scoping is '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE
+    models the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and
+    routed_expert_norm -- on Kimi-K3 that is 552 extra modules the vendor deliberately left
+    unquantized, one of which is an RMSNorm. '*.experts.*' matches only the routed expert
+    projections; shared_experts is missed because its path contains '_experts.' rather than
+    '.experts.' (fnmatch semantics, conversion.py).
+
+    An interrupted run resumes without recalibrating or re-exporting finished layers, losing
+    at most the in-flight one -- the point of the combination for a run that outlasts its GPU
+    session. The resume state lives beside the checkpoint at <export_path>.layerwise_resume
+    unless you set layerwise.checkpoint_dir yourself; do not point it at container-local
+    storage, or a run that survives its session comes back to a wiped manifest.
+
+    A resumed run never recalibrates the layers it skipped, so the exported checkpoint is
+    complete but the in-memory model is not and must not be used for inference.
+quantize:
+  algorithm:
+    method: max
+    layerwise:
+      enable: true
+      # max only updates _amax, so the exported shard stays valid for its layer.
+      calib_mutates_weights: false
+      # Presence enables per-layer export; the value is replaced with --export_path.
+      export_dir: /tmp/modelopt_layerwise_export
+  quant_cfg:
+    - $import: base_disable_all
+    - quantizer_name: '*.experts.*weight_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*.experts.*input_quantizer'
+      cfg:
+        $import: nvfp4
+    - $import: kv_fp8
+    - $import: default_disabled_quantizers
diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md
index 1735289ac3f..3f85166158e 100644
--- a/modelopt_recipes/ptq.md
+++ b/modelopt_recipes/ptq.md
@@ -29,7 +29,7 @@ supported combinations.
 ### The shipped recipes
 
 <details>
-<summary>All 25 <code>general/ptq/</code> recipes (click to expand)</summary>
+<summary>All 26 <code>general/ptq/</code> recipes (click to expand)</summary>
 
 | Recipe | Model body | KV cache | Calibration |
 |--------|-----------|----------|-------------|
@@ -49,6 +49,7 @@ supported combinations.
 | `nvfp4_experts_only-kv_fp8_layerwise` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise |
 | `nvfp4_experts_only-kv_fp8_layerwise_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (non-mutating, for disk offload) |
 | `nvfp4_experts_only-kv_fp8_layerwise_export` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (exports each layer as it is calibrated) |
+| `nvfp4_experts_only-kv_fp8_layerwise_export_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise, tuned for accelerate-offloaded models |
 | `nvfp4_experts_only_mse-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | MSE + FP8 sweep |
 | `nvfp4_experts_only_input_scale1-kv_fp8_cast` | NVFP4 W4A4, MoE experts only, expert `input_scale` pinned to 1.0 | FP8 (constant amax) | max (weights); expert activations uncalibrated |
 | `nvfp4_omlp_only-kv_fp8` | NVFP4 W4A4, o_proj + MLP/MoE | FP8 (calibrated) | max |
diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py
index ca40f72e875..4634c8af420 100644
--- a/tests/gpu/torch/export/test_layerwise_export.py
+++ b/tests/gpu/torch/export/test_layerwise_export.py
@@ -22,11 +22,20 @@
 
 import pytest
 import torch
-from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe
+from _test_utils.torch.transformers_models import (
+    get_tiny_gemma3vl,
+    get_tiny_llama,
+    get_tiny_qwen3_moe,
+)
 from safetensors.torch import load_file
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.layerwise_export import LayerwiseExporter, layer_shard_name
+from modelopt.torch.export.layerwise_export import (
+    LayerwiseExporter,
+    export_parent,
+    layer_shard_name,
+)
+from modelopt.torch.export.model_utils import get_language_model_from_vl
 from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
 
 NUM_LAYERS = 4
@@ -477,3 +486,54 @@ def test_awq_is_refused(tmp_path, make_cfg, method, match):
 
     with pytest.raises(NotImplementedError, match=match):
         mtq.quantize(_build_model(), cfg, _calib)
+
+
+def _build_vlm():
+    torch.manual_seed(0)
+    # tie_word_embeddings=False: per-layer export does not dedup tied weights on
+    # transformers 4 (see the TODO in unified_export_hf_streaming), and Kimi-K3 does not tie,
+    # so keep this test on the namespace/towers/config behaviour it exists for.
+    model = get_tiny_gemma3vl(tie_word_embeddings=False).cuda().eval()
+    # is_multimodal_model reads this, and the tiny fixtures leave it unset.
+    model.config.architectures = ["Gemma3ForConditionalGeneration"]
+    return model
+
+
+def _calib_vlm(language_model):
+    # use_cache=False: this calls the text model directly rather than the CausalLM wrapper,
+    # so a retained KV cache would make batch 2 build a mask for 31 positions against 16 keys.
+    for batch in CALIB_BATCHES:
+        language_model(batch.cuda(), use_cache=False)
+
+
+def test_vlm_export_matches_whole_model_export(tmp_path):
+    """A VLM calibrates its language model but must export the whole model.
+
+    Without the parent link the shards lose the ``language_model.`` namespace, the vision
+    tower and projector are never written at all, and config.json describes the submodel.
+    """
+    baseline_vlm = _build_vlm()
+    baseline_lm = get_language_model_from_vl(baseline_vlm)[-1]
+    cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}}
+    mtq.quantize(baseline_lm, cfg, _calib_vlm)
+    baseline_dir = tmp_path / "baseline"
+    export_hf_checkpoint(baseline_vlm, export_dir=baseline_dir)
+
+    vlm = _build_vlm()
+    language_model = get_language_model_from_vl(vlm)[-1]
+    export_dir = tmp_path / "fused"
+    with export_parent(vlm):
+        mtq.quantize(language_model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib_vlm)
+
+    exported = _load_checkpoint(export_dir)
+    _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported)
+    _assert_same_quant_config(baseline_dir, export_dir)
+
+    assert any(k.startswith("language_model.") for k in exported), "lost the VLM namespace"
+    assert any(k.startswith(("vision_tower", "multi_modal_projector")) for k in exported), (
+        "the towers outside the language model were not exported"
+    )
+    assert "vision_config" in json.loads((export_dir / "config.json").read_text()), (
+        "config.json describes the submodel, not the VLM"
+    )

From 96f5459783987ae7442c4fb50198e7dc07b80ae7 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Mon, 31 Aug 2026 20:46:16 +0000
Subject: [PATCH 2/5] docs(export): trim the K3 comments to what the code
 cannot say

Drops four comments that restated the code beside them -- the legacy-mapper
fallback, the _force_attn_implementation call site, the multimodal refusal whose
raise message already says it, and a test docstring paragraph each assert repeats
-- and trims seven others. What stays is the rejected alternatives
(build_reverse_name_mapper raising on 4.x, save_pretrained owning the name
reversal) and the hazards: longest-prefix ordering, the retained cache that
mis-sizes the second calibration batch, and the wrapper nesting order.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 examples/hf_ptq/example_utils.py              |  8 ++-----
 examples/hf_ptq/hf_ptq.py                     |  6 ++----
 modelopt/torch/export/layerwise_export.py     | 21 +++++++------------
 .../torch/quantization/plugins/huggingface.py |  4 ++--
 .../gpu/torch/export/test_layerwise_export.py | 15 +++++--------
 5 files changed, 18 insertions(+), 36 deletions(-)

diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py
index 6035fbba603..0db85c574be 100755
--- a/examples/hf_ptq/example_utils.py
+++ b/examples/hf_ptq/example_utils.py
@@ -707,10 +707,8 @@ def _resolve_init_config(hf_config, auto_model_module, ckpt_path, config_kwargs)
 def _force_attn_implementation(model, attn_implementation: str) -> None:
     """Re-apply the caller's ``attn_implementation`` after ``__init__``.
 
-    Kimi-K3's remote code rewrites ``text_config._attn_implementation`` to
-    ``flash_attention_2`` unconditionally, so an uninstalled backend fails later at the
-    export trace forward instead of at load. Drop this once flash-attn is installed
-    reliably here and the checkpoint is fixed upstream.
+    Kimi-K3's remote code overwrites it unconditionally, so an uninstalled backend fails at
+    the export forward rather than at load. Drop once the checkpoint is fixed upstream.
     """
     sub_configs = (v for v in vars(model.config).values() if isinstance(v, PretrainedConfig))
     for cfg in (model.config, *sub_configs):
@@ -983,8 +981,6 @@ def has_pack_quantized_config(config):
         )
     model.eval()
 
-    # Honour the caller's explicit choice even when remote modeling code overwrote it
-    # during __init__ (see _force_attn_implementation).
     if attn_implementation is not None:
         _force_attn_implementation(model, attn_implementation)
 
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 590ebee5193..f0483131f6c 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -785,8 +785,6 @@ def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) ->
     calibration begins, the user has already paid for the whole run.
     """
     if is_multimodal_model(full_model):
-        # Calibration runs on the extracted language model, so the exporter is pointed at the
-        # full model below. That needs the submodel to be reachable from it by identity.
         lineage = get_language_model_from_vl(full_model)
         language_model = lineage[-1] if lineage else None
         if language_model is None or all(m is not language_model for m in full_model.modules()):
@@ -869,8 +867,8 @@ def export_quantized(
         is_vlm = is_multimodal_model(full_model)
 
         if is_vlm:
-            # Skipped under per-layer export: calibration already wrote a config carrying the
-            # quantization_config, and the source config is unquantized.
+            # Not under per-layer export: it already wrote a config with the
+            # quantization_config, and this source config is unquantized.
             if not args.layerwise_export:
                 print(f"Saving original model config to {export_path}")
 
diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py
index ea7306722a0..9ab5118759c 100644
--- a/modelopt/torch/export/layerwise_export.py
+++ b/modelopt/torch/export/layerwise_export.py
@@ -156,10 +156,8 @@ def assert_layerwise_export_supported(model: nn.Module) -> None:
 def export_parent(parent: nn.Module):
     """Export the checkpoint for ``parent`` while calibration runs on one of its submodules.
 
-    Multimodal pipelines calibrate the extracted language model, but the shards and config
-    have to describe the whole VLM. The decoder layers are the same objects either way, so
-    walking the parent yields parent-namespace tensor names, the full config and the
-    untouched towers with no prefixing.
+    The decoder layers are the same objects either way, so walking the parent yields
+    parent-namespace names, the full config and the untouched towers with no prefixing.
     """
     token = _export_parent.set(parent)
     try:
@@ -184,15 +182,12 @@ def _resolve_export_parent(model: nn.Module) -> nn.Module:
 def build_legacy_name_mapper(model: nn.Module):
     """Hub-name mapper for transformers < 5, or ``None``.
 
-    The whole-model path writes through ``save_pretrained``, and it is save_pretrained that
-    reverses ``_checkpoint_conversion_mapping`` (in-memory ``model.language_model.*`` ->
-    published ``language_model.model.*``). Per-layer export writes shards directly, so it
-    never passes through that. ``build_reverse_name_mapper`` does not cover it either: it
-    reads transformers 5's ``conversion_mapping`` module and raises on 4.x.
+    ``save_pretrained`` is what reverses ``_checkpoint_conversion_mapping``, and per-layer
+    export writes shards directly without it. ``build_reverse_name_mapper`` is no help
+    either: it reads transformers 5's ``conversion_mapping`` and raises on 4.x.
 
-    The mapping is stored hub-pattern -> in-memory-prefix, so it is inverted here. Longest
-    in-memory prefix first, or a short rule shadows a longer one (``lm_head`` inside
-    ``model.language_model...``).
+    Rules are inverted (the mapping is stored hub -> in-memory) and applied longest-prefix
+    first, or a short one shadows a longer (``lm_head`` inside ``model.language_model...``).
     """
     import re
 
@@ -277,8 +272,6 @@ def __init__(
         try:
             self._name_mapper = build_reverse_name_mapper(model)
         except Exception as exc:
-            # transformers < 5 has no conversion_mapping module; fall back to the legacy
-            # mapping save_pretrained would have applied.
             self._name_mapper = build_legacy_name_mapper(model)
             if self._name_mapper is None:
                 warnings.warn(
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index b06e70e6f5d..7f62eb3ac8d 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1800,8 +1800,8 @@ def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None:
     if not _is_supported_hf_model(model):
         return None
 
-    # Descend through every wrapper before reading ``layers``: multimodal models nest them
-    # in either order, e.g. Kimi-K3 keeps its layers at ``language_model.model.layers``.
+    # Descend through every wrapper before reading ``layers``: the nesting order varies,
+    # e.g. Kimi-K3 keeps its layers at ``language_model.model.layers``.
     decoder, seen = model, set()
     while id(decoder) not in seen:
         seen.add(id(decoder))
diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py
index 4634c8af420..e91a6b0274e 100644
--- a/tests/gpu/torch/export/test_layerwise_export.py
+++ b/tests/gpu/torch/export/test_layerwise_export.py
@@ -490,9 +490,8 @@ def test_awq_is_refused(tmp_path, make_cfg, method, match):
 
 def _build_vlm():
     torch.manual_seed(0)
-    # tie_word_embeddings=False: per-layer export does not dedup tied weights on
-    # transformers 4 (see the TODO in unified_export_hf_streaming), and Kimi-K3 does not tie,
-    # so keep this test on the namespace/towers/config behaviour it exists for.
+    # tie_word_embeddings=False: per-layer export does not dedup tied weights, and K3 does
+    # not tie -- keep this test on the namespace/towers/config behaviour it exists for.
     model = get_tiny_gemma3vl(tie_word_embeddings=False).cuda().eval()
     # is_multimodal_model reads this, and the tiny fixtures leave it unset.
     model.config.architectures = ["Gemma3ForConditionalGeneration"]
@@ -500,18 +499,14 @@ def _build_vlm():
 
 
 def _calib_vlm(language_model):
-    # use_cache=False: this calls the text model directly rather than the CausalLM wrapper,
-    # so a retained KV cache would make batch 2 build a mask for 31 positions against 16 keys.
+    # use_cache=False: this calls the text model directly, so a retained cache would make
+    # batch 2 build a mask for 31 positions against 16 keys.
     for batch in CALIB_BATCHES:
         language_model(batch.cuda(), use_cache=False)
 
 
 def test_vlm_export_matches_whole_model_export(tmp_path):
-    """A VLM calibrates its language model but must export the whole model.
-
-    Without the parent link the shards lose the ``language_model.`` namespace, the vision
-    tower and projector are never written at all, and config.json describes the submodel.
-    """
+    """A VLM calibrates its language model but must export the whole model."""
     baseline_vlm = _build_vlm()
     baseline_lm = get_language_model_from_vl(baseline_vlm)[-1]
     cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)

From efac201e6f23ed0e4c7cbf864f1238282572fbac Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Mon, 31 Aug 2026 21:00:49 +0000
Subject: [PATCH 3/5] fix(export): refuse weight-tied modules whether or not
 they are quantized

The refusal only looked at quantized modules, which held while everything exported
came from the calibrated model: conversion quantizes every nn.Linear and
nn.Embedding, so a tie always involved one. Walking a VLM parent breaks that.
Gemma3's lm_head lives on the outer wrapper, outside the calibrated language
model, so it is untouched by conversion -- the tie went unseen and the alias was
written as a duplicate key that save_pretrained drops, leaving the per-layer
checkpoint with an extra language_model.lm_head.weight.

The tiny VLM fixture hid it in the other direction: get_tiny_gemma3vl forwards
tie_word_embeddings only to text_config, so the test asked for an untied model and
got a tied one. It now unties for real, which keeps it on the namespace, towers
and config behaviour it exists for, and a tied VLM is refused instead of silently
exported.

Tied-weight support is coming; until then this fails loudly rather than writing a
checkpoint that differs from the whole-model path.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 modelopt/torch/export/layerwise_export.py     | 21 +++++++++----------
 .../gpu/torch/export/test_layerwise_export.py | 12 ++++++++---
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py
index 9ab5118759c..933446bce92 100644
--- a/modelopt/torch/export/layerwise_export.py
+++ b/modelopt/torch/export/layerwise_export.py
@@ -83,8 +83,8 @@ def _module_formats(model: nn.Module) -> set:
     }
 
 
-def _tied_quantized_modules(model: nn.Module) -> list[str]:
-    """Quantized modules sharing a weight with another.
+def _tied_weight_modules(model: nn.Module) -> list[str]:
+    """Modules sharing a weight with another, quantized or not.
 
     Grouped by name, which survives offload: a ``data_ptr`` grouping sees nothing when the
     weights are on meta and would pass vacuously. Falls back to ``data_ptr`` when the model
@@ -95,7 +95,7 @@ def _tied_quantized_modules(model: nn.Module) -> list[str]:
     by_ptr: dict[int, list[str]] = {}
     for name, module in model.named_modules():
         weight = getattr(module, "weight", None)
-        if weight is None or not _is_quantized_module(module):
+        if weight is None:
             continue
         key = tied_map.group_key(f"{name}.weight")
         if key is not None:
@@ -128,16 +128,15 @@ def assert_layerwise_export_supported(model: nn.Module) -> None:
     """Raise unless per-layer export is valid for this model."""
     assert_formats_supported(model, "before calibration")
 
-    tied = _tied_quantized_modules(model)
+    tied = _tied_weight_modules(model)
     if tied:
         raise NotImplementedError(
-            f"layerwise export does not support weight-tied quantized modules {tied[:6]}: "
-            "the whole-model path merges their input_quantizer amaxes via "
-            "sync_tied_input_amax so both sides share one input_scale, which a per-layer "
-            "pass cannot do because a tie partner may be uncalibrated or already written. "
-            "Conversion quantizes every nn.Linear and nn.Embedding, so disabling their "
-            "quantizers does not lift this -- tie_word_embeddings models need "
-            "export_hf_checkpoint()."
+            f"layerwise export does not support weight-tied modules {tied[:6]}: quantized, "
+            "the whole-model path merges their input_quantizer amaxes via sync_tied_input_amax "
+            "so both sides share one input_scale, which a per-layer pass cannot do because a "
+            "tie partner may be uncalibrated or already written; unquantized, save_pretrained "
+            "drops the duplicate key and writing shards directly does not. "
+            "tie_word_embeddings models need export_hf_checkpoint()."
         )
 
     if dist.is_initialized() and dist.size() > 1:
diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py
index e91a6b0274e..12c806232f6 100644
--- a/tests/gpu/torch/export/test_layerwise_export.py
+++ b/tests/gpu/torch/export/test_layerwise_export.py
@@ -22,6 +22,7 @@
 
 import pytest
 import torch
+import torch.nn as nn
 from _test_utils.torch.transformers_models import (
     get_tiny_gemma3vl,
     get_tiny_llama,
@@ -418,7 +419,7 @@ def test_tied_embeddings_are_refused(tmp_path, make_cfg):
     model.config.architectures = ["LlamaForCausalLM"]
 
     cfg = _layerwise_cfg(tmp_path / "fused", tmp_path / "ckpt", base=make_cfg())
-    with pytest.raises(NotImplementedError, match="weight-tied quantized modules"):
+    with pytest.raises(NotImplementedError, match="weight-tied modules"):
         mtq.quantize(model, cfg, _calib)
 
 
@@ -490,9 +491,14 @@ def test_awq_is_refused(tmp_path, make_cfg, method, match):
 
 def _build_vlm():
     torch.manual_seed(0)
-    # tie_word_embeddings=False: per-layer export does not dedup tied weights, and K3 does
-    # not tie -- keep this test on the namespace/towers/config behaviour it exists for.
     model = get_tiny_gemma3vl(tie_word_embeddings=False).cuda().eval()
+    # The kwarg only reaches text_config; the outer config still ties lm_head to the
+    # embedding, and tied weights are refused. Untie for real so this test stays on the
+    # namespace/towers/config behaviour it exists for.
+    model.config.tie_word_embeddings = False
+    model._tied_weights_keys = {}
+    model.all_tied_weights_keys = {}
+    model.lm_head.weight = nn.Parameter(model.lm_head.weight.detach().clone())
     # is_multimodal_model reads this, and the tiny fixtures leave it unset.
     model.config.architectures = ["Gemma3ForConditionalGeneration"]
     return model

From 937d220c316b9c74cc98928a7aa01c0924b31f53 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Mon, 31 Aug 2026 21:13:57 +0000
Subject: [PATCH 4/5] fix(export): close the review findings on the VLM export
 path

Four fixes from PR review.

The legacy hub-name mapper treated the mapping's hub side as a literal when it is
a regex. On Qwen2-VL / Qwen2.5-VL / GLM-4V, whose mapping carries a lookahead
group, an in-memory key came back with the raw pattern embedded in it -- silently,
no exception -- and a mapping with an alphanumeric escape raised re.error
mid-export. It now strips the group exactly as save_pretrained does when reversing
the same mapping, and neutralises backslashes in the remainder so nothing is read
as a template escape. Verified against both mappings.

The export parent was installed on only one of the two calibration entry points.
--low_memory_mode routes through mtq.calibrate, which reaches the same layerwise
machinery with the same algorithm block, so a VLM in low-memory mode exported the
submodel: no language_model. prefix, no towers, submodel config -- the exact
failure export_parent exists to prevent, silently. The context manager now spans
both branches.

Decoder discovery descended past a level that had both its own layers and a
wrapper child, returning None (reported as an unsupported architecture) or an
unrelated deeper ModuleList. It now takes the shallowest layers, restoring the
precedence the pre-existing code had, and returns only an nn.ModuleList as the
annotation promises.

The VLM test calibrated a bare extracted language model, but hf_ptq first runs
mtq.quantize(tower, disabled) over every non-language sibling, so in production
the towers carry disabled quantizers by the time the exporter walks the parent --
which changes what _module_formats and the tail dispatch see. The test now
reproduces that arrangement.

Also corrects the offload recipe's description: nothing in it configures offload,
which comes from --offload_folder and the memory budgets. Its only functional
difference from the resident recipe is the narrower expert scope.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 examples/hf_ptq/hf_ptq.py                     | 17 +++++++-------
 modelopt/torch/export/layerwise_export.py     | 16 +++++++++----
 .../torch/quantization/plugins/huggingface.py |  8 ++++---
 ..._only-kv_fp8_layerwise_export_offload.yaml |  9 +++++---
 .../gpu/torch/export/test_layerwise_export.py | 23 +++++++++++++++++--
 5 files changed, 53 insertions(+), 20 deletions(-)

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index f0483131f6c..b60f5b420b5 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -754,14 +754,15 @@ def mono_quantize(
                     else None,
                 )
 
-        if calibration_only:
-            language_model = mtq.calibrate(
-                language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop
-            )
-        else:
-            # A VLM calibrates its language model but must export the whole thing.
-            parent = full_model if args.layerwise_export else language_model
-            with export_parent(parent):
+        # A VLM calibrates its language model but must export the whole thing. Both branches
+        # reach the same layerwise machinery, so the link is set over both.
+        parent = full_model if args.layerwise_export else language_model
+        with export_parent(parent):
+            if calibration_only:
+                language_model = mtq.calibrate(
+                    language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop
+                )
+            else:
                 language_model = mtq.quantize(
                     language_model, quant_cfg, forward_loop=calibrate_loop
                 )
diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py
index 933446bce92..a1ddf58f3f5 100644
--- a/modelopt/torch/export/layerwise_export.py
+++ b/modelopt/torch/export/layerwise_export.py
@@ -18,6 +18,7 @@
 import contextlib
 import contextvars
 import json
+import re
 import warnings
 from pathlib import Path
 
@@ -179,7 +180,7 @@ def _resolve_export_parent(model: nn.Module) -> nn.Module:
 
 
 def build_legacy_name_mapper(model: nn.Module):
-    """Hub-name mapper for transformers < 5, or ``None``.
+    r"""Hub-name mapper for transformers < 5, or ``None``.
 
     ``save_pretrained`` is what reverses ``_checkpoint_conversion_mapping``, and per-layer
     export writes shards directly without it. ``build_reverse_name_mapper`` is no help
@@ -187,14 +188,21 @@ def build_legacy_name_mapper(model: nn.Module):
 
     Rules are inverted (the mapping is stored hub -> in-memory) and applied longest-prefix
     first, or a short one shadows a longer (``lm_head`` inside ``model.language_model...``).
+    The hub side is a regex: its groups are stripped exactly as save_pretrained strips them
+    when reversing the same mapping, and the remainder has its backslashes neutralised so a
+    pattern like ``layers\.(\d+)`` substitutes literally instead of raising on ``\d``.
     """
-    import re
-
     mapping = getattr(model, "_checkpoint_conversion_mapping", None)
     if not mapping:
         return None
     rules = sorted(
-        ((re.compile("^" + re.escape(mem)), hub.lstrip("^")) for hub, mem in mapping.items()),
+        (
+            (
+                re.compile("^" + re.escape(mem)),
+                re.sub(r"\(.*\)", "", hub.lstrip("^")).replace("\\", "\\\\"),
+            )
+            for hub, mem in mapping.items()
+        ),
         key=lambda r: -len(r[0].pattern),
     )
 
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 7f62eb3ac8d..f0a2f0733ff 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1800,18 +1800,20 @@ def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None:
     if not _is_supported_hf_model(model):
         return None
 
-    # Descend through every wrapper before reading ``layers``: the nesting order varies,
-    # e.g. Kimi-K3 keeps its layers at ``language_model.model.layers``.
+    # Take the shallowest ``layers`` but keep descending while wrappers remain: the nesting
+    # order varies, e.g. Kimi-K3 keeps its layers at ``language_model.model.layers``.
     decoder, seen = model, set()
     while id(decoder) not in seen:
         seen.add(id(decoder))
+        if isinstance(getattr(decoder, "layers", None), nn.ModuleList):
+            return decoder.layers
         for attr in ("model", "language_model"):
             inner = getattr(decoder, attr, None)
             if isinstance(inner, nn.Module):
                 decoder = inner
                 break
         else:
-            return getattr(decoder, "layers", None)
+            return None
 
     return None
 
diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
index b041123cfa7..eeba1cc5c3e 100644
--- a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
+++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
@@ -23,9 +23,12 @@ metadata:
   recipe_type: ptq
   description: >
     NVFP4 W4A4 on routed experts only, FP8 KV cache, max layerwise calibration, with each
-    decoder layer exported to its own shard as soon as it is calibrated. Same intent as
-    nvfp4_experts_only-kv_fp8_layerwise_export, but scoped and paired for a model too large
-    to hold resident: use with --offload_folder and per-device memory budgets.
+    decoder layer exported to its own shard as soon as it is calibrated.
+
+    The only functional difference from nvfp4_experts_only-kv_fp8_layerwise_export is the
+    narrower expert scope below; nothing here configures offload, which comes from
+    --offload_folder and the per-device memory budgets. The suffix records the pairing this
+    recipe was validated in -- a model too large to hold resident -- not a setting it sets.
 
     Expert scoping is '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE
     models the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and
diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py
index 12c806232f6..0e6e4cc4ca5 100644
--- a/tests/gpu/torch/export/test_layerwise_export.py
+++ b/tests/gpu/torch/export/test_layerwise_export.py
@@ -489,6 +489,25 @@ def test_awq_is_refused(tmp_path, make_cfg, method, match):
         mtq.quantize(_build_model(), cfg, _calib)
 
 
+def _disable_quant_on_towers(vlm):
+    """Mirror ``extract_and_prepare_language_model_from_vl``: towers get disabled quantizers.
+
+    hf_ptq runs this before calibration, so by the time the exporter walks the parent the
+    vision tower and projector carry ``TensorQuantizer`` children. That changes what
+    ``_module_formats`` and the tail dispatch see, so the test has to reproduce it.
+    """
+    lineage = get_language_model_from_vl(vlm)
+    language_model, ancestors = lineage[-1], lineage[:-1]
+    disabled = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"}
+    memo = set(ancestors) | {language_model}
+    for ancestor in ancestors:
+        for _, module in ancestor.named_children():
+            if module not in memo:
+                mtq.quantize(module, copy.deepcopy(disabled), forward_loop=None)
+                memo.add(module)
+    return language_model
+
+
 def _build_vlm():
     torch.manual_seed(0)
     model = get_tiny_gemma3vl(tie_word_embeddings=False).cuda().eval()
@@ -514,7 +533,7 @@ def _calib_vlm(language_model):
 def test_vlm_export_matches_whole_model_export(tmp_path):
     """A VLM calibrates its language model but must export the whole model."""
     baseline_vlm = _build_vlm()
-    baseline_lm = get_language_model_from_vl(baseline_vlm)[-1]
+    baseline_lm = _disable_quant_on_towers(baseline_vlm)
     cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
     cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}}
     mtq.quantize(baseline_lm, cfg, _calib_vlm)
@@ -522,7 +541,7 @@ def test_vlm_export_matches_whole_model_export(tmp_path):
     export_hf_checkpoint(baseline_vlm, export_dir=baseline_dir)
 
     vlm = _build_vlm()
-    language_model = get_language_model_from_vl(vlm)[-1]
+    language_model = _disable_quant_on_towers(vlm)
     export_dir = tmp_path / "fused"
     with export_parent(vlm):
         mtq.quantize(language_model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib_vlm)

From 5ee399edf4b24873aa3541157ca5c67781247ca6 Mon Sep 17 00:00:00 2001
From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Date: Mon, 31 Aug 2026 21:32:09 +0000
Subject: [PATCH 5/5] refactor(recipes): move the K3 layerwise-export recipe
 out of the general tier

ptq.md's own rule is that general/ptq/ holds recipes whose wildcards work "on any
architecture whose module names follow the usual conventions", and that a recipe
earns a model-tier place only when it must deviate. This one exists precisely
because it deviates: '*block_sparse_moe*' over-matches on Kimi-K3's fine-grained
MoE naming, hitting 552 modules the vendor left unquantized. That is the
architecture-aware quant_cfg case the doc lists, so it belongs beside the other
K3 recipe rather than as a 26th near-duplicate of a general one it differs from
in two lines.

Renamed to match its new sibling's prefix while dropping two tokens that were not
earning their place: _only, and _offload, which named something the recipe does
not configure -- offload comes from --offload_folder and the memory budgets.

general/ptq/ returns to 25 and the table row moves to the checkpoint-mirrors
prose, where the existing K3 entry already is.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
---
 .../ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml}  |  8 ++++----
 modelopt_recipes/ptq.md                              | 12 ++++++++++--
 2 files changed, 14 insertions(+), 6 deletions(-)
 rename modelopt_recipes/{general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml => huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml} (88%)

diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml
similarity index 88%
rename from modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
rename to modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml
index eeba1cc5c3e..ea1e26d897e 100644
--- a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml
+++ b/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml
@@ -25,10 +25,10 @@ metadata:
     NVFP4 W4A4 on routed experts only, FP8 KV cache, max layerwise calibration, with each
     decoder layer exported to its own shard as soon as it is calibrated.
 
-    The only functional difference from nvfp4_experts_only-kv_fp8_layerwise_export is the
-    narrower expert scope below; nothing here configures offload, which comes from
-    --offload_folder and the per-device memory budgets. The suffix records the pairing this
-    recipe was validated in -- a model too large to hold resident -- not a setting it sets.
+    Kimi-K3's counterpart to general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export, from
+    which it differs only in the narrower expert scope below -- the reason it does not live
+    in the general tier. It was validated on a model too large to hold resident, but nothing
+    here configures offload: that comes from --offload_folder and the memory budgets.
 
     Expert scoping is '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE
     models the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and
diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md
index 3f85166158e..e201f6a0d6f 100644
--- a/modelopt_recipes/ptq.md
+++ b/modelopt_recipes/ptq.md
@@ -29,7 +29,7 @@ supported combinations.
 ### The shipped recipes
 
 <details>
-<summary>All 26 <code>general/ptq/</code> recipes (click to expand)</summary>
+<summary>All 25 <code>general/ptq/</code> recipes (click to expand)</summary>
 
 | Recipe | Model body | KV cache | Calibration |
 |--------|-----------|----------|-------------|
@@ -49,7 +49,6 @@ supported combinations.
 | `nvfp4_experts_only-kv_fp8_layerwise` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise |
 | `nvfp4_experts_only-kv_fp8_layerwise_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (non-mutating, for disk offload) |
 | `nvfp4_experts_only-kv_fp8_layerwise_export` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (exports each layer as it is calibrated) |
-| `nvfp4_experts_only-kv_fp8_layerwise_export_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise, tuned for accelerate-offloaded models |
 | `nvfp4_experts_only_mse-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | MSE + FP8 sweep |
 | `nvfp4_experts_only_input_scale1-kv_fp8_cast` | NVFP4 W4A4, MoE experts only, expert `input_scale` pinned to 1.0 | FP8 (constant amax) | max (weights); expert activations uncalibrated |
 | `nvfp4_omlp_only-kv_fp8` | NVFP4 W4A4, o_proj + MLP/MoE | FP8 (calibrated) | max |
@@ -361,6 +360,15 @@ checkpoint's** quant config verbatim:
   cache remain BF16. Because the 2.8T source uses packed MXFP4 expert tensors,
   use the calibration-free streaming converter in `examples/kimi/` rather than
   the in-memory `hf_ptq.py` flow.
+- **`models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export`** is the
+  in-memory counterpart to the above: NVFP4 W4A4 on the routed experts with an FP8
+  KV cache, calibrated and **exported one decoder layer at a time**, so a run that
+  outlasts its GPU session resumes without recalibrating or re-exporting finished
+  layers. It scopes experts as `*.experts.*` rather than the general recipes'
+  `*block_sparse_moe*`, which on K3 also matches `shared_experts.*` and
+  `routed_expert_*` -- 552 modules the vendor left unquantized. Pair it with
+  `--offload_folder` and per-device memory budgets; nothing in the recipe itself
+  configures offload.
 - **`models/mistralai/Mistral-Medium-3.5-128B/ptq/nvfp4-max-calib`** mirrors
   `nvidia/Mistral-Medium-3.5-128B-NVFP4`: decoder MLP layers 4–86 use NVFP4
   W4A4, edge MLP layers 0–3 and 87 use FP8 W8A8, and all attention projections