Fix: Release per-layer expert weights in layerwise export under offload (#2466)

### What does this PR do?

Type of change: Bug fix

Layerwise export leaks one layer's worth of quantized tensors per layer
on offloaded models, and dies of OOM partway through a large MoE. Two
things accumulate, both because the offload window cannot reclaim what
the export pass adds inside it.

**1. Per-expert holder modules.** `_export_fused_experts` splits a fused
MoE experts module into per-expert holders and attaches them to the live
model:

```python
proj = nn.Module();  proj.weight = wrapper.weight   # packed U8
expert.add_module(proj_name, proj)
module.add_module(str(idx), expert)
```

They are plain `nn.Module`s built *inside* the weight-access window, so
they carry no accelerate `_hf_hook`.
`weight_access_and_writeback_context` closes by iterating the modules it
collected at entry and calling `hook.post_forward()` through each one's
own offload hook — the holders satisfy neither condition, so nothing
returns them to meta.

**2. Scale buffers.** `_export_quantized_weight` registers
`weight_scale` / `weight_scale_2` / `input_scale` on the layer's
pre-existing, hooked sub-modules, and `AlignDevicesHook.post_forward`
runs with `offload_buffers=False`:

```
after post_forward: {'weight': 'meta', 'weight_scale': 'cpu'}
```

The packed weight goes back to meta; the scales do not. For
`_QuantMoELinear` models this half is the whole leak on its own —
`_reconstruct_fused_moe_linear` restacks every expert's scales into one
`register_buffer` on the hooked wrapper.

A whole-model export never notices: one pass, write the state dict,
exit. Layerwise runs the same pass once per decoder layer, so every
finished layer stays resident.

### Measurements

Through unmodified `examples/hf_ptq/hf_ptq.py` on a Qwen3.5-MoE-shaped
model (10 layers, 64 experts), printing `torch.cuda.memory_allocated()`
after each exported layer:

| placement | per layer | over 9 layers |
| --- | --- | --- |
| offload | +0.052 GiB | 0.579 → 1.047 GiB |
| resident | -0.135 GiB | falls, as designed |

+0.052 GiB is exactly one layer's quantized experts: packed U8 100.7M/2
= 0.047 GiB plus FP8 block scales 100.7M/16 = 0.006 GiB.

At scale it is fatal rather than wasteful. Qwen/Qwen3.8-2.4T-A95B (92
layers, 512 experts) leaks 12.9 GB packed + 1.6 GB scales per layer, so
92 layers want 1.33 TB that no budget on a 283 GB card or 952 GB host
absorbs. The run died of CUDA OOM at layer 16/92 with
`--max_gpu_memory_gb 240`, and at `--max_gpu_memory_gb 30` leaked the
same 15 GB/layer onto the host instead.

### The fix

`_release_exported_tensors` (`model_utils.py`) is a context manager that
snapshots each sub-module's child-module and buffer names on entry, and
on exit drops whatever appeared. Persisting happens *inside* the block,
so "release only once it is on disk" is structural rather than a
comment.

Both packing sites use it: `LayerwiseExporter.export_layer` and the
offload decoder loop in `_export_transformers_checkpoint_streaming`.

The streaming writer had solved the same leak inline with a heuristic —
null every CUDA buffer, and every CUDA parameter on a hook-less module —
and that block is deleted in favour of the shared helper. Keying on
*what the pass added* rather than on device and hook presence drops two
assumptions that only held for a terminal, offloaded export: it no
longer nulls buffers the layer already had, nor parameters of
sub-modules accelerate simply did not hook. That is also what makes it
safe for the layerwise path, where resident models are supported and the
model outlives the export.

Deliberately out of scope: the FSDP2 per-unit loop in
`collect_export_tensors` keeps its per-unit holders. That predates this
PR, this PR does not touch that loop, and closing it needs its own
change and its own test.

### Usage

No API change. Existing layerwise export under offload simply stops
growing:

```bash
python examples/hf_ptq/hf_ptq.py \
    --pyt_ckpt_path Qwen/Qwen3.8-2.4T-A95B \
    --qformat nvfp4 \
    --export_path /path/to/export \
    --max_gpu_memory_gb 240
```

### Testing

- `tests/unit/torch/export` and `tests/unit/torch/quantization` — 1246
passed
- `tests/gpu/torch/export/test_layerwise_export.py` +
`test_offload_export.py` — 35 passed
- `cuda_alloc` over 10 offloaded layers: 0.526 → 0.518 GiB (-0.008), was
+0.468
- Exported checkpoint byte-identical to the unfixed run: 7841 tensors, 0
mismatches, max abs diff 0.0; `hf_quant_config.json` / `config.json` /
index identical
- Full Qwen3.8-2.4T-A95B PTQ then completed all 92 layers: peak GPU 117
GB of 283, peak host RSS 71 GB of 952, flat across 50 consecutive layers
at 23-26 s/layer

Coverage added: the existing `test_export_creates_per_expert_submodules`
now runs the export inside the context manager and asserts the holders
are gone on exit, and a new test in `test_offload_export.py` pins the
`offload_buffers=False` behaviour the buffer half exists for —
export-registered scales dropped, pre-existing buffers untouched.

`tests/gpu/torch/export/test_fsdp2_export.py` reports 34 failures in my
environment. They are **pre-existing and unrelated**: the same 34 fail
identically on this branch and on the merge-base (`2b1f33d0ef`), with
byte-identical failure sets and runtimes within 3 s. All 34 are `Failed:
Timeout (>120.0s)` from hung NCCL collectives, with zero assertion
failures.

The GPU export suites and the whole-model measurements above were run at
`2999d7cfb7`. The two commits since — swapping the holder marker for a
child-name diff, and moving the helper to `model_utils` — are covered by
the unit suites; a GPU re-run before merge is worthwhile.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- 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 new offload test for
the buffer half, and the existing fused-experts export test now covers
holder release
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A — `layerwise.export_dir` is new in the unreleased 0.48.0, so this
bug was introduced and fixed within the same cycle
- Did you get Claude approval on this PR?: ✅

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

modelopt/torch/export/layerwise_export.py

@@ -36,7 +36,7 @@ from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCol
 from modelopt.torch.utils import distributed as dist
 
 from .layer_utils import sync_moe_gate_up_amax
-from .model_utils import TiedWeightMap, get_language_model_from_vl
+from .model_utils import TiedWeightMap, _release_exported_tensors, get_language_model_from_vl
 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 +309,19 @@ class LayerwiseExporter:
         )
         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/model_utils.py

