From ad662f75dc080eb364f67b815c8d14bbba86ed30 Mon Sep 17 00:00:00 2001
From: cehongwang <wangcehong@gmail.com>
Date: Mon, 12 Oct 2026 23:08:47 +0000
Subject: [PATCH 1/4] feat(dynamo): release model weights once their last
 engine is built

A new opt-in setting, release_weights_after_build. TensorRT holds a network's
weights only until its engine is built, and a built engine carries its own
copy, yet the model, the ExportedProgram and the partitioned graph modules
keep every weight referenced until compilation ends. With several engines
(resource partitioning, PyTorch fallbacks) compilation therefore holds all
weights while building each engine.

With the setting on, compile_module counts, for each weight storage, the
partitions that read it. After each TensorRT partition is built it releases
the storages no remaining partition reads, and before the first build the
lowered module's parameters that no partition reads (for example a model
head the traced function never calls). Storages read by PyTorch-run
partitions are counted but never released. The storage itself is released
(resize_(0)), because dropping references cannot free what the caller's
model still points to, so the source module must not be used afterwards.
Buffers are never released as unread: mutable buffers are lifted into
engine inputs and inlined back into the compiled module. A release-handler
hook lets an owner release its own storages (for example memory-mapped
weights) instead.

Measured on pi0.5 (lerobot, bf16, Anthony Shoumikhin's ExecuTorch export,
CPU-resident model, resource partitioner with 9 engines, A40): host + GPU
peak 25.7 -> 18.9 GiB (4.85 GiB released engine by engine, 1.90 GiB of
never-read weights, mostly the two lm_heads, before the first engine).
---
 py/torch_tensorrt/dynamo/_compiler.py         |  31 ++++
 py/torch_tensorrt/dynamo/_defaults.py         |   1 +
 py/torch_tensorrt/dynamo/_settings.py         |   8 +
 py/torch_tensorrt/dynamo/_weight_release.py   | 128 ++++++++++++++++
 .../test_release_weights_after_build.py       | 145 ++++++++++++++++++
 5 files changed, 313 insertions(+)
 create mode 100644 py/torch_tensorrt/dynamo/_weight_release.py
 create mode 100644 tests/py/dynamo/models/test_release_weights_after_build.py

diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index 8810ab8726..d001cbeabf 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -30,6 +30,11 @@
 )
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache, DiskEngineCache
 from torch_tensorrt.dynamo._exporter import replace_execute_engine_no_op_node
