diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 8e3bdaa5cd..f659ebfe34 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -31,6 +31,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, @@ -73,6 +78,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, @@ -526,6 +532,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 @@ -635,6 +642,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 @@ -850,6 +862,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, @@ -1314,6 +1327,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, @@ -1467,13 +1488,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) @@ -1592,6 +1625,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): @@ -1689,6 +1738,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, @@ -1698,6 +1751,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: diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 44992a291d..146a8b39de 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 1f7fabfbd0..7ffb1447b0 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -52,6 +52,7 @@ OPTIMIZATION_LEVEL, PASS_THROUGH_BUILD_FAILURES, REFIT_IDENTICAL_ENGINE_WEIGHTS, + RELEASE_WEIGHTS_AFTER_BUILD, REQUIRE_FULL_COMPILATION, REUSE_CACHED_ENGINES, SPARSE_WEIGHTS, @@ -152,6 +153,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 @@ -239,6 +248,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 @@ -295,6 +305,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: ) state.setdefault("fallback_data_dependent_ops", FALLBACK_DATA_DEPENDENT_OPS) state.setdefault("fold_shared_kv_attention", FOLD_SHARED_KV_ATTENTION) + 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..d0062b2e40 --- /dev/null +++ b/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 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 new file mode 100644 index 0000000000..6086585067 --- /dev/null +++ b/tests/py/dynamo/models/test_release_weights_after_build.py @@ -0,0 +1,190 @@ +# 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, + release_storage, + 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_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() + 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_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) + self.register_buffer("bias", torch.randn(16)) + + def forward(self, x): + 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, use_fast_partitioner): + torch.manual_seed(0) + 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}, + use_fast_partitioner=use_fast_partitioner, + release_weights_after_build=release, + ) + with torch.no_grad(): + 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 + # 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_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( + torch.nn.Linear(4, 4).cuda(), release_weights_after_build=True + )