# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- from logging import getLogger import numpy as np from fusion_gpt_attention import FusionGptAttentionPastBase from onnx import helper from onnx_model import OnnxModel logger = getLogger(__name__) def is_close(value, expected_value): return abs(value - expected_value) <= 1e-6 class FusionGptAttentionMegatron(FusionGptAttentionPastBase): """ Fuse GPT-2 Attention with past state subgraph from Megatron into one Attention node. """ def __init__(self, model: OnnxModel, num_heads: int): super().__init__(model, num_heads) def fuse_attention_node( self, matmul_before_split, add_before_split, past, present, input, reshape_qkv, mask, ): attention_node_name = self.model.create_node_name("GptAttention") int32_mask = self.cast_attention_mask(mask) output = reshape_qkv.output[0] i = 1 if (add_before_split.input[0] == matmul_before_split.output[0]) else 0 attention_node = helper.make_node( "Attention", inputs=[ input, matmul_before_split.input[1], add_before_split.input[i], int32_mask, past, ], outputs=[output, present], name=attention_node_name, ) attention_node.domain = "com.microsoft" attention_node.attribute.extend( [ helper.make_attribute("num_heads", self.num_heads), helper.make_attribute("unidirectional", 0), # unidirectional shall not be ON for 4D attention mask ] ) if self.mask_filter_value is not None: attention_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))]) nodes_to_add = [attention_node] self.nodes_to_add.extend(nodes_to_add) for node in nodes_to_add: self.node_name_to_graph_name[node.name] = self.this_graph_name self.nodes_to_remove.append(reshape_qkv) # we rely on prune_graph() to clean old subgraph nodes self.prune_graph = True def match_mask(self, sub_qk, mul_qk, matmul_qk, layernorm_before_attention): mask_nodes = self.model.match_parent_path(sub_qk, ["Mul", "Sub", "Slice", "Slice"], [1, 0, 1, 0]) if mask_nodes is None: logger.debug("fuse_attention: failed to match unidirectional mask path") return None (mul_mask, sub_mask, last_slice_mask, slice_mask) = mask_nodes if len(mask_nodes) > 1 and mask_nodes[0].op_type == "Mul": _, mul_val = self.model.get_constant_input(mask_nodes[0]) if mul_val != 10000: self.mask_filter_value = -mul_val if mul_qk.input[1] != last_slice_mask.output[0]: logger.debug("fuse_attention failed: mul_qk.input[1] != last_slice_mask.output[0]") return None if not self.utils.check_node_input_value(mul_mask, 1, 10000.0): logger.debug("fuse_attention failed: mul_mask input 1 is not constant 10000.0") return None if not self.utils.check_node_input_value(sub_mask, 0, 1.0): logger.debug("fuse_attention failed: sub_mask input 0 is not constant 1.0") return None if not self.model.find_graph_input(slice_mask.input[0]): logger.info("expect slick_mask input 0 to be graph input") return None if not self.utils.check_node_input_value(last_slice_mask, 1, [0]): logger.debug("fuse_attention failed: last_slice_mask input 1 (starts) is not constant [0]") return None if not self.utils.check_node_input_value(last_slice_mask, 3, [3]): logger.debug("fuse_attention failed: last_slice_mask input 3 (axes) is not constant [3]") return False if not self.utils.check_node_input_value(last_slice_mask, 4, [1]): logger.debug("fuse_attention failed: last_slice_mask input 4 (steps) is not constant [1]") return False if not self.utils.check_node_input_value(slice_mask, 3, [2]): logger.debug("fuse_attention failed: slice_mask input 3 (axes) is not constant [2]") return None if not self.utils.check_node_input_value(slice_mask, 4, [1]): logger.debug("fuse_attention failed: slice_mask input 4 (steps) is not constant [1]") return None last_slice_path = self.model.match_parent_path( last_slice_mask, ["Unsqueeze", "Gather", "Shape", "MatMul"], [2, 0, 0, 0] ) if last_slice_path is None or last_slice_path[-1] != matmul_qk: logger.debug("fuse_attention: failed to match last slice path") return None first_slice_path = self.model.match_parent_path( slice_mask, ["Unsqueeze", "Gather", "Shape", "MatMul"], [2, 0, 0, 0] ) if first_slice_path is None or first_slice_path[-1] != matmul_qk: logger.debug("fuse_attention: failed to match first slice path") return None first_slice_sub = self.model.match_parent_path( slice_mask, ["Unsqueeze", "Sub", "Gather", "Shape", "MatMul"], [1, 0, 0, 0, 0], ) if first_slice_sub is None or first_slice_sub[-1] != matmul_qk: logger.debug("fuse_attention: failed to match last slice sub path") return None first_slice_sub_1 = self.model.match_parent_path( slice_mask, ["Unsqueeze", "Sub", "Gather", "Shape", "LayerNormalization"], [1, 0, 1, 0, 0], ) if first_slice_sub_1 is None: first_slice_sub_1 = self.model.match_parent_path( slice_mask, ["Unsqueeze", "Sub", "Gather", "Shape", "SkipLayerNormalization"], [1, 0, 1, 0, 0], ) if first_slice_sub_1 is None or first_slice_sub_1[-1] != layernorm_before_attention: logger.debug("fuse_attention: failed to match last slice sub path 1") return None return slice_mask.input[0] def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node): past = None present = None is_normalize_node_skiplayernorm = normalize_node.op_type == "SkipLayerNormalization" qkv_nodes = None if not is_normalize_node_skiplayernorm: qkv_nodes = self.model.match_parent_path( normalize_node, ["Add", "Add", "MatMul", "Reshape", "Transpose", "MatMul"], [0, 1, None, 0, 0, 0], output_name_to_node=output_name_to_node, ) else: qkv_nodes = self.model.match_parent_path( normalize_node, ["Add", "MatMul", "Reshape", "Transpose", "MatMul"], [1, None, 0, 0, 0], output_name_to_node=output_name_to_node, ) if qkv_nodes is None: return skip_input = None if not is_normalize_node_skiplayernorm: ( add_skip, add_after_attention, matmul_after_attention, reshape_qkv, transpose_qkv, matmul_qkv, ) = qkv_nodes skip_input = add_skip.input[0] else: ( add_after_attention, matmul_after_attention, reshape_qkv, transpose_qkv, matmul_qkv, ) = qkv_nodes skip_input = normalize_node.input[0] v_nodes = self.model.match_parent_path( matmul_qkv, [ "Concat", "Transpose", "Reshape", "Split", "Add", "MatMul", "LayerNormalization", ], [1, 1, 0, 0, 0, None, 0], ) if v_nodes is None: v_nodes = self.model.match_parent_path( matmul_qkv, [ "Concat", "Transpose", "Reshape", "Split", "Add", "MatMul", "SkipLayerNormalization", ], [1, 1, 0, 0, 0, None, 0], ) if v_nodes is None: logger.debug("fuse_attention: failed to match v path") return ( concat_v, transpose_v, reshape_v, split_v, add_before_split, matmul_before_split, layernorm_before_attention, ) = v_nodes if ( layernorm_before_attention.op_type == "LayerNormalization" and skip_input != layernorm_before_attention.input[0] ): logger.debug("fuse_attention: skip_input != layernorm_before_attention.input[0]") return if ( layernorm_before_attention.op_type == "SkipLayerNormalization" and skip_input != layernorm_before_attention.output[3] ): logger.debug("fuse_attention: skip_input != layernorm_before_attention.input[0]") return qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Sub", "Mul", "MatMul"], [0, 0, 0, 0]) if qk_nodes is None: logger.debug("fuse_attention: failed to match qk path") return None (softmax_qk, sub_qk, mul_qk, matmul_qk) = qk_nodes if self.model.get_node_attribute(softmax_qk, "axis") != 3: logger.debug("fuse_attention failed: softmax_qk axis != 3") return None attention_mask = self.match_mask(sub_qk, mul_qk, matmul_qk, layernorm_before_attention) q_nodes = self.model.match_parent_path(matmul_qk, ["Div", "Transpose", "Reshape", "Split"], [0, 0, 0, 0]) if q_nodes is None: logger.debug("fuse_attention: failed to match q path") return (div_q, transpose_q, reshape_q, split_q) = q_nodes if split_v != split_q: logger.debug("fuse_attention: skip since split_v != split_q") return k_nodes = self.model.match_parent_path( matmul_qk, ["Div", "Transpose", "Concat", "Transpose", "Reshape", "Split"], [1, 0, 0, 1, 0, 0], ) if k_nodes is None: logger.debug("fuse_attention: failed to match k path") return (div_k, _, concat_k, transpose_k, reshape_k, split_k) = k_nodes if split_v != split_k: logger.debug("fuse_attention: skip since split_v != split_k") return i, value = self.model.get_constant_input(reshape_k) if not ( isinstance(value, np.ndarray) and list(value.shape) == [4] and value[0] == 0 and value[1] == 0 and value[2] > 0 and value[3] > 0 ): logger.debug("fuse_attention: reshape constant input is not [0, 0, N, H]") return num_heads = value[2] if num_heads != self.num_heads: logger.info(f"Detected num_heads={num_heads}. Ignore user specified value {self.num_heads}") self.num_heads = num_heads hidden_size_per_head = value[3] i, value = self.model.get_constant_input(div_k) expected_value = float(np.sqrt(np.sqrt(hidden_size_per_head))) if not is_close(value, expected_value): logger.debug(f"fuse_attention: div_k value={value} expected={expected_value}") return i, value = self.model.get_constant_input(div_q) if not is_close(value, expected_value): logger.debug(f"fuse_attention: div_q value={value} expected={expected_value}") return # Match past and present paths past = self.match_past_pattern_2(concat_k, concat_v, output_name_to_node) if past is None: logger.debug("fuse_attention: match past failed") return if not self.model.find_graph_input(past): logger.debug("fuse_attention: past is not graph input.") # For GPT2LMHeadModel_BeamSearchStep, there is an extra Gather node to select beam index so it is not graph input. present = self.match_present(concat_v, input_name_to_nodes) if present is None: logger.debug("fuse_attention: match present failed") return if not self.model.find_graph_output(present): logger.info("fuse_attention: expect present to be graph output") return self.fuse_attention_node( matmul_before_split, add_before_split, past, present, layernorm_before_attention.output[0], reshape_qkv, attention_mask, )