fix(dynamo): preserve explicit device placement across partitions

py/torch_tensorrt/dynamo/_compiler.py

@@ -1438,7 +1438,11 @@ def preserve_module_specs(
                 str(name),
                 str(submodule.graph),
             )
-            submodule.to(to_torch_device(settings.device))
+            # Only undo the explicit compilation-time offload. Otherwise this
+            # would relocate parameters that the source graph intentionally
+            # kept on the host.
+            if settings.offload_module_to_cpu:
+                submodule.to(to_torch_device(settings.device))
             continue
 
         if name not in submodule_node_dict:

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1721,6 +1721,27 @@ def validate_dtype(to_copy_node: Node) -> bool:
 
         Based on data type being casted to
         """
+        requested_device = to_copy_node.kwargs.get("device")
+        if requested_device is not None:
+            input_node = to_copy_node.args[0]
+            input_meta = (
+                input_node.meta.get("val") if isinstance(input_node, Node) else None
+            )
+            output_meta = to_copy_node.meta.get("val")
+            if (
+                not isinstance(input_meta, torch.Tensor)
+                or not isinstance(output_meta, torch.Tensor)
+                or input_meta.device != output_meta.device
+            ):
+                _LOGGER.debug(
+                    "_to_copy converter rejected node %s because TensorRT cannot "
+                    "represent a device transfer from %s to %s",
+                    to_copy_node,
+                    getattr(input_meta, "device", None),
+                    getattr(output_meta, "device", requested_device),
+                )
+                return False
+
         allowed_casts = {
             torch.float,
             torch.int32,

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -25,6 +25,7 @@
 from torch_tensorrt.dynamo.partitioning._global_partitioner import (
     TorchTensorRTOperatorSupport,
 )
+from torch_tensorrt.dynamo.utils import to_torch_device
 
 logger = logging.getLogger(__name__)
 
@@ -39,12 +40,33 @@ def __init__(self, torch_executed_ops: Collection[Target] = set()) -> None:
         self.supported_operators: Dict[str, int] = {}
         self.unsupported_operators: Dict[str, int] = {}
         self.torch_executed_ops = torch_executed_ops
+        self._non_target_device_cache: Dict[torch.fx.Node, bool] = {}
 
     def is_node_supported(
         self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node
     ) -> bool:
         node_name = ConverterRegistry.qualified_name_or_str(node.target)
 
+        settings = CONVERTERS.compilation_settings
+        if (
+            settings is not None
+            and TorchTensorRTOperatorSupport._is_explicit_non_target_region(
+                node,
+                to_torch_device(settings.device),
+                self._non_target_device_cache,
+            )
+        ):
+            if not node.is_impure():
+                self.unsupported_operators[node_name] = (
+                    self.unsupported_operators.get(node_name, 0) + 1
+                )
+            logger.debug(
+                "Operator %s is not supported because it belongs to an explicit "
+                "non-target device region",
+                node_name,
+            )
+            return False
+
         if TorchTensorRTOperatorSupport._has_complex_dtype(node):
             # Complex-dtype tensors are not supported by TensorRT; force PyTorch fallback
             if not node.is_impure():
@@ -53,7 +75,6 @@ def is_node_supported(
                 )
             return False
 
-        settings = CONVERTERS.compilation_settings
         if (
             settings is not None
             and settings.fallback_data_dependent_ops

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -6,6 +6,7 @@
 from torch.fx.node import Target
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupport, SupportDict
+from torch.utils._pytree import tree_flatten
 from torch_tensorrt.dynamo._defaults import (
     MIN_BLOCK_SIZE,
     REQUIRE_FULL_COMPILATION,
@@ -16,7 +17,7 @@
 from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
     ConverterRegistry,
 )
-from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES
+from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES, to_torch_device
 
 logger = logging.getLogger(__name__)
 
@@ -144,6 +145,7 @@ def __init__(
         self.supported_operators: Dict[str, int] = {}
         self.unsupported_operators: Dict[str, int] = {}
         self.torch_executed_ops: Collection[Target] = torch_executed_ops
+        self._non_target_device_cache: Dict[torch.fx.Node, bool] = {}
 
     @staticmethod
     def _has_complex_dtype(node: torch.fx.Node) -> bool:
@@ -166,6 +168,60 @@ def _dtype(n: torch.fx.Node) -> Optional[torch.dtype]:
                 return True
         return False
 
+    @staticmethod
+    def _tensor_devices(node: torch.fx.Node) -> List[torch.device]:
+        leaves, _ = tree_flatten(node.meta.get("val"))
+        return [leaf.device for leaf in leaves if isinstance(leaf, torch.Tensor)]
+
+    @staticmethod
+    def _same_device(left: torch.device, right: torch.device) -> bool:
+        return left.type == right.type and (
+            left.index is None or right.index is None or left.index == right.index
+        )
+
+    @classmethod
+    def _is_explicit_non_target_region(
+        cls,
+        node: torch.fx.Node,
+        target_device: torch.device,
+        cache: Dict[torch.fx.Node, bool],
+    ) -> bool:
+        """Return True for non-target tensors produced by an explicit transfer.
+
+        Torch-TensorRT historically accepts models traced with CPU example inputs
+        and relocates their engines to the requested CUDA device. Preserve that
+        behavior while preventing a device-changing op inside an otherwise CUDA
+        graph from starting a CPU region that a later converter absorbs.
+        """
+
+        if node in cache:
+            return cache[node]
+
+        output_devices = cls._tensor_devices(node)
+        if not output_devices or all(
+            cls._same_device(device, target_device) for device in output_devices
+        ):
+            cache[node] = False
+            return False
+
+        runtime_inputs = [
+            input_node
+            for input_node in node.all_input_nodes
+            if input_node.op != "get_attr"
+        ]
+        crosses_device = any(
+            not cls._same_device(input_device, output_device)
+            for input_node in runtime_inputs
+            for input_device in cls._tensor_devices(input_node)
+            for output_device in output_devices
+        )
+        inherited_transfer = any(
+            cls._is_explicit_non_target_region(input_node, target_device, cache)
+            for input_node in runtime_inputs
+        )
+        cache[node] = crosses_device or inherited_transfer
+        return cache[node]
+
     @staticmethod
     def _requires_output_allocator(node: torch.fx.Node) -> bool:
         # True if the converter selected for this node needs a TRT output allocator,
@@ -181,6 +237,23 @@ def is_node_supported(
     ) -> bool:
         node_name = ConverterRegistry.qualified_name_or_str(node.target)
 
+        settings = CONVERTERS.compilation_settings
+        if settings is not None and self._is_explicit_non_target_region(
+            node,
+            to_torch_device(settings.device),
+            self._non_target_device_cache,
+        ):
+            if not node.is_impure():
+                self.unsupported_operators[node_name] = (
+                    self.unsupported_operators.get(node_name, 0) + 1
+                )
+            logger.debug(
+                "Operator %s is not supported because it belongs to an explicit "
+                "non-target device region",
+                node_name,
+            )
+            return False
+
         if self._has_complex_dtype(node):
             # Complex-dtype tensors are not supported by TensorRT; force PyTorch fallback
             # so the graph breaks around the complex cluster inserted by complex_graph_detection.
@@ -190,7 +263,6 @@ def is_node_supported(
                 )
             return False
 
-        settings = CONVERTERS.compilation_settings
         if (
             settings is not None
             and settings.fallback_data_dependent_ops

tests/py/dynamo/conversion/test_casts.py

@@ -1,13 +1,11 @@
 # type: ignore
 
-import unittest
-
 import torch
 import torch.nn as nn
-import torch_tensorrt
 from torch.testing._internal.common_utils import run_tests
-from torch_tensorrt import dtype
-from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException
+from torch_tensorrt.dynamo.conversion.aten_ops_converters import (
+    to_copy_dtype_validator,
+)
 
 from .harness import DispatchTestCase
 
@@ -109,6 +107,31 @@ def forward(self, x):
             precision=torch.float,
         )
 
+    def test_to_copy_validator_rejects_device_transfer(self):
+        class DeviceAndDtypeCopy(nn.Module):
+            def __init__(self, device):
+                super().__init__()
+                self.device = device
+
+            def forward(self, x):
+                y = x + 1
+                return torch.ops.aten._to_copy.default(
+                    y, device=self.device, dtype=torch.int32
+                )
+
+        def to_copy_node(device):
+            x = torch.randn(4, device="cuda")
+            exported = torch.export.export(DeviceAndDtypeCopy(device), (x,))
+            return next(
+                node
+                for node in exported.graph.nodes
+                if node.target is torch.ops.aten._to_copy.default
+            )
+
+        validator = to_copy_dtype_validator(placeholder_only=False)
+        self.assertFalse(validator(to_copy_node("cpu")))
+        self.assertTrue(validator(to_copy_node("cuda")))
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/models/test_device_placement.py

@@ -0,0 +1,117 @@
+import pytest
+import torch
+import torch_tensorrt
+from torch import nn
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.experimental.symbolic_shapes import ShapeEnv
+from torch_tensorrt.dynamo._exporter import transform
+
+
+class HostBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(8, 8))
+
+    def forward(self, x):
+        return x.to(self.weight.device) @ self.weight
+
+
+class HybridHostModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.host = HostBlock()
+
+    def forward(self, x):
+        y = x * 2.0 + 1.0
+        host = self.host(y.to("cpu"))
+        return host.to(y.device) + y
+
+
+@pytest.mark.unit
+def test_compile_preserves_host_fallback_parameter_placement():
+    model = HybridHostModel().eval()
+    x = torch.randn(8, 8, device="cuda")
+    expected = model(x)
+
+    exported = torch.export.export(model, (x,))
+    compiled = torch_tensorrt.dynamo.compile(
+        exported,
+        inputs=(x,),
+        min_block_size=1,
+        pass_through_build_failures=True,
+        torch_executed_ops={
+            "torch.ops.aten.matmul.default",
+            "torch.ops.aten._to_copy.default",
+        },
+    )
+
+    fallback_devices = [
+        parameter.device
+        for name, child in compiled.named_children()
+        if "_run_on_acc" not in name
+        for parameter in child.parameters()
+    ]
+    assert fallback_devices
+    assert all(device.type == "cpu" for device in fallback_devices)
+    torch.testing.assert_close(compiled(x), expected)
+
+    # Re-export exercises FakeTensor propagation through the hybrid boundary.
+    torch_tensorrt.dynamo.export(compiled, arg_inputs=(x,))
+
+
+class HostOutput(nn.Module):
+    def forward(self, x):
+        return (x * 2.0).to(device="cpu", dtype=torch.int32)
+
+
+@pytest.mark.unit
+def test_device_changing_to_copy_stays_outside_engine():
+    model = HostOutput().eval().cuda()
+    x = torch.randn(16, device="cuda")
+    expected = model(x)
+
+    exported = torch.export.export(model, (x,))
+    compiled = torch_tensorrt.dynamo.compile(
+        exported,
+        inputs=(x,),
+        min_block_size=1,
+        pass_through_build_failures=True,
+    )
+
+    actual = compiled(x)
+    assert actual.device.type == "cpu"
+    torch.testing.assert_close(actual, expected)
+
+    inlined = transform(compiled)
+    inlined.recompile()
+    with FakeTensorMode(shape_env=ShapeEnv()):
+        fake_output = inlined(torch.empty(16, device="cuda"))
+    assert fake_output.device.type == "cpu"
+
+
+class HostComputeAfterCopy(nn.Module):
+    def forward(self, x):
+        host = (x * 2.0).to("cpu")
+        return host + 1.0
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("use_fast_partitioner", [True, False])
+def test_cpu_compute_after_device_copy_stays_outside_engine(use_fast_partitioner):
+    model = HostComputeAfterCopy().eval().cuda()
+    x = torch.randn(16, device="cuda")
+    expected = model(x)
+
+    exported = torch.export.export(model, (x,))
+    compiled = torch_tensorrt.dynamo.compile(
+        exported,
+        inputs=(x,),
+        min_block_size=1,
+        pass_through_build_failures=True,
+        use_fast_partitioner=use_fast_partitioner,
+    )
+
+    assert any("_run_on_acc" in name for name, _ in compiled.named_children())
+    actual = compiled(x)
+    assert actual.device.type == "cpu"
+    torch.testing.assert_close(actual, expected)