diff --git a/py/torch_tensorrt/dynamo/utils.py b/py/torch_tensorrt/dynamo/utils.py index df984d46e5..e685299d21 100644 --- a/py/torch_tensorrt/dynamo/utils.py +++ b/py/torch_tensorrt/dynamo/utils.py @@ -141,11 +141,25 @@ def unified_dtype_converter( raise TypeError("%s is not a supported dtype" % dtype) +def _module_occupies_cuda(module: torch.nn.Module) -> bool: + """True if any parameter or buffer still lives on CUDA.""" + for tensor in module.parameters(): + if tensor.is_cuda: + return True + for tensor in module.buffers(): + if tensor.is_cuda: + return True + return False + + def deallocate_module(module: torch.fx.GraphModule) -> None: + """Move the FX module to CPU and free cached CUDA blocks for the TRT builder. + + No-op when nothing is on CUDA, so a second call after compile() already + offloaded (and CPU-only compiles) skip ``to("cpu")`` / ``empty_cache`` / ``gc``. """ - This is a helper function to delete the instance of module. We first move it to CPU and then - delete the object. This function ensures the GPU memory occupied by the module is released effectively after this call - """ + if not torch.cuda.is_available() or not _module_occupies_cuda(module): + return module.to(CPU_DEVICE) torch.cuda.empty_cache() gc.collect() @@ -1002,10 +1016,11 @@ def get_cpu_memory_usage() -> Any: def release_host_and_device_memory() -> None: gc.collect() if torch.cuda.is_available(): + # One sync so builder work has finished before empty_cache; a second + # sync after ipc_collect does not free extra blocks. torch.cuda.synchronize() torch.cuda.empty_cache() torch.cuda.ipc_collect() - torch.cuda.synchronize() if ( platform.system() == "Linux" diff --git a/tests/py/dynamo/runtime/test_000_compiler_utils.py b/tests/py/dynamo/runtime/test_000_compiler_utils.py index abbc1db882..da375a4f3b 100644 --- a/tests/py/dynamo/runtime/test_000_compiler_utils.py +++ b/tests/py/dynamo/runtime/test_000_compiler_utils.py @@ -2,10 +2,12 @@ # SPDX-License-Identifier: BSD-3-Clause import unittest +from unittest import mock import torch import torch_tensorrt from torch_tensorrt.dynamo.utils import ( + deallocate_module, get_torch_tensor, prepare_inputs, to_torch_device, @@ -190,5 +192,37 @@ def forward(self, x): torch_tensorrt.save(M(), arg_inputs=(), inputs=(torch.randn(2, 3),)) +class _TinyLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + +class TestDeallocateModule(unittest.TestCase): + def test_cpu_module_skips_empty_cache(self) -> None: + gm = torch.fx.symbolic_trace(_TinyLinear().eval()) + with mock.patch("torch.cuda.empty_cache") as empty_cache: + deallocate_module(gm) + empty_cache.assert_not_called() + self.assertEqual(next(gm.parameters()).device.type, "cpu") + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_cuda_module_moves_to_cpu(self) -> None: + gm = torch.fx.symbolic_trace(_TinyLinear().eval().cuda()) + deallocate_module(gm) + self.assertEqual(next(gm.parameters()).device.type, "cpu") + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_second_call_is_noop(self) -> None: + gm = torch.fx.symbolic_trace(_TinyLinear().eval().cuda()) + deallocate_module(gm) + with mock.patch("torch.cuda.empty_cache") as empty_cache: + deallocate_module(gm) + empty_cache.assert_not_called() + + if __name__ == "__main__": unittest.main()