feat(dynamo): release model weights once their last engine is built

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,11 @@ 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, 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
@@ -836,6 +848,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,
@@ -1282,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,
@@ -1435,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)
@@ -1560,6 +1593,22 @@ 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:
+        _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(
+            [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))
+
     # 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 +1706,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 +1719,10 @@ def preserve_module_specs(
             )
 
             trt_modules[name] = trt_module
+            if weight_releaser is not None and weight_releaser.built(
+                submodule_weights, name
+            ):
+                release_host_and_device_memory()
 
             if _debugger_config:
                 if _debugger_config.save_engine_profile:

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

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,14 @@ 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, 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), 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
@@ -236,6 +245,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 +300,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)

py/torch_tensorrt/dynamo/_weight_release.py

@@ -0,0 +1,116 @@
+# 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 Dict, Iterable
+
+import torch
+from torch.utils._python_dispatch import is_traceable_wrapper_subclass
+
+logger = logging.getLogger(__name__)
+
+
+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 _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 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":
+            continue
+        try:
+            value = _get_attr(gm, str(node.target))
+        except AttributeError:
+            continue
+        if isinstance(value, torch.Tensor):
+            _add_storage(found, value)
+    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():
+        _add_storage(found, tensor)
+    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
+    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 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, readers: Iterable[torch.fx.GraphModule]):
+        self.uses: Counter[int] = Counter()
+        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 self.uses
+        )
+        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():
+            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:
+                # A later allocation at this address is a different storage.
+                del self.uses[ptr]
+                released += release_storage(storage)
+        if released:
+            logger.info(
+                f"Released {released / 2**20:.0f} MiB of weights after building {name}"
+            )
+        return released

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