# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- from logging import getLogger from typing import Tuple from fusion_attention import AttentionMask, FusionAttention from fusion_options import AttentionMaskFormat from onnx import NodeProto from onnx_model import OnnxModel logger = getLogger(__name__) class FusionAttentionClip(FusionAttention): """ Fuse Attention subgraph of Clip into one Attention node. """ def __init__( self, model: OnnxModel, hidden_size: int, num_heads: int, ): attention_mask = AttentionMask(model) attention_mask.mask_format = AttentionMaskFormat.NoMask super().__init__( model, hidden_size, num_heads, attention_mask, use_multi_head_attention=False, search_op_types=["SkipLayerNormalization"], ) def get_num_heads_and_hidden_size(self, reshape_q: NodeProto) -> Tuple[int, int]: """Detect num_heads and hidden_size for ONNX model from MiDaS Args: reshape_q (NodeProto): reshape node for q Returns: Tuple[int, int]: num_heads and hidden_size """ concat = self.model.match_parent(reshape_q, "Concat", 1) if concat is None or len(concat.input) != 4: return self.num_heads, self.hidden_size # The shape is a tensor like [?, ?, num_heads, head_size] num_head_value = self.model.get_constant_value(concat.input[2]) if num_head_value is None: return self.num_heads, self.hidden_size # Fall back to user specified value if len(num_head_value) != 1 or num_head_value[0] <= 0: return self.num_heads, self.hidden_size # Fall back to user specified value num_heads = num_head_value[0] head_size_value = self.model.get_constant_value(concat.input[3]) if head_size_value is None: return self.num_heads, self.hidden_size # Fall back to user specified value if len(head_size_value) != 1 or head_size_value[0] <= 0: return self.num_heads, self.hidden_size # Fall back to user specified value head_size = head_size_value[0] hidden_size = num_heads * head_size if self.num_heads > 0 and num_heads != self.num_heads: if self.num_heads_warning: logger.warning(f"--num_heads is {self.num_heads}. Detected value is {num_heads}. Using detected value.") self.num_heads_warning = False # Do not show the warning more than once if self.hidden_size > 0 and hidden_size != self.hidden_size: if self.hidden_size_warning: logger.warning( f"--hidden_size is {self.hidden_size}. Detected value is {hidden_size}. Using detected value." ) self.hidden_size_warning = False # Do not show the warning more than once return num_heads, hidden_size def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node): skip_input_index = None node_before_layer_norm = None for i in [1, 0]: parent = self.model.match_parent(normalize_node, "SkipLayerNormalization", i) if parent is not None: skip_input_index = i node_before_layer_norm = parent root_input = None if node_before_layer_norm is not None: root_input = node_before_layer_norm.output[0] else: # Deal with the first attention after the embedding layer. for i in [0, 1]: node_before_layer_norm = None node_before_layer_norm_1 = self.model.match_parent(normalize_node, "Add", i) node_before_layer_norm_2 = self.model.match_parent(normalize_node, "LayerNormalization", i) if node_before_layer_norm_1 is not None: # Add -----------+ # | | # LayerNorm | # | | # LayerNorm | # | | # Attention subgraph | # | | # SkipLayerNorm ------+ node_before_layer_norm = node_before_layer_norm_1 elif node_before_layer_norm_2 is not None: # Add # | # LayerNorm --------+ # | | # LayerNorm | # | | # Attention subgraph | # | | # SkipLayerNorm ------+ node_before_layer_norm = node_before_layer_norm_2 if node_before_layer_norm is None: continue child = self.model.find_first_child_by_type( node_before_layer_norm, "LayerNormalization", input_name_to_nodes, False ) if child is None: continue root_input = child.output[0] skip_input_index = i break if skip_input_index is None: return qkv_nodes = self.model.match_parent_path( normalize_node, ["Add", "MatMul", "Reshape", "Transpose", "Reshape", "MatMul"], [1 - skip_input_index, None, None, 0, 0, 0], ) if qkv_nodes is None: return (_, _, reshape_qkv, transpose_qkv, _, matmul_qkv) = qkv_nodes v_nodes = self.model.match_parent_path( matmul_qkv, ["Reshape", "Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0, None] ) if v_nodes is None: logger.debug("fuse_attention: failed to match v path") return (_, _, reshape_v, add_v, matmul_v) = v_nodes add_mask = None add_mask_indices = [] qk_nodes = None qk_nodes_1 = self.model.match_parent_path( matmul_qkv, ["Softmax", "Reshape", "Add", "Reshape", "MatMul"], [0, 0, 0, None, 0], return_indice=add_mask_indices, ) qk_nodes_2 = self.model.match_parent_path( matmul_qkv, ["Softmax", "MatMul"], [0, 0], ) if qk_nodes_1 is not None: qk_nodes = qk_nodes_1 assert len(add_mask_indices) == 1 causal_mask_input_index = 1 - add_mask_indices[0] (_softmax_qk, _, add_mask, _, matmul_qk) = qk_nodes elif qk_nodes_2 is not None: qk_nodes = qk_nodes_2 (_softmax_qk, matmul_qk) = qk_nodes else: logger.debug("fuse_attention: failed to match qk path") return q_nodes = self.model.match_parent_path( matmul_qk, ["Reshape", "Transpose", "Reshape", "Mul", "Add", "MatMul"], [0, 0, 0, 0, None, None] ) if q_nodes is None: logger.debug("fuse_attention: failed to match q path") return (_, _transpose_q, reshape_q, mul_q, add_q, matmul_q) = q_nodes k_nodes = self.model.match_parent_path( matmul_qk, ["Transpose", "Reshape", "Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0, 0, None] ) if k_nodes is None: logger.debug("fuse_attention: failed to match k path") return (_transpose_k, _reshape_k, _, _, add_k, matmul_k) = k_nodes if matmul_q.input[0] != root_input or matmul_k.input[0] != root_input or matmul_v.input[0] != root_input: logger.debug("fuse_attention: expect to have same input to q, k and v matmul") return num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q) if num_heads <= 0 or hidden_size <= 0: logger.debug("fuse_attention: failed to detect num_heads or hidden_size") return attention_last_node = reshape_qkv if add_mask is not None: # Here we do not match the whole subgraph since it is very complex. Instead, we just check whether a key path # of computing causal mask. causal_mask_nodes = self.model.match_parent_path( add_mask, ["Concat", "Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"], [causal_mask_input_index, 0, 0, 0, 0, 0], ) if causal_mask_nodes is None: # If the model is exported with batch_size == 1, there is no Concat node causal_mask_nodes = self.model.match_parent_path( add_mask, ["Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"], [causal_mask_input_index, 0, 0, 0, 0], ) if causal_mask_nodes is None: logger.debug("fuse_attention: failed to match causal mask subgraph") return new_node = self.create_attention_node( mask_index=None, q_matmul=matmul_q, k_matmul=matmul_k, v_matmul=matmul_v, q_add=add_q, k_add=add_k, v_add=add_v, num_heads=num_heads, hidden_size=hidden_size, input=root_input, output=attention_last_node.output[0], add_qk_str=None, scale=None, causal=(add_mask is not None), ) if new_node is None: return self.nodes_to_add.append(new_node) self.node_name_to_graph_name[new_node.name] = self.this_graph_name self.nodes_to_remove.extend([attention_last_node, transpose_qkv]) # Use prune graph to remove nodes since they are shared by all attention nodes. self.prune_graph = True