Skip deallocate_module when the FX module is already on CPU.

py/torch_tensorrt/dynamo/utils.py

@@ -137,11 +137,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()
@@ -993,10 +1007,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"

tests/py/dynamo/runtime/test_000_compiler_utils.py

@@ -1,8 +1,10 @@
 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,
@@ -139,5 +141,37 @@ def test_prepare_scalar_inputs(self):
         self.assertIsInstance(bool_result, torch_tensorrt.Input)
 
 
+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()