# Borrowed from ModelOpt AutoCast's nodeclassifier.py, modified to fit Torch-TensorRT's needs. import abc import logging import operator import re from typing import Collection, Optional import torch from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( ConverterRegistry, ) logger = logging.getLogger(__name__) class NodeRuleBase: """Base class for node classification rules. This class defines the interface for rules that determine whether a node should be kept in high precision or converted to low precision. """ @abc.abstractmethod def _check_inner(self, node): """Implement this method to check if node conversion should be skipped based on rule criteria.""" def _log_skipped(self, node, **kwargs): """Log information about skipped nodes.""" logger.info(f"Skipping node {node.name}: {self.__class__.__name__}") def check(self, node): """Check if a node should be skipped based on the rule. Args: node: The torch.fx.Node to check. Returns: bool: True if the node should be kept in high precision, False otherwise. """ result = self._check_inner(node) if result: self._log_skipped(node) return True return False class DisabledNodeNameRegexRule(NodeRuleBase): """Rule for keeping nodes with matching user-specified names in high precision.""" def __init__(self, disabled_node_name_regex): """Initialize the rule. Args: disabled_node_name_regex: List of regex patterns for user-specified node names to keep in high precision. """ self.disabled_node_name_regex = disabled_node_name_regex def _check_inner(self, node): stack = node.meta.get("nn_module_stack") try: # get the user specified name of the node node_name = stack.get(next(reversed(stack)), [""])[0] except Exception as e: raise ValueError( f"Failed to get the user specified name of the node {node.name} because {e}. Please file a bug with Torch-TensorRT." ) return any( re.match(regex, node_name) for regex in self.disabled_node_name_regex ) class DisabledOpTypes(NodeRuleBase): """Rule for keeping nodes with specific ATen ops in high precision.""" def __init__(self, excluded_ops): """Initialize the rule. Args: excluded_ops: List of ATen ops that should remain in FP32. """ self.excluded_ops = excluded_ops def _check_inner(self, node): node_name = ConverterRegistry.qualified_name_or_str(node.target) return node_name in self.excluded_ops class IORangeRule(NodeRuleBase): """Rule for keeping nodes with out-of-range inputs/outputs in high precision.""" def __init__(self, max_output_threshold, reference_data): """Initialize the rule. Args: max_output_threshold: Maximum absolute value allowed for node I/O. reference_data: Reference data for checking I/O ranges. """ self.max_output_threshold = max_output_threshold self.reference_data = reference_data self.output_data = None def _check_inner(self, node): def is_io_out_of_range(node): tensor_name = node.name if tensor_name not in self.reference_data: logger.debug( f"Node {node.name}: Tensor {tensor_name} not found in reference data. Skipping I/O range check." ) return False ref_data = self.reference_data[tensor_name] if ref_data.numel() == 0: logger.debug( f"Node {node.name}: Tensor {tensor_name} has 0 elements. Skipping I/O range check." ) return False logger.debug( f"Node {node.name}: reference data: min={ref_data.min()}, max={ref_data.max()}" ) if torch.any(torch.abs(ref_data) > self.max_output_threshold): self.output_data = ref_data return True if self.reference_data: for in_node in node.all_input_nodes: if is_io_out_of_range(in_node): return True for out_node in list(node.users): if is_io_out_of_range(out_node): return True return False def _log_skipped(self, node, **kwargs): """Log information about skipped nodes with I/O range violations.""" if self.output_data is not None: logger.info( f"Skipping node {node.name}: reference IO out of range: min={torch.min(self.output_data)}, " f"max={torch.max(self.output_data)}, range=[{-self.max_output_threshold}, {self.max_output_threshold}]" ) else: super()._log_skipped(node, **kwargs) class DepthOfReductionRule(NodeRuleBase): """ Rule for keeping nodes with high depth of reduction in high precision. This helps prevent excessive accuracy loss in operations particularly sensitive to reduced precision, as higher-depth reductions may amplify computation errors in low precision formats. Reduction ops are those that aggregate data across one or more axes, decreasing the dimensionality of the input tensor, such as convolution, gemm, etc. """ def __init__(self, max_depth_of_reduction, reference_data): """Initialize the rule. Args: max_depth_of_reduction: Maximum depth of reduction allowed in low precision. reference_data: Reference data for checking I/O ranges. """ self.max_depth_of_reduction = max_depth_of_reduction self.reference_data = reference_data self.reduction_depth = 0 def _get_tensor_shape(self, tensor_name): """Get tensor shape from reference data.""" if tensor_name in self.reference_data: return self.reference_data[tensor_name].shape return None def _log_skipped(self, node, **kwargs): """Log information about skipped nodes with depth of reduction violations.""" if self.reduction_depth > 0: logger.info( f"Skipping node {node.name}: depth of reduction {self.reduction_depth} exceeds " f"{self.max_depth_of_reduction}." ) else: super()._log_skipped(node, **kwargs) def _check_inner(self, node): # All reduction ops rely on shape of input[0] input_0_dims = ( self._get_tensor_shape(node.all_input_nodes[0].name) if len(node.all_input_nodes) > 0 else None ) if input_0_dims is None: return False self.reduction_depth = 0 if node.target in [ torch.ops.aten.scaled_dot_product_attention.default, ]: # Attention: input (batch_size, sequence_length, hidden_size) # or (batch_size, kv_num_heads, total_sequence_length, head_size) assert len(input_0_dims) == 3 or len(input_0_dims) == 4 hidden_size = ( input_0_dims[2] if len(input_0_dims) == 3 else input_0_dims[1] * input_0_dims[3] ) self.reduction_depth = hidden_size elif node.target in [ torch.ops.aten.convolution.default, torch.ops.aten.conv1d.default, torch.ops.aten.conv2d.default, torch.ops.aten.conv3d.default, ]: # Conv: input (N x C x D1 x D2 ... x Dn) # weight (out_channels, in_channels, kD1, kD2, ... kDn) # Reduction depth = in_channels * kernel_volume weight_shape = ( self._get_tensor_shape(node.all_input_nodes[1].name) if len(node.all_input_nodes) > 1 else None ) if weight_shape is None: return False in_channels = weight_shape[1] kernel_volume = torch.prod(weight_shape[2:]) self.reduction_depth = in_channels * kernel_volume elif node.target in [ torch.ops.aten.matmul, torch.ops.aten.matmul.default, torch.ops.aten.dot.default, torch.ops.aten.mm.default, torch.ops.aten.mv.default, torch.ops.aten.bmm.default, ]: # GEMM: A (M, K) @ B (K, N) = C (M, N) self.reduction_depth = input_0_dims[-1] # TODO: Add more reduction ops here return self.reduction_depth > self.max_depth_of_reduction class NodeClassifier: """Main class for classifying nodes into high and low precision groups.""" def __init__( self, nodes, excluded_nodes: Collection[str] | None = None, excluded_ops: Collection[torch.fx.node.Target] | None = None, custom_rule: NodeRuleBase | None = None, max_output_threshold: float | None = 512, max_depth_of_reduction: int | None = None, ): """Initialize the node classifier. Args: nodes: The nodes to classify. nodes_to_exclude: Collection of regex patterns for node names to keep in high precision. targets_to_exclude: Collection of targets to keep in high precision. custom_rule: Optional custom classification rule. max_output_threshold: Maximum absolute value allowed for node I/O. max_depth_of_reduction: Maximum depth of reduction allowed in low precision. """ self.nodes = nodes self.excluded_nodes = excluded_nodes self.excluded_ops = excluded_ops self.custom_rule = custom_rule self.max_output_threshold = max_output_threshold self.max_depth_of_reduction = max_depth_of_reduction def _gen_block_node_rules(self, reference_data): """Generate list of rules for blocking nodes from precision conversion. Args: reference_data: Reference data for checking I/O ranges. Returns: list[NodeRuleBase]: List of rules to apply. """ block_node_rules: list[NodeRuleBase] = [] if self.excluded_nodes: block_node_rules.append(DisabledNodeNameRegexRule(self.excluded_nodes)) if self.excluded_ops: block_node_rules.append(DisabledOpTypes(self.excluded_ops)) if reference_data: block_node_rules.append( IORangeRule(self.max_output_threshold, reference_data) ) if self.max_depth_of_reduction is not None: block_node_rules.append( DepthOfReductionRule( self.max_depth_of_reduction, reference_data, ) ) if self.custom_rule: block_node_rules.append(self.custom_rule) return block_node_rules def run( self, ref_outputs_dict: Optional[dict[str, torch.Tensor]] = None ) -> tuple[list[str], list[str]]: """Run node classification. Args: ref_outputs_dict: Optional tensors' reference data. Returns: tuple: Lists of node names (low_precision_nodes, high_precision_nodes). """ block_node_rules = self._gen_block_node_rules(ref_outputs_dict) low_precision_nodes = [] high_precision_nodes = [] for node in self.nodes: if node.op == "call_function": if ( node.target == torch.ops.higher_order.wrap_with_autocast or node.target == operator.getitem ): continue # If any condition is met - node will be executed in high precision if any(rule.check(node) for rule in block_node_rules): high_precision_nodes.append(node.name) else: low_precision_nodes.append(node.name) logger.debug(f"Low Precision Nodes: {low_precision_nodes}") logger.debug(f"High Precision Nodes: {high_precision_nodes}") return low_precision_nodes, high_precision_nodes