refactor(export): scope the VLM export parent to a context manager

Review follow-ups on the per-layer VLM path:

- Replace the _layerwise_export_parent module attribute with an
  export_parent() context manager over a ContextVar. The attribute held an
  nn.Module, so it had to bypass nn.Module.__setattr__ to avoid a
  parent-child cycle, and it was never cleared.
- Reduce _force_attn_implementation to the model config and its direct
  sub-configs. Kimi-K3's remote code only rewrites text_config, so the
  recursive walk and its bookkeeping were unused.
- Drop _MAX_DECODER_UNWRAP_DEPTH: descend through the wrappers first and
  read `layers` at the bottom, which generalises main's fixed two-level
  unwrap instead of bounding a shallow-first search.
- Drop the tied-weight alias handling. Kimi-K3 sets
  tie_word_embeddings=False and tied weights are unsupported on main, so
  it belongs in a separate change.

Testing: tests/gpu/torch/export/test_layerwise_export.py 25 passed;
test_layerwise_calibrate.py and tests/examples/hf_ptq/test_example_utils.py
pass. test_offload_export.py has 2 failures that reproduce on the branch
without these changes.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

examples/hf_ptq/example_utils.py

@@ -41,6 +41,7 @@
     AutoModelForCausalLM,
     AutoProcessor,
     AutoTokenizer,
+    PretrainedConfig,
     PreTrainedTokenizerBase,
     ProcessorMixin,
 )
@@ -704,41 +705,17 @@ def _resolve_init_config(hf_config, auto_model_module, ckpt_path, config_kwargs)
 
 
 def _force_attn_implementation(model, attn_implementation: str) -> None:
-    """Set ``_attn_implementation`` on the model config and every nested sub-config.
-
-    Some remote modeling code overrides the requested attention implementation inside
-    ``__init__`` -- Kimi-K3 rewrites it to ``flash_attention_2`` unconditionally, ignoring
-    ``--attn_implementation``. Export runs a trace forward, so a backend whose compiled
-    extension is unavailable in this environment fails there rather than at load.
-
-    Sub-configs are walked because multimodal models keep separate ones per tower, and
-    remote code typically rewrites the *nested* config (Kimi-K3 rewrites ``text_config``
-    from ``KimiLinearModel.__init__``). Layer modules hold a reference to the same config
-    object, so overriding it here takes effect on the next forward. Only applied when the
-    caller asked for an implementation explicitly.
-    """
-    pending, seen, changed = [model.config], set(), []
-    while pending:
-        cfg = pending.pop()
-        if cfg is None or id(cfg) in seen:
-            continue
-        seen.add(id(cfg))
-        current = getattr(cfg, "_attn_implementation", None)
-        if current is not None and current != attn_implementation:
-            try:
-                cfg._attn_implementation = attn_implementation
-                changed.append(f"{type(cfg).__name__}: {current} -> {attn_implementation}")
-            except Exception as e:  # pragma: no cover - depends on remote config class
-                warnings.warn(f"Could not apply attn_implementation on {type(cfg).__name__}: {e}")
-        pending.extend(
-            getattr(cfg, sub, None)
-            for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder")
-        )
+    """Re-apply the caller's ``attn_implementation`` after ``__init__``.
 
-    if changed:
-        print("Re-applied the requested attention implementation after model init:")
-        for line in changed:
-            print(f"  {line}")
+    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):

examples/hf_ptq/hf_ptq.py

@@ -78,7 +78,7 @@
     has_spec_opt,
     save_expert_token_count_table,
 )
-from modelopt.torch.export.layerwise_export import EXPORT_PARENT_ATTR
+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
@@ -759,11 +759,12 @@ def mono_quantize(
                 language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop
             )
         else:
-            if args.layerwise_export and language_model is not full_model:
-                # Into __dict__: nn.Module.__setattr__ would register the parent as a
-                # submodule of its own child and make named_modules() recurse forever.
-                language_model.__dict__[EXPORT_PARENT_ATTR] = full_model
-            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:

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,26 +147,36 @@ def assert_layerwise_export_supported(model: nn.Module) -> None:
         )
 
 
-#: Set by the caller on the calibrated submodel when the checkpoint must describe a larger
-#: model. Multimodal pipelines calibrate the extracted language model, but the shards and
-#: config have to describe the whole VLM.
-EXPORT_PARENT_ATTR = "_layerwise_export_parent"
+_export_parent: contextvars.ContextVar[nn.Module | None] = contextvars.ContextVar(
+    "layerwise_export_parent", default=None
+)
 
 
-def resolve_export_parent(model: nn.Module) -> nn.Module:
-    """Return the model the checkpoint should describe: the marked parent, or ``model``.
+@contextlib.contextmanager
+def export_parent(parent: nn.Module):
+    """Export the checkpoint for ``parent`` while calibration runs on one of its submodules.
 
-    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. Membership is checked by identity, not by name.
+    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.
     """
