Skip to content

fix(dynamo): make truncate_double metadata match engine precision #4630

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
narendasan wants to merge 1 commit into from
Open

fix(dynamo): make truncate_double metadata match engine precision #4630

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion py/torch_tensorrt/dynamo/conversion/_conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,9 @@ def interpret_module_to_result(
SerializedInterpreterResult
"""

symbolic_shape_expressions = extract_symbolic_shape_expressions(module)
symbolic_shape_expressions = extract_symbolic_shape_expressions(
module, truncate_double=settings.truncate_double
)
if symbolic_shape_expressions is None:
raise RuntimeError(
"Failed to extract symbolic shape expressions from source FX graph partition"
Expand Down
15 changes: 13 additions & 2 deletions py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

def extract_symbolic_shape_expressions(
module: torch.fx.GraphModule,
truncate_double: bool = False,
) -> Optional[Dict[str, List[Dict[str, Any]]]]:
"""
Extract symbolic shape expressions from an FX graph.
Expand All @@ -25,6 +26,8 @@ def extract_symbolic_shape_expressions(

Args:
module: FX GraphModule with symbolic shapes in node metadata
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,
Expand Down Expand Up @@ -64,7 +67,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,
}
)
Expand Down Expand Up @@ -115,7 +122,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)):
Expand Down
149 changes: 81 additions & 68 deletions py/torch_tensorrt/dynamo/conversion/truncate_double.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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__)

Expand Down Expand Up @@ -40,124 +40,140 @@ 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 _repair_64bit_input(
gm: torch.fx.GraphModule,
position: int,
submodule_name: str,
submodule_outputs: Optional[torch.Tensor | Sequence[torch.Tensor]],
submodule_output_metadata: Optional[Sequence[Dict[str, Any]]],
is_collection_output: bool,
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
"""Fix a single double input and any double outputs at a TRT boundary.

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)
The output dtypes come from the partition's FX metadata. Compilation must not
execute the partition merely to discover information already recorded there.
"""
assert dtype in (
torch.float64,
), f"dtype argument must be torch.float64, got {dtype}"
assert dtype == torch.float64, f"dtype argument must be torch.float64, got {dtype}"

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

# 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
for node in gm.graph.nodes:
if node.op == "call_module" and str(node.target) == submodule_name:
module_node = node
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()

# 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)
output_positions_64bit: Set[int] = set()
original_output_metadata = list(submodule_output_metadata or ())
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())

# The call_module node describes the actual engine boundary. Preserve its
# container convention while correcting tensor dtypes to what TRT emits.
if truncated_output_metadata:
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]

# 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):
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()

# 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)
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
)

# For each downstream user, append a cast node back to the 64-bit precision
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=(get_node,),
kwargs={"dtype": torch.float64},
kwargs={"dtype": dtype_64bit},
)
cast_node_64bit.meta = original_output_metadata[
output_position
].copy()

get_node.replace_all_uses_with(
cast_node_64bit,
delete_user_cb=lambda user: (user != 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()
Expand Down Expand Up @@ -188,24 +204,21 @@ def repair_double_inputs(
submodule_torch_inputs = get_torch_inputs(submodule_inputs, device)
num_submodule_inputs = len(submodule_inputs)
repaired_outputs_once = False
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,
None if repaired_outputs_once else submodule_output_metadata,
is_collection_output,
param.dtype,
)

Expand Down
Loading
Loading