From 292d096b259c206610f858cf8844360ac50c263d Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 21 Aug 2026 18:28:17 +0000 Subject: [PATCH] feedback --- .../lowering/passes/complex_graph_rewrite.py | 31 +++++- .../dynamo/lowering/test_complex_rewrite.py | 96 ++++++++++++++++++- 2 files changed, 124 insertions(+), 3 deletions(-) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py index c8b7ad8571..72e097ef4a 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py @@ -1833,7 +1833,36 @@ def propagate_metadata( if existing_fake_mode is not None else FakeTensorMode(allow_non_fake_inputs=True) ) - FakeTensorProp(self.gm, mode=prop_mode).propagate(*fake_inputs) + + # offload_module_to_cpu leaves a folded constant on CPU while its node + # metadata still records cuda, which FakeTensorProp would then mix with + # fake cuda activations. Realign each attr to the device its own meta + # claims, so genuinely-CPU attrs are left alone. + attr_devices: Dict[str, torch.device] = {} + for node in self.gm.graph.nodes: + if node.op != "get_attr": + continue + attr_val = node.meta.get("val", None) + if isinstance(attr_val, torch.Tensor): + attr_devices[str(node.target)] = attr_val.device + + if not attr_devices: + FakeTensorProp(self.gm, mode=prop_mode).propagate(*fake_inputs) + return + + class _DeviceAligningFakeTensorProp(FakeTensorProp): # type: ignore[misc] + def fetch_attr(self, target: str) -> Any: + attr = super().fetch_attr(target) + device = attr_devices.get(target) + if ( + device is not None + and isinstance(attr, torch.Tensor) + and attr.device != device + ): + return attr.to(device) + return attr + + _DeviceAligningFakeTensorProp(self.gm, mode=prop_mode).propagate(*fake_inputs) def extract_real_imag(input, placeholder_or_func: bool = True): # type: ignore diff --git a/tests/py/dynamo/lowering/test_complex_rewrite.py b/tests/py/dynamo/lowering/test_complex_rewrite.py index bae2a6bfcd..9912067831 100644 --- a/tests/py/dynamo/lowering/test_complex_rewrite.py +++ b/tests/py/dynamo/lowering/test_complex_rewrite.py @@ -38,6 +38,7 @@ from torch_tensorrt.dynamo.lowering.passes.complex_graph_rewrite import ( complex_graph_detection, ) +from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold # --------------------------------------------------------------------------- # 1. Infrastructure @@ -1251,7 +1252,9 @@ def forward(self, z, opt): @pytest.mark.unit -@pytest.mark.parametrize("scale", [torch.tensor(2.0), 2, True], ids=["tensor", "int", "bool"]) +@pytest.mark.parametrize( + "scale", [torch.tensor(2.0), 2, True], ids=["tensor", "int", "bool"] +) def test_non_tensor_scalar_placeholder(scale): class RotaryComplex(nn.Module): def forward(self, xq, freqs_cis, scale): @@ -1261,4 +1264,93 @@ def forward(self, xq, freqs_cis, scale): xq = torch.randn(1, 2, 4, 8) freqs = torch.polar(torch.ones(1, 2, 4, 4), torch.randn(1, 2, 4, 4)) - _export_and_lower(RotaryComplex(), (xq, freqs, scale)) + gm = _export_and_lower(RotaryComplex(), (xq, freqs, scale)) + + placeholder_vals = [ + node.meta["val"] + for node in gm.graph.nodes + if node.op == "placeholder" and "val" in node.meta + ] + if isinstance(scale, torch.Tensor): + # A 0-dim tensor stays a tensor: smoke test only, not guard coverage + assert placeholder_vals, "export produced no placeholder metadata" + else: + # The guard only runs for non-tensor vals; fail if export specialized it away + kinds = [type(val).__name__ for val in placeholder_vals] + assert any(not isinstance(val, torch.Tensor) for val in placeholder_vals), ( + f"scale={scale!r} did not survive export as a non-tensor " + f"placeholder; placeholder vals were {kinds}" + ) + + +@pytest.mark.unit +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_cpu_offloaded_frozen_param_is_aligned(): + """A CPU-offloaded folded constant keeping cuda metadata must be realigned.""" + + class M(nn.Module): + def __init__(self) -> None: + super().__init__() + # Non-scalar, so scalar promotion cannot excuse a device mismatch + self.register_buffer("w", torch.randn(3, 4)) + + def forward(self, z): + return (z * z) * (self.w * 2) + + model = M().cuda().eval() + z = torch.randn(3, 4, dtype=torch.complex64, device="cuda") + with torch.no_grad(): + exp = torch.export.export(model, (z,)) + gm = exp.module() + + # Offloading stores the constant on CPU; its inherited meta still says cuda + gm = constant_fold(gm, CompilationSettings(offload_module_to_cpu=True)) + + offloaded = [ + node + for node in gm.graph.nodes + if node.op == "get_attr" + and str(node.target).startswith("_frozen_param") + and isinstance(node.meta.get("val"), torch.Tensor) + and node.meta["val"].device.type == "cuda" + and getattr(gm, str(node.target)).device.type == "cpu" + ] + assert offloaded, "expected a CPU-offloaded _frozen_param with cuda metadata" + + complex_graph_detection(gm, CompilationSettings()) + + for node in offloaded: + assert ( + getattr(gm, str(node.target)).device.type == "cpu" + ), f"{node.target}: propagation must not move the module's parameters" + + +@pytest.mark.unit +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_independent_cpu_branch_stays_on_cpu(): + """An unrelated CPU attribute must stay on CPU despite a cuda placeholder.""" + + class M(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("cpu_table", torch.arange(4)) + + def forward(self, z): + return z * z, self.cpu_table + 1 + + # Module stays on CPU so the table is genuinely CPU, not offloaded + z = torch.randn(3, 4, dtype=torch.complex64, device="cuda") + gm = _export_and_lower(M().eval(), (z,)) + + cpu_branch = [ + node + for node in gm.graph.nodes + if isinstance(node.meta.get("val"), torch.Tensor) + and node.meta["val"].dtype == torch.int64 + and tuple(node.meta["val"].shape) == (4,) + ] + assert cpu_branch, "the independent CPU branch disappeared from the graph" + for node in cpu_branch: + assert ( + node.meta["val"].device.type == "cpu" + ), f"{node.name}: CPU branch moved to {node.meta['val'].device}"