@@ -15,6 +15,7 @@
 """Utility functions for model type detection and classification."""
 
 import warnings
+from contextlib import contextmanager
 
 import torch.nn as nn
 
@@ -80,6 +81,28 @@ __all__ = [
 ]
 
 
+@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 existing sub-modules and attach per-expert holder
+    modules. Neither an accelerate offload window nor an FSDP2 reshard reclaims those, so a
+    caller that runs the pass once per unit accumulates them. An export that raises releases
+    nothing, leaving the unit 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 get_model_type(model):
     """Try get the model type from the model name. If not found, return None."""
     for k, v in MODEL_NAME_TO_TYPE.items():

modelopt/torch/export/unified_export_hf_streaming.py

@@ -43,7 +43,7 @@ from modelopt.torch.quantization.utils.core_utils import (
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 from modelopt.torch.utils import distributed as _dist
 
-from .model_utils import get_export_units
+from .model_utils import _release_exported_tensors, get_export_units
 from .quant_aware_conversion import _build_reverse_rules, build_reverse_name_mapper
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
@@ -453,7 +453,10 @@ def _export_transformers_checkpoint_streaming(
     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)
@@ -465,34 +468,6 @@ def _export_transformers_checkpoint_streaming(
                     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

@@ -36,7 +36,7 @@ from _test_utils.torch.quantization.tied_modules import (
 )
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.model_utils import TiedWeightMap
+from modelopt.torch.export.model_utils import TiedWeightMap, _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 +167,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,6 +26,7 @@ from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_par
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
+from modelopt.torch.export.model_utils import _release_exported_tensors
 from modelopt.torch.export.moe_utils import _export_fused_experts
 from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
@@ -481,24 +482,29 @@ class TestExportFusedExperts:
         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
+            # 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")
+
+        # 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)