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)