Fix: Release per-layer expert weights in layerwise export under offload

modelopt/torch/export/layerwise_export.py

@@ -37,6 +37,7 @@
 
 from .layer_utils import sync_moe_gate_up_amax
 from .model_utils import TiedWeightMap, get_language_model_from_vl
+from .moe_utils import _release_exported_tensors
 from .quant_aware_conversion import build_reverse_name_mapper, revert_quant_config_names
 from .quant_format import FUSION_FREE_FORMATS, QUANTIZATION_NVFP4
 from .quant_utils import (
@@ -309,16 +310,19 @@ def export_layer(
         )
         self._unify_shared_quantization_params(layer_module, layer_inputs)
 
-        for sub_name, sub_mod in layer_module.named_modules():
-            full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name
-            _dispatch_export_handler(full_name, sub_mod, self._ctx)
-        _reconstruct_fused_moe_linear(layer_module)
+        # The shard on disk is the artifact once this block closes; nothing reads the
+        # layer again.
+        with _release_exported_tensors(layer_module):
+            for sub_name, sub_mod in layer_module.named_modules():
+                full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name
+                _dispatch_export_handler(full_name, sub_mod, self._ctx)
+            _reconstruct_fused_moe_linear(layer_module)
 
-        prefix = f"{layer_name}." if layer_name else ""
-        for key, tensor in layer_module.state_dict().items():
-            self._collect(tensors, prefix + key, tensor)
+            prefix = f"{layer_name}." if layer_name else ""
+            for key, tensor in layer_module.state_dict().items():
+                self._collect(tensors, prefix + key, tensor)
 
-        save_file(tensors, str(self._export_dir / layer_shard_name(layer_idx)))
+            save_file(tensors, str(self._export_dir / layer_shard_name(layer_idx)))
 
     def _unify_shared_quantization_params(
         self, layer_module: nn.Module, layer_inputs: list | None

modelopt/torch/export/moe_utils.py

@@ -17,6 +17,7 @@
 
 import copy
 import warnings
+from contextlib import contextmanager
 from pathlib import Path
 
 import torch
@@ -231,6 +232,28 @@ def _export_fused_experts(
     _delete_fused_moe_source_attrs(module)
 
 
+@contextmanager
+def _release_exported_tensors(root: nn.Module):
+    """Drop what the export pass adds to ``root``, once the block has persisted it.
+
+    The handlers register scale buffers on ``root``'s existing sub-modules and
+    :func:`_export_fused_experts` attaches per-expert holder modules; an accelerate offload
+    window reclaims neither, so running the pass once per layer accumulates them. An export
+    that raises releases nothing, leaving the layer intact to be inspected.
+    """
+    before = {name: (set(mod._modules), set(mod._buffers)) for name, mod in root.named_modules()}
+
+    yield
+
+    # list(): deleting a child mutates the _modules dict the traversal walks.
+    for name, module in list(root.named_modules()):
+        children_before, buffers_before = before.get(name, (set(), set()))
+        for child_name in set(module._modules) - children_before:
+            delattr(module, child_name)
+        for buf_name in set(module._buffers) - buffers_before:
+            module._buffers[buf_name] = None
+
+
 def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | None = None):
     """Collect expert_token_count from all quantized MoE layers and save as an HTML table.
 

modelopt/torch/export/unified_export_hf_streaming.py

@@ -44,6 +44,7 @@
 from modelopt.torch.utils import distributed as _dist
 
 from .model_utils import get_export_units
+from .moe_utils import _release_exported_tensors
 from .quant_aware_conversion import _build_reverse_rules, build_reverse_name_mapper
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
@@ -456,7 +457,10 @@ def _is_persistent_buffer(name: str) -> bool:
     for layer_name, layer_module in model.named_modules():
         if id(layer_module) not in decoder_layer_ids:
             continue
-        with enable_weight_access_and_writeback(layer_module, model, names, writeback=False):
+        with (
+            enable_weight_access_and_writeback(layer_module, model, names, writeback=False),
+            _release_exported_tensors(layer_module),
+        ):
             for sub_name, sub_mod in layer_module.named_modules():
                 full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name
                 _dispatch_export_handler(full_name, sub_mod, ctx)
@@ -468,34 +472,6 @@ def _is_persistent_buffer(name: str) -> bool:
                     continue
                 seen_keys.add(full_key)
                 _stream_tensor(full_key, tensor)
-            # Release GPU tensors added by export handlers before hook.post_forward
-            # runs, to prevent cross-layer accumulation on disk-offloaded models.
-            #
-            # Two categories accumulate without explicit cleanup:
-            #
-            # 1. CUDA *buffers* on any sub-module (weight_scale, weight_scale_2,
-            #    input_scale): AlignDevicesHook.post_forward uses offload_buffers=False
-            #    by default, so it never offloads buffers.  Pre-existing buffers in
-            #    disk-offloaded layers live on CPU, so any CUDA buffer encountered here
-            #    was registered by the export handlers and is safe to drop.
-            #
-            # 2. CUDA *parameters* on modules WITHOUT _hf_hook: _export_fused_experts
-            #    creates fresh nn.Module objects (one per expert x projection) and adds
-            #    them to the layer via add_module() *after* weight_access_and_writeback
-            #    captured its materialized list.  hook.post_forward never visits these
-            #    new modules, so their packed NVFP4 weight parameters (~5 GB per MoE
-            #    layer) stay live on GPU.  Modules WITH _hf_hook are original model
-            #    modules whose parameters hook.post_forward will meta-ify; leave those
-            #    alone.
-            for sub_mod in layer_module.modules():
-                for buf_name in list(sub_mod._buffers):
-                    buf = sub_mod._buffers[buf_name]
-                    if buf is not None and buf.device.type == "cuda":
-                        sub_mod._buffers[buf_name] = None
-                if not hasattr(sub_mod, "_hf_hook"):
-                    for param_name, param in list(sub_mod._parameters.items()):
-                        if param is not None and param.device.type == "cuda":
-                            sub_mod._parameters[param_name] = None
         torch.cuda.empty_cache()
 
     # Non-decoder modules whose weights are not directly readable (embed_tokens, norm,

tests/unit/torch/export/test_offload_export.py

@@ -37,6 +37,7 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.model_utils import TiedWeightMap
+from modelopt.torch.export.moe_utils import _release_exported_tensors
 from modelopt.torch.export.quant_format import KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4
 from modelopt.torch.export.quant_utils import (
     _get_kv_cache_postprocess_config,
@@ -167,6 +168,34 @@ def test_meta_guard_not_raised_for_real_weight():
     _export_quantized_weight(linear, torch.float32)
 
 
+# ---------------------------------------------------------------------------
+# _release_exported_tensors
+# ---------------------------------------------------------------------------
+
+
+def test_release_exported_tensors_drops_what_the_offload_window_leaves():
+    """post_forward runs with offload_buffers=False, so export's scale buffers outlive it."""
+    layer = nn.Module()
+    layer.self_attn = nn.Linear(16, 16, bias=False)
+    layer.register_buffer("rotary_emb_inv_freq", torch.randn(8))
+    mtq.quantize(layer, mtq.FP8_DEFAULT_CFG, lambda m: m.self_attn(torch.randn(1, 16)))
+    _offload_module(layer.self_attn)
+
+    hook = layer.self_attn._hf_hook
+    assert hook.offload_buffers is False
+
+    with _release_exported_tensors(layer):
+        hook.pre_forward(layer.self_attn)
+        _export_quantized_weight(layer.self_attn, torch.float32)
+        assert layer.self_attn.weight_scale.device.type != "meta"
+        hook.post_forward(layer.self_attn, None)
+        assert layer.self_attn.weight.device.type == "meta"
+        assert layer.self_attn.weight_scale.device.type != "meta"  # the leak
+
+    assert layer.self_attn._buffers["weight_scale"] is None
+    assert layer._buffers["rotary_emb_inv_freq"] is not None
+
+
 # ---------------------------------------------------------------------------
 # _StreamingShardWriter
 # ---------------------------------------------------------------------------

tests/unit/torch/quantization/plugins/test_fused_experts.py

@@ -26,7 +26,7 @@
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
-from modelopt.torch.export.moe_utils import _export_fused_experts
+from modelopt.torch.export.moe_utils import _export_fused_experts, _release_exported_tensors
 from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name
@@ -481,24 +481,29 @@ def forward_loop(m):
         mtq.quantize(model, quant_cfg, forward_loop=forward_loop)
         converted = model.moe.experts
 
-        _export_fused_experts(converted, torch.float16)
+        with _release_exported_tensors(converted):
+            _export_fused_experts(converted, torch.float16)
+
+            # Verify per-expert submodules exist
+            for idx in range(NUM_EXPERTS):
+                expert_mod = getattr(converted, str(idx), None)
+                assert expert_mod is not None, f"Missing expert submodule {idx}"
+                assert hasattr(expert_mod, "gate_proj"), f"Expert {idx} missing gate_proj"
+                assert hasattr(expert_mod, "up_proj"), f"Expert {idx} missing up_proj"
+                assert hasattr(expert_mod, "down_proj"), f"Expert {idx} missing down_proj"
+
+                assert expert_mod.gate_proj.weight.shape == (INTERMEDIATE_DIM, HIDDEN_DIM)
+                assert expert_mod.up_proj.weight.shape == (INTERMEDIATE_DIM, HIDDEN_DIM)
+                assert expert_mod.down_proj.weight.shape == (HIDDEN_DIM, INTERMEDIATE_DIM)
+
+            # Verify fused params are removed
+            assert not hasattr(converted, "gate_up_proj")
+            assert not hasattr(converted, "down_proj")
+            assert not hasattr(converted, "gate_up_proj_weight_quantizers")
 
-        # Verify per-expert submodules exist
+        # Leaving the block releases the holders; nothing else can free them.
         for idx in range(NUM_EXPERTS):
-            expert_mod = getattr(converted, str(idx), None)
-            assert expert_mod is not None, f"Missing expert submodule {idx}"
-            assert hasattr(expert_mod, "gate_proj"), f"Expert {idx} missing gate_proj"
-            assert hasattr(expert_mod, "up_proj"), f"Expert {idx} missing up_proj"
-            assert hasattr(expert_mod, "down_proj"), f"Expert {idx} missing down_proj"
-
-            assert expert_mod.gate_proj.weight.shape == (INTERMEDIATE_DIM, HIDDEN_DIM)
-            assert expert_mod.up_proj.weight.shape == (INTERMEDIATE_DIM, HIDDEN_DIM)
-            assert expert_mod.down_proj.weight.shape == (HIDDEN_DIM, INTERMEDIATE_DIM)
-
-        # Verify fused params are removed
-        assert not hasattr(converted, "gate_up_proj")
-        assert not hasattr(converted, "down_proj")
-        assert not hasattr(converted, "gate_up_proj_weight_quantizers")
+            assert not hasattr(converted, str(idx)), f"Expert submodule {idx} survived release"
 
         self._cleanup_registry(expert_type)