diff --git a/py/torch_tensorrt/dynamo/conversion/_conversion.py b/py/torch_tensorrt/dynamo/conversion/_conversion.py index f70b732cbd..c717f1d11d 100644 --- a/py/torch_tensorrt/dynamo/conversion/_conversion.py +++ b/py/torch_tensorrt/dynamo/conversion/_conversion.py @@ -221,7 +221,9 @@ def interpret_module_to_result( """ symbolic_shape_expressions = extract_symbolic_shape_expressions( - module, inputs=inputs + module, + inputs=inputs, + truncate_double=settings.truncate_double, ) if symbolic_shape_expressions is None: raise RuntimeError( diff --git a/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py b/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py index 11213d9d98..b8c77dabaf 100644 --- a/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py +++ b/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py @@ -19,6 +19,7 @@ def extract_symbolic_shape_expressions( module: torch.fx.GraphModule, inputs: Optional[Sequence[Input]] = None, + truncate_double: bool = False, ) -> Optional[Dict[str, List[Dict[str, Any]]]]: """ Extract symbolic shape expressions from an FX graph. @@ -32,6 +33,8 @@ def extract_symbolic_shape_expressions( name. Used as the dtype source of truth for scalar inputs, which have no dtype of their own in FX metadata. Falls back to a best-effort default when not provided. + truncate_double: Record float64 tensor bindings as float32, matching + the precision TensorRT builds when double truncation is enabled Returns: Dict with 'inputs' and 'outputs' keys, each containing a list of dicts with shape_exprs and dtype, @@ -79,7 +82,11 @@ def extract_symbolic_shape_expressions( input_info.append( { "shape_exprs": shape_exprs, - "dtype": input_val.dtype, + "dtype": ( + torch.float32 + if truncate_double and input_val.dtype == torch.float64 + else input_val.dtype + ), "name": input_node.name, } ) @@ -134,7 +141,11 @@ def extract_symbolic_shape_expressions( output_info.append( { "shape_exprs": shape_exprs, - "dtype": out_val.dtype, + "dtype": ( + torch.float32 + if truncate_double and out_val.dtype == torch.float64 + else out_val.dtype + ), } ) elif isinstance(out_val, (torch.SymInt, torch.SymFloat, int, float, bool)): diff --git a/py/torch_tensorrt/dynamo/conversion/truncate_double.py b/py/torch_tensorrt/dynamo/conversion/truncate_double.py index 51e35a7840..36fd3b3689 100644 --- a/py/torch_tensorrt/dynamo/conversion/truncate_double.py +++ b/py/torch_tensorrt/dynamo/conversion/truncate_double.py @@ -1,13 +1,13 @@ from __future__ import annotations import logging -from typing import Optional, Sequence, Set +from typing import Any, Dict, Optional, Sequence, Set import torch from torch.fx.node import _get_qualified_name from torch_tensorrt._enums import dtype from torch_tensorrt._Input import Input -from torch_tensorrt.dynamo.utils import get_torch_inputs +from torch_tensorrt.dynamo.utils import get_output_metadata, get_torch_inputs logger = logging.getLogger(__name__) @@ -40,124 +40,150 @@ def _extract_downstream_get_nodes( return get_nodes +def _metadata_dtype(metadata: Dict[str, Any]) -> Optional[torch.dtype]: + """Return the dtype of tensor metadata, ignoring scalar outputs.""" + value = metadata.get("val") + if isinstance(value, torch.Tensor): + return value.dtype + + tensor_meta = metadata.get("tensor_meta") + return getattr(tensor_meta, "dtype", None) + + +def _metadata_to_dtype( + metadata: Dict[str, Any], target_dtype: torch.dtype +) -> Dict[str, Any]: + """Copy tensor metadata while changing its dtype.""" + updated = metadata.copy() + value = updated.get("val") + if isinstance(value, torch.Tensor): + updated["val"] = value.to(target_dtype) + + tensor_meta = updated.get("tensor_meta") + if tensor_meta is not None and hasattr(tensor_meta, "_replace"): + updated["tensor_meta"] = tensor_meta._replace(dtype=target_dtype) + + return updated + + +def _find_module_node(gm: torch.fx.GraphModule, submodule_name: str) -> torch.fx.Node: + for node in gm.graph.nodes: + if node.op == "call_module" and str(node.target) == submodule_name: + return node + + raise AssertionError( + f"Sought module node {submodule_name}, could not find in graph:\n{gm.graph}" + ) + + def _repair_64bit_input( gm: torch.fx.GraphModule, position: int, submodule_name: str, - submodule_outputs: Optional[torch.Tensor | Sequence[torch.Tensor]], - dtype: torch.dtype, ) -> None: - """Fixes a single Long/Double input to a TRT-accelerated subgraph - - In-Place modifies the provided graph - - Inserts a cast to the 32-bit equivalent type for TRT, then if necessary, - inserts an upcast back to the 64-bit type for subsequent Torch operations - - Args: - gm: FX GraphModule enclosing the TRT subgraph - position: Index in the submodule inputs at which the long or double input is found - submodule_name: Name of TRT-accelerated subgraph module in FX graph - submodule_outputs: Output tensor(s) of TRT-accelerated subgraph (used for dtypes/structure) - dtype: Data type of tensor at position in submodule (double/long) - """ - assert dtype in ( - torch.float64, - ), f"dtype argument must be torch.float64, got {dtype}" - + """Downcast a single double input at a TRT boundary to float32.""" logger.info( f"Downcasting a 64-bit input at position {position} of submodule {submodule_name}" ) - # Determine target data type in 32 and 64 bit forms - dtype_64bit = dtype dtype_32bit = torch.float32 + module_node = _find_module_node(gm, submodule_name) - # Find the node representing the submodule in the graph - module_node = None - - # Iterate over all nodes in the graph, seeking target module name match - for n in gm.graph.nodes: - if n.op == "call_module" and str(n.target) == submodule_name: - module_node = n - break - - if module_node is None: - raise AssertionError( - f"Sought module node {submodule_name}, could not find in graph:\n{gm.graph}" - ) - - # Extract the 64-bit node of the input node_64bit = module_node.all_input_nodes[position] - - # Prior to the module, insert a cast to the 32-bit equivalent node with gm.graph.inserting_before(module_node): node_32bit = gm.graph.call_function( torch.ops.aten._to_copy.default, args=(node_64bit,), kwargs={"dtype": dtype_32bit}, ) + node_32bit.meta = _metadata_to_dtype(node_64bit.meta, dtype_32bit) - # Replace 64-bit input to TRT module with new 32-bit cast node module_node.replace_input_with(node_64bit, node_32bit) - output_positions_64bit = set() + gm.graph.eliminate_dead_code() + gm.graph.lint() + gm.recompile() - # Determine if any outputs of the model are 64-bit type and store their indices - if submodule_outputs is not None: - outputs_list = ( - [submodule_outputs] - if isinstance(submodule_outputs, torch.Tensor) - else submodule_outputs - ) - for output_position, output in enumerate(outputs_list): - if output.dtype == dtype_64bit: - output_positions_64bit.add(output_position) - - # Only enter this code block if there exists a 64-bit output - # This implies a cast is needed, since TRT cannot output 64-bit tensors - if output_positions_64bit: - # Determine whether the outputs of the module are tuple-type or not - is_collection_output = False - if isinstance(submodule_outputs, tuple): - is_collection_output = True - - if not is_collection_output: - # If the output is a single tensor, insert a cast back to int64 - with gm.graph.inserting_after(module_node): +def _repair_64bit_outputs( + gm: torch.fx.GraphModule, + submodule_name: str, + submodule_output_metadata: Sequence[Dict[str, Any]], + is_collection_output: bool, +) -> None: + """Correct output metadata and insert restoring casts for double outputs. + + A partition can have a float64 output even when every runtime input is + float32 -- e.g. a float32 input added to a float64 weight/constant -- so + this must run independently of whether any input needed repair. The + output dtypes come from the partition's FX metadata; compilation must not + execute the partition merely to discover information already recorded + there. + """ + dtype_64bit = torch.float64 + dtype_32bit = torch.float32 + + output_positions_64bit: Set[int] = set() + original_output_metadata = list(submodule_output_metadata) + truncated_output_metadata = [] + for output_position, metadata in enumerate(original_output_metadata): + if _metadata_dtype(metadata) == dtype_64bit: + output_positions_64bit.add(output_position) + truncated_output_metadata.append(_metadata_to_dtype(metadata, dtype_32bit)) + else: + truncated_output_metadata.append(metadata.copy()) + + if not output_positions_64bit: + return + + module_node = _find_module_node(gm, submodule_name) + + # The call_module node describes the actual engine boundary. Preserve its + # container convention while correcting tensor dtypes to what TRT emits. + for key in ("val", "tensor_meta"): + values = [ + metadata[key] for metadata in truncated_output_metadata if key in metadata + ] + if not values: + continue + current = module_node.meta.get(key) + if isinstance(current, tuple): + module_node.meta[key] = tuple(values) + elif isinstance(current, list) or len(values) > 1: + module_node.meta[key] = values + else: + module_node.meta[key] = values[0] + + if not is_collection_output: + with gm.graph.inserting_after(module_node): + cast_node_64bit = gm.graph.call_function( + torch.ops.aten._to_copy.default, + args=(module_node,), + kwargs={"dtype": dtype_64bit}, + ) + cast_node_64bit.meta = original_output_metadata[0].copy() + + module_node.replace_all_uses_with( + cast_node_64bit, delete_user_cb=lambda user: user != cast_node_64bit + ) + else: + get_nodes = _extract_downstream_get_nodes(module_node, output_positions_64bit) + for get_node in get_nodes: + output_position = get_node.args[1] + get_node.meta = truncated_output_metadata[output_position].copy() + with gm.graph.inserting_after(get_node): cast_node_64bit = gm.graph.call_function( torch.ops.aten._to_copy.default, - args=(module_node,), + args=(get_node,), kwargs={"dtype": dtype_64bit}, ) + cast_node_64bit.meta = original_output_metadata[output_position].copy() - # Replace all uses of the TRT module (except the cast node) with the 64-bit equivalent - module_node.replace_all_uses_with( - cast_node_64bit, delete_user_cb=lambda user: (user != cast_node_64bit) - ) - - else: - # If the output is a tuple of tensors, extract downstream users for each 64-bit output - get_nodes = _extract_downstream_get_nodes( - module_node, output_positions_64bit + get_node.replace_all_uses_with( + cast_node_64bit, + delete_user_cb=lambda user: user != cast_node_64bit, ) - # For each downstream user, append a cast node back to the 64-bit precision - for get_node in get_nodes: - with gm.graph.inserting_after(get_node): - cast_node_64bit = gm.graph.call_function( - torch.ops.aten._to_copy.default, - args=(get_node,), - kwargs={"dtype": torch.float64}, - ) - - get_node.replace_all_uses_with( - cast_node_64bit, - delete_user_cb=lambda user: (user != cast_node_64bit), - ) - - # Clean up graph and ensure invariants are preserved gm.graph.eliminate_dead_code() gm.graph.lint() gm.recompile() @@ -170,12 +196,12 @@ def repair_double_inputs( device: torch.device, submodule_name: Optional[str] = None, ) -> Sequence[Input]: - """Fixes all Long/Double type inputs to a TRT-accelerated subgraph + """Repair float64 inputs and outputs at a TensorRT partition boundary. In-Place modifies the provided graph - Inserts a cast to the 32-bit equivalent type for TRT, then if necessary, - inserts an upcast back to the 64-bit type for subsequent Torch operations + Casts float64 runtime inputs to float32 for TensorRT and independently + restores float64 source outputs after the engine. Args: parent_graph: FX GraphModule enclosing the TRT subgraph @@ -183,33 +209,21 @@ def repair_double_inputs( submodule_inputs: Input tensor(s) of TRT-accelerated subgraph (used for dtypes/structure) submodule_name: Optionally specify the name of the submodule target in the parent graph Returns: - New submodule inputs, updated accordingly with long/double truncation + New submodule inputs, updated for float64 truncation """ submodule_torch_inputs = get_torch_inputs(submodule_inputs, device) num_submodule_inputs = len(submodule_inputs) - repaired_outputs_once = False + name = submodule_name if submodule_name is not None else submodule._get_name() + output_node = next(node for node in submodule.graph.nodes if node.op == "output") + is_collection_output = isinstance(output_node.args[0], (tuple, list)) + submodule_output_metadata = get_output_metadata(submodule) - # For each input to the TRT subgraph, check if its type is long/double + # For each input to the TRT subgraph, check if its type is double. for position in range(num_submodule_inputs): param = submodule_torch_inputs[position] - # If the data type of the input is long/double, insert necessary - # casts to replace the operation if isinstance(param, torch.Tensor) and param.dtype == torch.float64: - # Ensure outputs are only repaired once per submodule to avoid - # unnecessary ops showing up in the graph - if not repaired_outputs_once: - submodule_outputs = submodule(*submodule_torch_inputs) - - _repair_64bit_input( - parent_graph, - position, - submodule_name if submodule_name is not None else submodule._get_name(), - None if repaired_outputs_once else submodule_outputs, - param.dtype, - ) - - repaired_outputs_once = True + _repair_64bit_input(parent_graph, position, name) # Repair submodule inputs in accordance with inserted casts dtype_32bit = torch.float32 @@ -228,4 +242,10 @@ def repair_double_inputs( submodule_torch_inputs[idx].dtype ) + # A partition can have a float64 output (e.g. from a float64 weight) even + # when no runtime input was float64, so this runs unconditionally. + _repair_64bit_outputs( + parent_graph, name, submodule_output_metadata, is_collection_output + ) + return submodule_inputs diff --git a/tests/py/dynamo/conversion/test_truncate_double.py b/tests/py/dynamo/conversion/test_truncate_double.py new file mode 100644 index 0000000000..0def67b7fa --- /dev/null +++ b/tests/py/dynamo/conversion/test_truncate_double.py @@ -0,0 +1,241 @@ +import operator +import unittest +from unittest.mock import Mock + +import torch +import torch_tensorrt +from torch import nn +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt import Input +from torch_tensorrt._enums import dtype +from torch_tensorrt.dynamo.conversion._symbolic_shape_capture import ( + extract_symbolic_shape_expressions, +) +from torch_tensorrt.dynamo.conversion.truncate_double import repair_double_inputs + + +class TestTruncateDoubleMetadata(TestCase): + def _make_graphs( + self, *, with_scalar_output: bool + ) -> tuple[torch.fx.GraphModule, torch.fx.GraphModule, torch.fx.Node]: + tensor64 = torch.empty((2, 3), dtype=torch.float64) + + subgraph = torch.fx.Graph() + subgraph_input = subgraph.placeholder("x") + subgraph_input.meta["val"] = tensor64 + tensor_output = subgraph.call_function( + torch.ops.aten.add.Tensor, args=(subgraph_input, 1.0) + ) + tensor_output.meta["val"] = tensor64 + if with_scalar_output: + scalar_output = subgraph.call_function( + torch.ops.aten.sym_size.int, args=(tensor_output, 0) + ) + scalar_output.meta["val"] = 2 + subgraph.output((tensor_output, scalar_output)) + else: + subgraph.output(tensor_output) + submodule = torch.fx.GraphModule({}, subgraph) + + root = nn.Module() + root.add_module("run_on_acc_0", submodule) + parent_graph = torch.fx.Graph() + parent_input = parent_graph.placeholder("x") + parent_input.meta["val"] = tensor64 + engine_node = parent_graph.call_module("run_on_acc_0", args=(parent_input,)) + if with_scalar_output: + engine_node.meta["val"] = [tensor64, 2] + tensor_getitem = parent_graph.call_function( + operator.getitem, args=(engine_node, 0) + ) + tensor_getitem.meta["val"] = tensor64 + scalar_getitem = parent_graph.call_function( + operator.getitem, args=(engine_node, 1) + ) + scalar_getitem.meta["val"] = 2 + parent_graph.output((tensor_getitem, scalar_getitem)) + else: + engine_node.meta["val"] = tensor64 + parent_graph.output(engine_node) + + parent = torch.fx.GraphModule(root, parent_graph) + return parent, submodule, engine_node + + def _repair( + self, parent: torch.fx.GraphModule, submodule: torch.fx.GraphModule + ) -> Input: + # Any attempt to rediscover output dtypes by executing the partition is a + # regression: parameters may be offloaded and outputs may include scalars. + submodule.forward = Mock( + side_effect=AssertionError("truncate_double executed the partition") + ) + input_spec = Input((2, 3), dtype=torch.float64) + repaired = repair_double_inputs( + parent, + submodule, + [input_spec], + torch.device("cpu"), + "run_on_acc_0", + ) + return repaired[0] + + def test_single_output_metadata_matches_engine_boundary(self): + parent, submodule, engine_node = self._make_graphs(with_scalar_output=False) + + repaired_input = self._repair(parent, submodule) + + casts = [ + node + for node in parent.graph.nodes + if node.target == torch.ops.aten._to_copy.default + ] + self.assertEqual(len(casts), 2) + input_cast = next( + node for node in casts if node.kwargs["dtype"] == torch.float32 + ) + output_cast = next( + node for node in casts if node.kwargs["dtype"] == torch.float64 + ) + self.assertIs(engine_node.args[0], input_cast) + self.assertEqual(input_cast.meta["val"].dtype, torch.float32) + self.assertEqual(engine_node.meta["val"].dtype, torch.float32) + self.assertEqual(output_cast.meta["val"].dtype, torch.float64) + self.assertEqual(repaired_input.dtype, dtype.float32) + + def test_scalar_tuple_output_is_not_executed_or_retyped(self): + parent, submodule, engine_node = self._make_graphs(with_scalar_output=True) + + self._repair(parent, submodule) + + engine_values = engine_node.meta["val"] + self.assertEqual(engine_values[0].dtype, torch.float32) + self.assertEqual(engine_values[1], 2) + + getitems = [ + node for node in parent.graph.nodes if node.target == operator.getitem + ] + tensor_getitem = next(node for node in getitems if node.args[1] == 0) + scalar_getitem = next(node for node in getitems if node.args[1] == 1) + self.assertEqual(tensor_getitem.meta["val"].dtype, torch.float32) + self.assertEqual(scalar_getitem.meta["val"], 2) + + restoring_casts = [ + node + for node in parent.graph.nodes + if node.target == torch.ops.aten._to_copy.default + and node.kwargs["dtype"] == torch.float64 + ] + self.assertEqual(len(restoring_casts), 1) + self.assertIs(restoring_casts[0].args[0], tensor_getitem) + self.assertEqual(restoring_casts[0].meta["val"].dtype, torch.float64) + + def test_repairs_float64_output_with_only_float32_inputs(self): + # A float64 weight/constant (not a runtime input) can still produce a + # float64 output -- e.g. float32_input + float64_weight promotes to + # float64 in PyTorch. Output repair must not be gated on finding a + # float64 input. + tensor32 = torch.empty((2, 3), dtype=torch.float32) + tensor64_out = torch.empty((2, 3), dtype=torch.float64) + + subgraph = torch.fx.Graph() + subgraph_input = subgraph.placeholder("x") + subgraph_input.meta["val"] = tensor32 + tensor_output = subgraph.call_function( + torch.ops.aten.add.Tensor, args=(subgraph_input, 1.0) + ) + tensor_output.meta["val"] = tensor64_out + subgraph.output(tensor_output) + submodule = torch.fx.GraphModule({}, subgraph) + + root = nn.Module() + root.add_module("run_on_acc_0", submodule) + parent_graph = torch.fx.Graph() + parent_input = parent_graph.placeholder("x") + parent_input.meta["val"] = tensor32 + engine_node = parent_graph.call_module("run_on_acc_0", args=(parent_input,)) + engine_node.meta["val"] = tensor64_out + parent_graph.output(engine_node) + parent = torch.fx.GraphModule(root, parent_graph) + + submodule.forward = Mock( + side_effect=AssertionError("truncate_double executed the partition") + ) + input_spec = Input((2, 3), dtype=torch.float32) + repaired = repair_double_inputs( + parent, submodule, [input_spec], torch.device("cpu"), "run_on_acc_0" + ) + + self.assertEqual(repaired[0].dtype, dtype.float32) + + casts = [ + node + for node in parent.graph.nodes + if node.target == torch.ops.aten._to_copy.default + ] + self.assertEqual(len(casts), 1) + self.assertEqual(casts[0].kwargs["dtype"], torch.float64) + self.assertIs(casts[0].args[0], engine_node) + self.assertEqual(engine_node.meta["val"].dtype, torch.float32) + self.assertEqual(casts[0].meta["val"].dtype, torch.float64) + + def test_symbolic_shape_metadata_uses_truncated_binding_dtype(self): + _, submodule, _ = self._make_graphs(with_scalar_output=True) + + metadata = extract_symbolic_shape_expressions(submodule, truncate_double=True) + + self.assertIsNotNone(metadata) + self.assertEqual(metadata["inputs"][0]["dtype"], torch.float32) + self.assertEqual(metadata["outputs"][0]["dtype"], torch.float32) + self.assertEqual(metadata["outputs"][1]["dtype"], torch.int64) + self.assertTrue(metadata["outputs"][1]["is_scalar"]) + + @unittest.skipIf(not torch.cuda.is_available(), "requires CUDA") + def test_float64_weight_restores_output_without_float64_input(self): + class Float64WeightModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.bias = nn.Parameter( + torch.randn((1, 8), dtype=torch.float64), + requires_grad=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.bias + + model = Float64WeightModule().eval().cuda() + input_tensor = torch.randn((4, 8), dtype=torch.float32, device="cuda") + expected = model(input_tensor) + exported = torch.export.export(model, (input_tensor,)) + + compiled = torch_tensorrt.dynamo.compile( + exported, + arg_inputs=[input_tensor], + min_block_size=1, + truncate_double=True, + pass_through_build_failures=True, + ) + actual = compiled(input_tensor) + + self.assertEqual(expected.dtype, torch.float64) + self.assertEqual(actual.dtype, expected.dtype) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + engine_node = next( + node for node in compiled.graph.nodes if node.op == "call_module" + ) + engine_values = engine_node.meta["val"] + if not isinstance(engine_values, (tuple, list)): + engine_values = [engine_values] + self.assertEqual(engine_values[0].dtype, torch.float32) + + restoring_casts = [ + node + for node in compiled.graph.nodes + if node.target == torch.ops.aten._to_copy.default + and node.kwargs.get("dtype") == torch.float64 + ] + self.assertEqual(len(restoring_casts), 1) + + +if __name__ == "__main__": + run_tests()