-    parent = getattr(model, EXPORT_PARENT_ATTR, None)
+    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_ATTR} was set to a {type(parent).__name__} that does not "
-            "contain the calibrated model."
+            f"export_parent() was given a {type(parent).__name__} that does not contain the "
+            "calibrated model."
         )
     return parent
 
@@ -220,7 +231,7 @@ def __init__(
 
         Runs before calibration, so nothing amax-dependent exists yet.
         """
-        model = resolve_export_parent(model)
+        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)
@@ -275,16 +286,6 @@ def __init__(
                     "match the original HF hub checkpoint."
                 )
 
-        # Tied aliases the whole-model export omits and lets the loader re-tie.
-        raw_tied = (
-            set(getattr(model, "_tied_weights_keys", None) or [])
-            if getattr(model.config, "tie_word_embeddings", False)
-            else set()
-        )
-        self._tied_alias_keys = (
-            {self._name_mapper(k) for k in raw_tied} if self._name_mapper else raw_tied
-        )
-
     def export_layer(
         self,
         layer_idx: int,
@@ -480,8 +481,6 @@ def _collect(self, out: dict[str, torch.Tensor], full_key: str, tensor: torch.Te
             return
         if self._name_mapper is not None:
             new_key = self._name_mapper(new_key)
-        if new_key in self._tied_alias_keys:
-            return
         out[new_key] = new_value.detach().contiguous().cpu()
 
     def _write_index(self) -> None:

modelopt/torch/quantization/plugins/huggingface.py

@@ -1796,30 +1796,22 @@ def is_homogeneous_hf_model(model: nn.Module) -> bool:
     return len(layer_classes) == 1
 
 
-#: How deep to unwrap before giving up. Each level is one of the wrappers below, and no
-#: released architecture nests more than a handful; the bound only stops a cycle.
-_MAX_DECODER_UNWRAP_DEPTH = 8
-
-
 def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None:
     if not _is_supported_hf_model(model):
         return None
 
-    # Unwrap iteratively rather than testing each wrapper once: multimodal models nest them
-    # in either order and to arbitrary depth. Kimi-K3 keeps its layers at
-    # ``language_model.model.layers``, so a single ``.model`` then ``.language_model`` walk
-    # stops on the intermediate wrapper and reports the architecture unsupported.
-    decoder = model
-    for _ in range(_MAX_DECODER_UNWRAP_DEPTH):
-        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 None
+            return getattr(decoder, "layers", None)
 
     return None
 

tests/gpu/torch/export/test_layerwise_export.py

@@ -31,8 +31,8 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.layerwise_export import (
-    EXPORT_PARENT_ATTR,
     LayerwiseExporter,
+    export_parent,
     layer_shard_name,
 )
 from modelopt.torch.export.model_utils import get_language_model_from_vl
@@ -490,7 +490,10 @@ def test_awq_is_refused(tmp_path, make_cfg, method, match):
 
 def _build_vlm():
     torch.manual_seed(0)
-    model = get_tiny_gemma3vl().cuda().eval()
+    # 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
@@ -519,11 +522,9 @@ def test_vlm_export_matches_whole_model_export(tmp_path):
 
     vlm = _build_vlm()
     language_model = get_language_model_from_vl(vlm)[-1]
-    # Into __dict__: nn.Module.__setattr__ would register the parent as a submodule of its
-    # own child and make named_modules() recurse forever.
-    language_model.__dict__[EXPORT_PARENT_ATTR] = vlm
     export_dir = tmp_path / "fused"
-    mtq.quantize(language_model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib_vlm)
+    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)