+from torch_tensorrt.dynamo._weight_release import (
+    WeightReleaser,
+    parameter_storages,
+    weight_storages,
+)
 from torch_tensorrt.dynamo.conversion import (
     CompilationSettings,
     UnsupportedOperatorException,
@@ -71,6 +76,7 @@
     parse_graph_io,
     pin_torch_executed_state,
     prepare_inputs,
+    release_host_and_device_memory,
     to_torch_device,
     to_torch_tensorrt_device,
     validate_optimization_profiles,
@@ -515,6 +521,7 @@ def compile(
     cpu_memory_budget: Optional[int] = _defaults.CPU_MEMORY_BUDGET,
     enable_resource_partitioning: bool = _defaults.ENABLE_RESOURCE_PARTITIONING,
     dynamically_allocate_resources: bool = _defaults.DYNAMICALLY_ALLOCATE_RESOURCES,
+    release_weights_after_build: bool = _defaults.RELEASE_WEIGHTS_AFTER_BUILD,
     decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION,
     disabled_constant_fold_exclusions: Collection[
         str
@@ -623,6 +630,9 @@ def compile(
         enable_resource_partitioning (bool): Enable resource-aware partitioning. This is useful when the model is large and the CPU memory is limited.
         cpu_memory_budget (Optional[int]): The maximum amount of CPU memory to use for the compilation. If the compilation requires more memory than this budget, the compilation will fail.
         dynamically_allocate_resources (bool): Dynamically allocate resources during engine execution.
+        release_weights_after_build (bool): Release the memory of each model weight once the last TensorRT engine that reads it is built,
+            and of weights no partition reads before the first build, so compilation holds only the weights still to be built. The source
+            module shares that memory and must not be used afterwards. Weights read by subgraphs that run in PyTorch are kept.
         decompose_attention (bool): Whether to decompose attention layers into smaller operations
             instead of using the attention converters. When combined with ``use_fp32_acc=True``,
             decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only
@@ -836,6 +846,7 @@ def compile(
         "enable_resource_partitioning": enable_resource_partitioning,
         "cpu_memory_budget": cpu_memory_budget,
         "dynamically_allocate_resources": dynamically_allocate_resources,
+        "release_weights_after_build": release_weights_after_build,
         "decompose_attention": decompose_attention,
         "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions,
         "attn_bias_is_causal": attn_bias_is_causal,
@@ -1560,6 +1571,19 @@ def preserve_module_specs(
     # Iterate over all components that can be accelerated
     # Generate the corresponding TRT Module for those
 
+    weight_releaser = None
+    if settings.release_weights_after_build and not settings.dryrun:
+        weight_releaser = WeightReleaser(
+            sub
+            for _, sub in partitioned_module.named_children()
+            if isinstance(sub, torch.fx.GraphModule)
+        )
+        # The lowered module still holds weights no partition reads (for example a model head the
+        # traced function never calls); nothing needs them once partitioning is done.
+        weight_releaser.release_unread(
+            parameter_storages(gm), keep=weight_storages(partitioned_module)
+        )
+
     # Here we delete the frozen parameters from the graph module. Note this does not affect the submodules. We are going to delete the frozen parameters from the submodules in the convert_module function.
     # This is done to release CPU memory.
     for attr in dir(gm):
@@ -1657,6 +1681,10 @@ def preserve_module_specs(
         torch.cuda.empty_cache()
         # Create TRT engines from submodule
         if not settings.dryrun:
+            # Collected before conversion, which deletes the submodule's frozen parameters.
+            submodule_weights = (
+                weight_storages(submodule) if weight_releaser is not None else {}
+            )
             trt_module = convert_module(
                 submodule,
                 submodule_inputs,
@@ -1666,6 +1694,9 @@ def preserve_module_specs(
             )
 
             trt_modules[name] = trt_module
+            if weight_releaser is not None:
+                weight_releaser.built(submodule_weights, name)
+                release_host_and_device_memory()
 
             if _debugger_config:
                 if _debugger_config.save_engine_profile:
diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py
index cffa5af521..af34d746df 100644
--- a/py/torch_tensorrt/dynamo/_defaults.py
+++ b/py/torch_tensorrt/dynamo/_defaults.py
@@ -79,6 +79,7 @@
 ENABLE_RESOURCE_PARTITIONING = False
 CPU_MEMORY_BUDGET = None
 DYNAMICALLY_ALLOCATE_RESOURCES = False
+RELEASE_WEIGHTS_AFTER_BUILD = False
 DECOMPOSE_ATTENTION = False
 DISABLED_CONSTANT_FOLD_EXCLUSIONS = frozenset[str]()
 ATTN_BIAS_IS_CAUSAL = True
diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py
index 8314e51abe..9b72c9cdfd 100644
--- a/py/torch_tensorrt/dynamo/_settings.py
+++ b/py/torch_tensorrt/dynamo/_settings.py
@@ -51,6 +51,7 @@
     OPTIMIZATION_LEVEL,
     PASS_THROUGH_BUILD_FAILURES,
     REFIT_IDENTICAL_ENGINE_WEIGHTS,
+    RELEASE_WEIGHTS_AFTER_BUILD,
     REQUIRE_FULL_COMPILATION,
     REUSE_CACHED_ENGINES,
     SPARSE_WEIGHTS,
@@ -151,6 +152,11 @@ class CompilationSettings:
         autocast_calibration_dataloader (Optional[torch.utils.data.DataLoader]): The dataloader to use for autocast calibration. Default is None.
         offload_module_to_cpu (bool): Offload the model to CPU to reduce memory footprint during compilation
         dynamically_allocate_resources (bool): Dynamically allocate resources for TensorRT engines
+        release_weights_after_build (bool): Release the memory of each model weight once the last
+            TensorRT engine that reads it is built, and of weights no partition reads before the first
+            build, so compilation holds only the weights still to be built. The model's tensors share
+            that memory, so the source module (and an ``ExportedProgram`` exported from it) must not be
+            used after compilation. Weights read by subgraphs that run in PyTorch are kept. Default is False.
         decompose_attention (bool): Whether to decompose attention layers into smaller operations
             instead of using the attention converters. When combined with ``use_fp32_acc=True``,
             decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only
@@ -236,6 +242,7 @@ class CompilationSettings:
     enable_resource_partitioning: bool = ENABLE_RESOURCE_PARTITIONING
     cpu_memory_budget: Optional[int] = CPU_MEMORY_BUDGET
     dynamically_allocate_resources: bool = DYNAMICALLY_ALLOCATE_RESOURCES
+    release_weights_after_build: bool = RELEASE_WEIGHTS_AFTER_BUILD
     decompose_attention: bool = DECOMPOSE_ATTENTION
     disabled_constant_fold_exclusions: Collection[str] = (
         DISABLED_CONSTANT_FOLD_EXCLUSIONS
@@ -290,6 +297,7 @@ def __setstate__(self, state: dict[str, Any]) -> None:
             )
         )
         state.setdefault("fallback_data_dependent_ops", FALLBACK_DATA_DEPENDENT_OPS)
+        state.setdefault("release_weights_after_build", RELEASE_WEIGHTS_AFTER_BUILD)
         state.setdefault("build_route", BUILD_ROUTE)
         state.setdefault("tune_build_routes", TUNE_BUILD_ROUTES)
         state.setdefault("tune_build_route_file", TUNE_BUILD_ROUTE_FILE)
diff --git a/py/torch_tensorrt/dynamo/_weight_release.py b/py/torch_tensorrt/dynamo/_weight_release.py
new file mode 100644
index 0000000000..fd0f1194eb
--- /dev/null
+++ b/py/torch_tensorrt/dynamo/_weight_release.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Release model weights during compilation, as soon as no TensorRT engine still to be built reads them.
+
+TensorRT holds a network's weights by reference only until its engine is built, and a built engine carries
+its own copy. With several engines (resource partitioning, PyTorch fallbacks), a weight is dead once the
+last engine that reads it is built, yet the model, the ``ExportedProgram`` and the partitioned graph
+modules keep it referenced until compilation ends. ``compile_module`` uses these helpers, when
+``release_weights_after_build`` is set, to free that memory: the storage itself is released, because
+dropping references cannot free what the caller's model still points to.
+"""
+
+import logging
+from collections import Counter
+from typing import Callable, Dict, Iterable, List
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# A handler returns True if it released the storage (for example, by dropping the pages of a
+# memory-mapped weight); release_storage falls back to resizing the storage to zero.
+ReleaseHandler = Callable[[torch.UntypedStorage], bool]
+_release_handlers: List[ReleaseHandler] = []
+
+
+def register_release_handler(handler: ReleaseHandler) -> None:
+    """Let a storage owner release its own storages (for example, file-backed weights)."""
+    if handler not in _release_handlers:
+        _release_handlers.append(handler)
+
+
+def unregister_release_handler(handler: ReleaseHandler) -> None:
+    if handler in _release_handlers:
+        _release_handlers.remove(handler)
+
+
+def _get_attr(module: torch.nn.Module, target: str) -> object:
+    obj: object = module
+    for part in target.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def weight_storages(gm: torch.fx.GraphModule) -> Dict[int, torch.UntypedStorage]:
+    """The storages read by ``gm``'s ``get_attr`` nodes, keyed by data pointer (tied weights collapse)."""
+    found: Dict[int, torch.UntypedStorage] = {}
+    for node in gm.graph.nodes:
+        if node.op != "get_attr":
+            continue
+        try:
+            value = _get_attr(gm, str(node.target))
+        except AttributeError:
+            continue
+        if isinstance(value, torch.Tensor) and not value.is_meta:
+            storage = value.untyped_storage()
+            if storage.nbytes():
+                found[storage.data_ptr()] = storage
+    return found
+
+
+def parameter_storages(module: torch.nn.Module) -> Dict[int, torch.UntypedStorage]:
+    """The storages of ``module``'s parameters, keyed by data pointer. Buffers are left out: mutable
+    buffers are lifted into engine inputs (no partition reads them through ``get_attr``) and are
+    inlined back into the compiled module afterwards, so they must survive compilation.
+    """
+    found: Dict[int, torch.UntypedStorage] = {}
+    for tensor in module.parameters():
+        if not tensor.is_meta:
+            storage = tensor.untyped_storage()
+            if storage.nbytes():
+                found[storage.data_ptr()] = storage
+    return found
+
+
+def release_storage(storage: torch.UntypedStorage) -> int:
+    """Release a storage's memory; return the bytes released (0 if it could not be released)."""
+    nbytes = int(storage.nbytes())
+    if nbytes == 0:
+        return 0
+    for handler in _release_handlers:
+        if handler(storage):
+            return nbytes
+    try:
+        storage.resize_(0)
+    except RuntimeError as e:  # e.g. a non-resizable storage over an external buffer
+        logger.debug(f"Could not release a {nbytes}-byte weight storage: {e}")
+        return 0
+    return nbytes
+
+
+class WeightReleaser:
+    """Counts, for each weight storage, the partitions that read it, and releases a storage once the
+    last TensorRT partition that reads it has been built. Storages read by PyTorch partitions are
+    counted but never decremented, so they are never released."""
+
+    def __init__(self, partitions: Iterable[torch.fx.GraphModule]):
+        self.uses: Counter[int] = Counter()
+        for partition in partitions:
+            self.uses.update(weight_storages(partition).keys())
+
+    def release_unread(
+        self, candidates: Dict[int, torch.UntypedStorage], keep: Iterable[int] = ()
+    ) -> int:
+        """Release the candidate storages that no partition reads (and that are not in ``keep``)."""
+        kept = set(self.uses) | set(keep)
+        released = sum(
+            release_storage(s) for ptr, s in candidates.items() if ptr not in kept
+        )
+        if released:
+            logger.info(
+                f"Released {released / 2**20:.0f} MiB of weights that no partition reads"
+            )
+        return released
+
+    def built(self, storages: Dict[int, torch.UntypedStorage], name: str = "") -> int:
+        """Record that a TensorRT partition reading ``storages`` has been built; release the ones done."""
+        released = 0
+        for ptr, storage in storages.items():
+            self.uses[ptr] -= 1
+            if self.uses[ptr] <= 0:
+                released += release_storage(storage)
+        if released:
+            logger.info(
+                f"Released {released / 2**20:.0f} MiB of weights after building {name}"
+            )
+        return released
diff --git a/tests/py/dynamo/models/test_release_weights_after_build.py b/tests/py/dynamo/models/test_release_weights_after_build.py
new file mode 100644
index 0000000000..ed6db9e2d0
--- /dev/null
+++ b/tests/py/dynamo/models/test_release_weights_after_build.py
@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import pytest
+import torch
+import torch_tensorrt
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo._weight_release import (
+    WeightReleaser,
+    parameter_storages,
+    register_release_handler,
+    release_storage,
+    unregister_release_handler,
+    weight_storages,
+)
+
+
+def test_release_weights_after_build_default_off():
+    assert CompilationSettings().release_weights_after_build is False
+
+
+def test_old_pickle_without_field_defaults_to_false():
+    state = CompilationSettings().__dict__.copy()
+    state.pop("release_weights_after_build", None)
+    restored = CompilationSettings.__new__(CompilationSettings)
+    restored.__setstate__(state)
+    assert restored.release_weights_after_build is False
+
+
+class _Linear(torch.nn.Module):
+    def __init__(self, weight):
+        super().__init__()
+        self.weight = weight
+
+    def forward(self, x):
+        return x @ self.weight
+
+
+def _traced(weight: torch.nn.Parameter) -> torch.fx.GraphModule:
+    return torch.fx.symbolic_trace(_Linear(weight))
+
+
+def _param(n: int = 4) -> torch.nn.Parameter:
+    return torch.nn.Parameter(torch.randn(n, n), requires_grad=False)
+
+
+def test_releaser_releases_after_the_last_reader():
+    own, shared, torch_side = _param(), _param(), _param()
+    first, second = _traced(own), _traced(shared)
+    also_shared, in_torch = _traced(shared), _traced(torch_side)
+    releaser = WeightReleaser([first, second, also_shared, in_torch])
+
+    releaser.built(weight_storages(first), "first")
+    assert own.untyped_storage().nbytes() == 0
+    assert shared.untyped_storage().nbytes() > 0
+
+    releaser.built(weight_storages(second), "second")
+    assert shared.untyped_storage().nbytes() > 0  # also_shared still reads it
+    releaser.built(weight_storages(also_shared), "also_shared")
+    assert shared.untyped_storage().nbytes() == 0
+
+    # in_torch runs in PyTorch: it is counted but never built, so its weight stays.
+    assert torch_side.untyped_storage().nbytes() > 0
+
+
+def test_releaser_releases_unread_parameters_only():
+    used, unread = _param(), _param()
+    holder = torch.nn.Module()
+    holder.used, holder.unread = used, unread
+    holder.register_buffer("buf", torch.randn(4))
+    releaser = WeightReleaser([_traced(used)])
+    released = releaser.release_unread(parameter_storages(holder))
+    assert released == unread.numel() * unread.element_size()
+    assert unread.untyped_storage().nbytes() == 0
+    assert used.untyped_storage().nbytes() > 0
+    assert holder.buf.untyped_storage().nbytes() > 0  # buffers are never candidates
+
+
+def test_release_handler_takes_precedence():
+    weight = _param()
+    seen = []
+
+    def handler(storage):
+        seen.append(storage.data_ptr())
+        return True
+
+    register_release_handler(handler)
+    try:
+        assert release_storage(weight.untyped_storage()) == weight.numel() * 4
+    finally:
+        unregister_release_handler(handler)
+    assert seen == [weight.untyped_storage().data_ptr()]
+    assert (
+        weight.untyped_storage().nbytes() > 0
+    )  # the handler, not resize_, owned the release
+
+
+def test_non_resizable_storage_is_skipped():
+    weight = torch.frombuffer(bytearray(64), dtype=torch.float32)
+    assert release_storage(weight.untyped_storage()) == 0
+    assert weight.untyped_storage().nbytes() == 64
+
+
+class _TwoEngines(torch.nn.Module):
+    """Two TensorRT partitions split by ops forced to PyTorch, a weight shared by both, one read only
+    by the PyTorch partition, and one the forward never reads."""
+
+    def __init__(self):
+        super().__init__()
+        p = lambda *s: torch.nn.Parameter(
+            torch.randn(*s), requires_grad=False
+        )  # noqa: E731
+        self.w_a, self.w_b, self.w_shared = p(16, 16), p(16, 16), p(16)
+        self.w_torch, self.unused_head = p(16), p(16, 1000)
+
+    def forward(self, x):
+        y = x @ self.w_a + self.w_shared
+        y = torch.sin(y * self.w_torch)  # aten.mul / aten.sin run in PyTorch
+        return y @ self.w_b + self.w_shared
+
+
+@pytest.mark.parametrize("release", [True, False])
+def test_compile_releases_weights(release):
+    model = _TwoEngines().eval().cuda()
+    x = torch.randn(4, 16, device="cuda")
+    with torch.no_grad():
+        expected = model(x)
+    ep = torch.export.export(model, (x,))
+    trt_gm = torch_tensorrt.dynamo.compile(
+        ep,
+        arg_inputs=[x],
+        min_block_size=1,
+        torch_executed_ops={torch.ops.aten.sin.default, torch.ops.aten.mul.Tensor},
+        release_weights_after_build=release,
+    )
+    with torch.no_grad():
+        actual = trt_gm(x)
+    torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3)
+
+    state = ep.state_dict
+    for name in ("w_a", "w_b", "w_shared", "unused_head"):
+        assert (state[name].untyped_storage().nbytes() == 0) == release, name
+    assert (
+        state["w_torch"].untyped_storage().nbytes() > 0
+    )  # read by the PyTorch partition

From 73e6ec4cb5b364421fd106365575b0b1c17b36f8 Mon Sep 17 00:00:00 2001
From: cehongwang <wangcehong@gmail.com>
Date: Tue, 13 Oct 2026 05:27:45 +0000
Subject: [PATCH 2/4] review: keep uncounted storages, count the root graph,
 guard MutableTorchTensorRTModule, release on the single-engine path

---
 py/torch_tensorrt/dynamo/_compiler.py         | 46 ++++++++++++----
 py/torch_tensorrt/dynamo/_settings.py         |  4 +-
 py/torch_tensorrt/dynamo/_weight_release.py   | 54 +++++++++++--------
 .../runtime/_MutableTorchTensorRTModule.py    |  5 ++
 .../test_release_weights_after_build.py       | 53 ++++++++++++++++--
 5 files changed, 123 insertions(+), 39 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index d001cbeabf..f3246dd0bb 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -632,7 +632,9 @@ def compile(
         dynamically_allocate_resources (bool): Dynamically allocate resources during engine execution.
         release_weights_after_build (bool): Release the memory of each model weight once the last TensorRT engine that reads it is built,
             and of weights no partition reads before the first build, so compilation holds only the weights still to be built. The source
-            module shares that memory and must not be used afterwards. Weights read by subgraphs that run in PyTorch are kept.
+            module shares that memory and must not be used afterwards, even if compilation fails. Weights read by subgraphs that run in
+            PyTorch are kept. Applies to ``torch_tensorrt.dynamo.compile``; no effect under ``torch.compile``, and not supported by
+            ``MutableTorchTensorRTModule``.
         decompose_attention (bool): Whether to decompose attention layers into smaller operations
             instead of using the attention converters. When combined with ``use_fp32_acc=True``,
             decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only
@@ -1293,6 +1295,14 @@ def _build_user_symbol_bounds(
     return user_symbol_bounds
 
 
+def _warn_if_refit_source_released(settings: CompilationSettings) -> None:
+    if settings.strip_engine_weights and not settings.immutable_weights:
+        logger.warning(
+            "release_weights_after_build releases the source weights; engines built with "
+            "strip_engine_weights must be refit from a freshly loaded checkpoint."
+        )
+
+
 @fn_supports_debugger  # type: ignore[misc]
 def compile_module(
     gm: torch.fx.GraphModule,
@@ -1446,13 +1456,25 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
             gm,
             user_symbol_bounds=user_symbol_bounds,
         )
-        return convert_module(
+        single_releaser = None
+        single_weights: Dict[int, torch.UntypedStorage] = {}
+        if settings.release_weights_after_build:
+            _warn_if_refit_source_released(settings)
+            single_releaser = WeightReleaser([gm])
+            single_releaser.release_unread(parameter_storages(gm))
+            single_weights = weight_storages(gm)
+        single_module = convert_module(
             gm,
             convert_inputs,
             settings=settings,
             name="_run_on_acc_0",
             engine_cache=engine_cache,
         )
+        if single_releaser is not None and single_releaser.built(
+            single_weights, "_run_on_acc_0"
+        ):
+            release_host_and_device_memory()
+        return single_module
 
     # Store the original input spec for later use
     original_in_spec = getattr(gm, "_in_spec", None)
@@ -1573,16 +1595,19 @@ def preserve_module_specs(
 
     weight_releaser = None
     if settings.release_weights_after_build and not settings.dryrun:
+        _warn_if_refit_source_released(settings)
+        # The partitioned module's own graph runs in PyTorch: it counts as a reader that is never built.
         weight_releaser = WeightReleaser(
-            sub
-            for _, sub in partitioned_module.named_children()
-            if isinstance(sub, torch.fx.GraphModule)
+            [partitioned_module]
+            + [
+                sub
+                for _, sub in partitioned_module.named_children()
+                if isinstance(sub, torch.fx.GraphModule)
+            ]
         )
         # The lowered module still holds weights no partition reads (for example a model head the
         # traced function never calls); nothing needs them once partitioning is done.
-        weight_releaser.release_unread(
-            parameter_storages(gm), keep=weight_storages(partitioned_module)
-        )
+        weight_releaser.release_unread(parameter_storages(gm))
 
     # Here we delete the frozen parameters from the graph module. Note this does not affect the submodules. We are going to delete the frozen parameters from the submodules in the convert_module function.
     # This is done to release CPU memory.
@@ -1694,8 +1719,9 @@ def preserve_module_specs(
             )
 
             trt_modules[name] = trt_module
-            if weight_releaser is not None:
-                weight_releaser.built(submodule_weights, name)
+            if weight_releaser is not None and weight_releaser.built(
+                submodule_weights, name
+            ):
                 release_host_and_device_memory()
 
             if _debugger_config:
diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py
index 9b72c9cdfd..81894a5894 100644
--- a/py/torch_tensorrt/dynamo/_settings.py
+++ b/py/torch_tensorrt/dynamo/_settings.py
@@ -156,7 +156,9 @@ class CompilationSettings:
             TensorRT engine that reads it is built, and of weights no partition reads before the first
             build, so compilation holds only the weights still to be built. The model's tensors share
             that memory, so the source module (and an ``ExportedProgram`` exported from it) must not be
-            used after compilation. Weights read by subgraphs that run in PyTorch are kept. Default is False.
+            used after compilation, even if compilation fails. Weights read by subgraphs that run in
+            PyTorch are kept. Applies to ``torch_tensorrt.dynamo.compile``; no effect under
+            ``torch.compile`` (parameters arrive as graph inputs). Default is False.
         decompose_attention (bool): Whether to decompose attention layers into smaller operations
             instead of using the attention converters. When combined with ``use_fp32_acc=True``,
             decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only
diff --git a/py/torch_tensorrt/dynamo/_weight_release.py b/py/torch_tensorrt/dynamo/_weight_release.py
index fd0f1194eb..198fa1712c 100644
--- a/py/torch_tensorrt/dynamo/_weight_release.py
+++ b/py/torch_tensorrt/dynamo/_weight_release.py
@@ -16,6 +16,7 @@
 from typing import Callable, Dict, Iterable, List
 
 import torch
+from torch.utils._python_dispatch import is_traceable_wrapper_subclass
 
 logger = logging.getLogger(__name__)
 
@@ -43,8 +44,17 @@ def _get_attr(module: torch.nn.Module, target: str) -> object:
     return obj
 
 
+def _add_storage(found: Dict[int, torch.UntypedStorage], tensor: torch.Tensor) -> None:
+    if tensor.is_meta or is_traceable_wrapper_subclass(tensor):
+        return  # a wrapper subclass's storage is a placeholder; its inner tensors are not reachable here
+    storage = tensor.untyped_storage()
+    if storage.nbytes():
+        found[storage.data_ptr()] = storage
+
+
 def weight_storages(gm: torch.fx.GraphModule) -> Dict[int, torch.UntypedStorage]:
-    """The storages read by ``gm``'s ``get_attr`` nodes, keyed by data pointer (tied weights collapse)."""
+    """The storages read by ``gm``'s own ``get_attr`` nodes (not its submodules'), keyed by data
+    pointer (tied weights collapse)."""
     found: Dict[int, torch.UntypedStorage] = {}
     for node in gm.graph.nodes:
         if node.op != "get_attr":
@@ -53,10 +63,8 @@ def weight_storages(gm: torch.fx.GraphModule) -> Dict[int, torch.UntypedStorage]
             value = _get_attr(gm, str(node.target))
         except AttributeError:
             continue
-        if isinstance(value, torch.Tensor) and not value.is_meta:
-            storage = value.untyped_storage()
-            if storage.nbytes():
-                found[storage.data_ptr()] = storage
+        if isinstance(value, torch.Tensor):
+            _add_storage(found, value)
     return found
 
 
@@ -67,10 +75,7 @@ def parameter_storages(module: torch.nn.Module) -> Dict[int, torch.UntypedStorag
     """
     found: Dict[int, torch.UntypedStorage] = {}
     for tensor in module.parameters():
-        if not tensor.is_meta:
-            storage = tensor.untyped_storage()
-            if storage.nbytes():
-                found[storage.data_ptr()] = storage
+        _add_storage(found, tensor)
     return found
 
 
@@ -91,22 +96,20 @@ def release_storage(storage: torch.UntypedStorage) -> int:
 
 
 class WeightReleaser:
-    """Counts, for each weight storage, the partitions that read it, and releases a storage once the
-    last TensorRT partition that reads it has been built. Storages read by PyTorch partitions are
-    counted but never decremented, so they are never released."""
+    """Counts, for each weight storage, the graphs that read it, and releases a storage once the last
+    TensorRT partition that reads it has been built. Graphs that run in PyTorch (PyTorch partitions,
+    the partitioned module's own graph) are counted but never built, so their storages are never
+    released."""
 
-    def __init__(self, partitions: Iterable[torch.fx.GraphModule]):
+    def __init__(self, readers: Iterable[torch.fx.GraphModule]):
         self.uses: Counter[int] = Counter()
-        for partition in partitions:
-            self.uses.update(weight_storages(partition).keys())
-
-    def release_unread(
-        self, candidates: Dict[int, torch.UntypedStorage], keep: Iterable[int] = ()
-    ) -> int:
-        """Release the candidate storages that no partition reads (and that are not in ``keep``)."""
-        kept = set(self.uses) | set(keep)
+        for reader in readers:
+            self.uses.update(weight_storages(reader).keys())
+
+    def release_unread(self, candidates: Dict[int, torch.UntypedStorage]) -> int:
+        """Release the candidate storages that no graph reads."""
         released = sum(
-            release_storage(s) for ptr, s in candidates.items() if ptr not in kept
+            release_storage(s) for ptr, s in candidates.items() if ptr not in self.uses
         )
         if released:
             logger.info(
@@ -118,8 +121,13 @@ def built(self, storages: Dict[int, torch.UntypedStorage], name: str = "") -> in
         """Record that a TensorRT partition reading ``storages`` has been built; release the ones done."""
         released = 0
         for ptr, storage in storages.items():
+            if ptr not in self.uses:
+                continue  # not counted at partition time (e.g. rebound since): another reader may hold it
             self.uses[ptr] -= 1
-            if self.uses[ptr] <= 0:
+            if self.uses[ptr] == 0:
+                del self.uses[
+                    ptr
+                ]  # a later allocation at this address is a different storage
                 released += release_storage(storage)
         if released:
             logger.info(
diff --git a/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py
index 2bec183c31..e4d9abf871 100644
--- a/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py
+++ b/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py
@@ -235,6 +235,11 @@ def __init__(
         self.arg_inputs: tuple[Any, ...] = tuple()
         self.kwarg_inputs: dict[str, Any] = {}
         self.additional_settings = kwargs
+        if kwargs.get("release_weights_after_build", False):
+            raise ValueError(
+                "release_weights_after_build is not supported by MutableTorchTensorRTModule, "
+                "which refits and recompiles from the source module's weights."
+            )
         self.strict = strict
         self.prefer_deferred_runtime_asserts_over_guards = (
             prefer_deferred_runtime_asserts_over_guards
diff --git a/tests/py/dynamo/models/test_release_weights_after_build.py b/tests/py/dynamo/models/test_release_weights_after_build.py
index ed6db9e2d0..77f267a422 100644
--- a/tests/py/dynamo/models/test_release_weights_after_build.py
+++ b/tests/py/dynamo/models/test_release_weights_after_build.py
@@ -63,6 +63,17 @@ def test_releaser_releases_after_the_last_reader():
     assert torch_side.untyped_storage().nbytes() > 0
 
 
+def test_storage_rebound_after_counting_is_kept():
+    shared = _param()
+    first, second = _traced(shared), _traced(shared)
+    releaser = WeightReleaser([first, second])
+    shared.data = shared.data.clone()  # e.g. moved between devices after partitioning
+    releaser.built(weight_storages(first), "first")
+    assert (
+        shared.untyped_storage().nbytes() > 0
+    )  # unknown to the releaser: second may still read it
+
+
 def test_releaser_releases_unread_parameters_only():
     used, unread = _param(), _param()
     holder = torch.nn.Module()
@@ -112,15 +123,18 @@ def __init__(self):
         )  # noqa: E731
         self.w_a, self.w_b, self.w_shared = p(16, 16), p(16, 16), p(16)
         self.w_torch, self.unused_head = p(16), p(16, 1000)
+        self.register_buffer("bias", torch.randn(16))
 
     def forward(self, x):
-        y = x @ self.w_a + self.w_shared
+        y = x @ self.w_a + self.w_shared + self.bias
         y = torch.sin(y * self.w_torch)  # aten.mul / aten.sin run in PyTorch
         return y @ self.w_b + self.w_shared
 
 
+@pytest.mark.parametrize("use_fast_partitioner", [True, False])
 @pytest.mark.parametrize("release", [True, False])
-def test_compile_releases_weights(release):
+def test_compile_releases_weights(release, use_fast_partitioner):
+    torch.manual_seed(0)
     model = _TwoEngines().eval().cuda()
     x = torch.randn(4, 16, device="cuda")
     with torch.no_grad():
@@ -131,15 +145,44 @@ def test_compile_releases_weights(release):
         arg_inputs=[x],
         min_block_size=1,
         torch_executed_ops={torch.ops.aten.sin.default, torch.ops.aten.mul.Tensor},
+        use_fast_partitioner=use_fast_partitioner,
         release_weights_after_build=release,
     )
     with torch.no_grad():
-        actual = trt_gm(x)
-    torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3)
+        for _ in range(2):  # TF32 matmuls in the engines
+            torch.testing.assert_close(trt_gm(x), expected, rtol=1e-2, atol=1e-2)
 
     state = ep.state_dict
-    for name in ("w_a", "w_b", "w_shared", "unused_head"):
+    # A read-only buffer read by an engine is a weight like any other: released after its build.
+    for name in ("w_a", "w_b", "w_shared", "unused_head", "bias"):
         assert (state[name].untyped_storage().nbytes() == 0) == release, name
     assert (
         state["w_torch"].untyped_storage().nbytes() > 0
     )  # read by the PyTorch partition
+
+
+def test_single_engine_fast_path_releases_weights():
+    model = torch.nn.Linear(16, 16).eval().cuda()
+    model.unused_head = torch.nn.Linear(16, 1000).cuda()
+    x = torch.randn(4, 16, device="cuda")
+    with torch.no_grad():
+        expected = model(x)
+    ep = torch.export.export(model, (x,))
+    trt_gm = torch_tensorrt.dynamo.compile(
+        ep,
+        arg_inputs=[x],
+        min_block_size=1,
+        require_full_compilation=True,
+        release_weights_after_build=True,
+    )
+    with torch.no_grad():
+        torch.testing.assert_close(trt_gm(x), expected, rtol=1e-3, atol=1e-3)
+    for name, tensor in ep.state_dict.items():
+        assert tensor.untyped_storage().nbytes() == 0, name
+
+
+def test_mutable_module_rejects_release():
+    with pytest.raises(ValueError, match="release_weights_after_build"):
+        torch_tensorrt.MutableTorchTensorRTModule(
+            torch.nn.Linear(4, 4).cuda(), release_weights_after_build=True
+        )

From 5d9420aec0056ab00fc897039622730adfd3097c Mon Sep 17 00:00:00 2001
From: cehongwang <wangcehong@gmail.com>
Date: Tue, 13 Oct 2026 06:01:10 +0000
Subject: [PATCH 3/4] review 2: mutated-buffer test, settings docstring, tidy

---
 py/torch_tensorrt/dynamo/_settings.py         |  3 ++-
 py/torch_tensorrt/dynamo/_weight_release.py   |  5 ++--
 .../test_release_weights_after_build.py       | 23 +++++++++++++++++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py
index 81894a5894..4263d32774 100644
--- a/py/torch_tensorrt/dynamo/_settings.py
+++ b/py/torch_tensorrt/dynamo/_settings.py
@@ -158,7 +158,8 @@ class CompilationSettings:
             that memory, so the source module (and an ``ExportedProgram`` exported from it) must not be
             used after compilation, even if compilation fails. Weights read by subgraphs that run in
             PyTorch are kept. Applies to ``torch_tensorrt.dynamo.compile``; no effect under
-            ``torch.compile`` (parameters arrive as graph inputs). Default is False.
+            ``torch.compile`` (parameters arrive as graph inputs), and not supported by
+            ``MutableTorchTensorRTModule``. Default is False.
         decompose_attention (bool): Whether to decompose attention layers into smaller operations
             instead of using the attention converters. When combined with ``use_fp32_acc=True``,
             decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only
diff --git a/py/torch_tensorrt/dynamo/_weight_release.py b/py/torch_tensorrt/dynamo/_weight_release.py
index 198fa1712c..e8a2f6b6f4 100644
--- a/py/torch_tensorrt/dynamo/_weight_release.py
+++ b/py/torch_tensorrt/dynamo/_weight_release.py
@@ -125,9 +125,8 @@ def built(self, storages: Dict[int, torch.UntypedStorage], name: str = "") -> in
                 continue  # not counted at partition time (e.g. rebound since): another reader may hold it
             self.uses[ptr] -= 1
             if self.uses[ptr] == 0:
-                del self.uses[
-                    ptr
-                ]  # a later allocation at this address is a different storage
+                # A later allocation at this address is a different storage.
+                del self.uses[ptr]
                 released += release_storage(storage)
         if released:
             logger.info(
diff --git a/tests/py/dynamo/models/test_release_weights_after_build.py b/tests/py/dynamo/models/test_release_weights_after_build.py
index 77f267a422..c583dedd01 100644
--- a/tests/py/dynamo/models/test_release_weights_after_build.py
+++ b/tests/py/dynamo/models/test_release_weights_after_build.py
@@ -181,6 +181,29 @@ def test_single_engine_fast_path_releases_weights():
         assert tensor.untyped_storage().nbytes() == 0, name
 
 
+def test_mutated_buffer_survives_release():
+    class _State(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.w = torch.nn.Parameter(torch.randn(16, 16), requires_grad=False)
+            self.register_buffer("state", torch.zeros(4, 16))
+
+        def forward(self, x):
+            y = x @ self.w + self.state
+            self.state.copy_(y)
+            return y * 2  # returning y itself trips an unrelated lowering error on main
+
+    x = torch.randn(4, 16, device="cuda")
+    ep = torch.export.export(_State().eval().cuda(), (x,))
+    trt_gm = torch_tensorrt.dynamo.compile(
+        ep, arg_inputs=[x], min_block_size=1, release_weights_after_build=True
+    )
+    assert ep.state_dict["w"].untyped_storage().nbytes() == 0
+    # Lifted into an engine input, not read through get_attr: never released.
+    buffers = list(trt_gm.buffers())
+    assert buffers and all(b.untyped_storage().nbytes() > 0 for b in buffers)
+
+
 def test_mutable_module_rejects_release():
     with pytest.raises(ValueError, match="release_weights_after_build"):
         torch_tensorrt.MutableTorchTensorRTModule(

From 71d15fc0c59347d02911db6ec1911f44bab6ea81 Mon Sep 17 00:00:00 2001
From: cehongwang <wangcehong@gmail.com>
Date: Tue, 13 Oct 2026 07:56:36 +0000
Subject: [PATCH 4/4] drop the release-handler hook (moved to #4854, its only
 user)

---
 py/torch_tensorrt/dynamo/_weight_release.py   | 21 +------------------
 .../test_release_weights_after_build.py       | 21 -------------------
 2 files changed, 1 insertion(+), 41 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/_weight_release.py b/py/torch_tensorrt/dynamo/_weight_release.py
index e8a2f6b6f4..d0062b2e40 100644
--- a/py/torch_tensorrt/dynamo/_weight_release.py
+++ b/py/torch_tensorrt/dynamo/_weight_release.py
@@ -13,29 +13,13 @@
 
 import logging
 from collections import Counter
-from typing import Callable, Dict, Iterable, List
+from typing import Dict, Iterable
 
 import torch
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
 
 logger = logging.getLogger(__name__)
 
-# A handler returns True if it released the storage (for example, by dropping the pages of a
-# memory-mapped weight); release_storage falls back to resizing the storage to zero.
-ReleaseHandler = Callable[[torch.UntypedStorage], bool]
-_release_handlers: List[ReleaseHandler] = []
-
-
-def register_release_handler(handler: ReleaseHandler) -> None:
-    """Let a storage owner release its own storages (for example, file-backed weights)."""
-    if handler not in _release_handlers:
-        _release_handlers.append(handler)
-
-
-def unregister_release_handler(handler: ReleaseHandler) -> None:
-    if handler in _release_handlers:
-        _release_handlers.remove(handler)
-
 
 def _get_attr(module: torch.nn.Module, target: str) -> object:
     obj: object = module
@@ -84,9 +68,6 @@ def release_storage(storage: torch.UntypedStorage) -> int:
     nbytes = int(storage.nbytes())
     if nbytes == 0:
         return 0
-    for handler in _release_handlers:
-        if handler(storage):
-            return nbytes
     try:
         storage.resize_(0)
     except RuntimeError as e:  # e.g. a non-resizable storage over an external buffer
diff --git a/tests/py/dynamo/models/test_release_weights_after_build.py b/tests/py/dynamo/models/test_release_weights_after_build.py
index c583dedd01..6086585067 100644
--- a/tests/py/dynamo/models/test_release_weights_after_build.py
+++ b/tests/py/dynamo/models/test_release_weights_after_build.py
@@ -8,9 +8,7 @@
 from torch_tensorrt.dynamo._weight_release import (
     WeightReleaser,
     parameter_storages,
-    register_release_handler,
     release_storage,
-    unregister_release_handler,
     weight_storages,
 )
 
@@ -87,25 +85,6 @@ def test_releaser_releases_unread_parameters_only():
     assert holder.buf.untyped_storage().nbytes() > 0  # buffers are never candidates
 
 
-def test_release_handler_takes_precedence():
-    weight = _param()
-    seen = []
-
-    def handler(storage):
-        seen.append(storage.data_ptr())
-        return True
-
-    register_release_handler(handler)
-    try:
-        assert release_storage(weight.untyped_storage()) == weight.numel() * 4
-    finally:
-        unregister_release_handler(handler)
-    assert seen == [weight.untyped_storage().data_ptr()]
-    assert (
-        weight.untyped_storage().nbytes() > 0
-    )  # the handler, not resize_, owned the release
-
-
 def test_non_resizable_storage_is_skipped():
     weight = torch.frombuffer(bytearray(64), dtype=torch.float32)
     assert release_storage(weight.untyped_storage()) == 0