diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..dc82d93c8f 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -5,7 +5,6 @@ import os import platform import warnings - from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union import sympy @@ -1263,6 +1262,9 @@ def preserve_module_specs( torch_executed_ops=settings.torch_executed_ops, require_full_compilation=settings.require_full_compilation, skip_fusion=(num_supported_ops == total_ops), + assume_full_support=( + settings.require_full_compilation and num_supported_ops == total_ops + ), ) except torch.fx.passes.splitter_base.FxNetSplitterInternalError: diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 60dec56f3b..8f91ee9606 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -18,6 +18,30 @@ logger = logging.getLogger(__name__) +# Skip installing large folded tensors that already share storage with an +# existing module attribute (typically weight permutes/views). On Flux NVFP4 +# those dominate replace cost via cpu().contiguous() under offload_module_to_cpu. +# Do NOT skip materialized folds (e.g. VLA position embedding outputs) — those +# can be required for TRT legality (int64 indices never reach the converter). +_MAX_CONSTANT_FOLD_BYTES = 1 << 20 # 1 MiB + + +def _tensor_reuses_module_storage( + gm: torch.fx.GraphModule, constant: torch.Tensor +) -> bool: + try: + ptr = constant.untyped_storage().data_ptr() + except Exception: + return False + for tensor in gm.state_dict().values(): + if isinstance(tensor, torch.Tensor): + try: + if tensor.untyped_storage().data_ptr() == ptr: + return True + except Exception: + continue + return False + @torch.utils._python_dispatch._disable_current_modes() # type: ignore def constant_fold( @@ -35,17 +59,38 @@ def constant_fold( # The constants are created on CPU to save GPU memory for TensorRT compilation. # For TRT INetwork construction the constants are moved to CPU in get_attr call. + skipped_alias = 0 for node, constant in cf.node_replacements.items(): + if isinstance(constant, torch.Tensor): + nbytes = int(constant.numel() * constant.element_size()) + if nbytes > _MAX_CONSTANT_FOLD_BYTES and _tensor_reuses_module_storage( + gm, constant + ): + skipped_alias += 1 + logger.debug( + "Skipping constant-fold install for aliased %s (%d bytes > %d)", + node.name, + nbytes, + _MAX_CONSTANT_FOLD_BYTES, + ) + continue + # Register folded values as plain tensors (buffers), matching Inductor. if settings.offload_module_to_cpu: replace_node_with_constant( gm, node, - torch.nn.Parameter(constant.cpu().contiguous(), requires_grad=False), + constant.cpu().contiguous(), ) else: - replace_node_with_constant( - gm, node, torch.nn.Parameter(constant, requires_grad=False) - ) + replace_node_with_constant(gm, node, constant) + + if skipped_alias: + logger.info( + "Skipped installing %d large aliased folded constant(s) (>%d bytes); " + "leaving original view/permute ops in the graph", + skipped_alias, + _MAX_CONSTANT_FOLD_BYTES, + ) erased_params = [] for node in gm.graph.nodes: @@ -71,7 +116,7 @@ def replace_node_with_constant( """Adapted from: https://github.com/pytorch/pytorch/blob/bcf35c6ae62bb6560befa3550e37a8283944e5f4/torch/_inductor/constant_folding.py#L17-L43 - Modified to register parameters, instead of buffers for frozen constants + Registers frozen constants as buffers (same as Inductor), not Parameters. """ g = gm.graph @@ -95,7 +140,7 @@ def replace_node_with_constant( g.erase_node(node) # Needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning - gm.register_parameter(qualname, constant) + gm.register_buffer(qualname, constant) setattr(gm, qualname, constant) diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 098e9b2685..979d2c476f 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -138,6 +138,7 @@ def __init__( require_full_compilation: bool = REQUIRE_FULL_COMPILATION, return_tuple: bool = False, skip_fusion: bool = False, + assume_full_support: bool = False, ): """ Preprocesses graph before splitting: @@ -157,20 +158,30 @@ def __init__( skip_fusion=skip_fusion, ) self.operator_support = operator_support - - # Get all accelerated nodes based on operator support conditions - self.acc_nodes = FxNetAccNodesFinder( - self.module, self.operator_support, self.settings.allow_non_tensor - )() - - if self.settings.skip_fusion: - self.fusions = {} + self.assume_full_support = assume_full_support + + if self.assume_full_support: + # The caller already walked the graph and verified converter support. + # Avoid repeating support discovery and dependency construction. + self.acc_nodes = { + node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS + } + self.fusions: Dict[torch.fx.Node, NodeSet] = {} + self.deps: Dict[torch.fx.Node, NodeSet] = {} else: - self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))() + # Get all accelerated nodes based on operator support conditions + self.acc_nodes = FxNetAccNodesFinder( + self.module, self.operator_support, self.settings.allow_non_tensor + )() - # Modify deps to add more deps for fused nodes - self.deps = self.find_deps() - self.update_deps_for_fusions() + if self.settings.skip_fusion: + self.fusions = {} + else: + self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))() + + # Modify deps to add more deps for fused nodes + self.deps = self.find_deps() + self.update_deps_for_fusions() self.non_acc_submodule_name = "_run_on_gpu_" self._node_submodule_map: Dict[str, str] = {} @@ -223,6 +234,40 @@ def partition_graph(self) -> torch.fx.GraphModule: Returns a GraphModule with submodules for each segment """ + unsupported = getattr(self.operator_support, "unsupported_operators", None) + # The explicit assumption comes from the compiler's earlier support walk. + # Otherwise, an empty dict means AccNodesFinder found no unsupported ops. + fully_supported = self.assume_full_support or ( + isinstance(unsupported, dict) and len(unsupported) == 0 + ) + + # Fast path: user demanded a single TRT engine and every op is convertible. + # Skip adjacency splitting; emit one ACC block and tag/split as usual so the + # rest of compile still sees `_run_on_acc_*`. + if self.require_full_compilation and fully_supported: + if self.settings.min_acc_module_size != MIN_BLOCK_SIZE: + logger.warning( + "Detected both require_full_compilation and min_block_size compilation " + "arguments were specified. Disregarding min_block_size argument for " + "fully supported model." + ) + nodes = [ + node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS + ] + if not nodes: + raise AssertionError( + "require_full_compilation=True was specified, but no accelerated " + "operators were found in the graph" + ) + logger.info( + "require_full_compilation + full operator support: " + "wrapping graph as a single TRT submodule (skipping adjacency split)" + ) + subgraphs = [Subgraph(is_acc=True, nodes=nodes)] + self.num_trt_accelerated_subgraphs = 1 + self.tag(subgraphs) + return self.split(remove_tag=True) + # Delegate nodes based on operator coverage subgraphs = self.put_nodes_into_subgraphs() @@ -280,6 +325,7 @@ def partition( torch_executed_ops: Collection[Target] = set(), require_full_compilation: bool = REQUIRE_FULL_COMPILATION, skip_fusion: bool = False, + assume_full_support: bool = False, ) -> Tuple[torch.fx.GraphModule, OpSupportTester]: """Partition an FX GraphModule with aten ops into TRT engines Partitioning is based on converter operator support @@ -290,6 +336,8 @@ def partition( torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage require_full_compilation: Require that all computational operators be run in TRT skip_fusion: Skip fusions found by FxNetAccFusionsFinder + assume_full_support: Skip repeated support/dependency discovery because + the caller already verified that every computational op is supported Returns: torch.fx.GraphModule, OpSupportTester """ @@ -306,6 +354,7 @@ def partition( min_block_size=min_block_size, require_full_compilation=require_full_compilation, skip_fusion=skip_fusion, + assume_full_support=assume_full_support, ) partitioned_graph = partitioner.partition_graph()