mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-05 04:17:53 +00:00
70b8cda979
2nd round of cherry pick LLaMA related changes to 1.16.2 release. --------- Co-authored-by: aciddelgado <139922440+aciddelgado@users.noreply.github.com> Co-authored-by: Frank Dong <123416088+frank-dong-ms@users.noreply.github.com>
1376 lines
56 KiB
Python
1376 lines
56 KiB
Python
# -------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
# --------------------------------------------------------------------------
|
|
import logging
|
|
from typing import Optional, Union
|
|
|
|
from fusion_attention import FusionAttention
|
|
from fusion_base import Fusion
|
|
from onnx import FunctionProto, NodeProto, TensorProto, helper, numpy_helper
|
|
from onnx_model import OnnxModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FusionRotaryAttention(FusionAttention):
|
|
"""
|
|
Fuse Attention subgraph with rotary positional embeddings into one MultiHeadAttention node.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: OnnxModel,
|
|
hidden_size: int,
|
|
num_heads: int,
|
|
):
|
|
super().__init__(
|
|
model,
|
|
hidden_size,
|
|
num_heads,
|
|
use_multi_head_attention=True,
|
|
search_op_types=["SimplifiedLayerNormalization", "SkipSimplifiedLayerNormalization", "Add"],
|
|
)
|
|
|
|
def create_mha_node(
|
|
self,
|
|
input: str,
|
|
output: str,
|
|
q_rotary: NodeProto,
|
|
k_rotary: NodeProto,
|
|
v_matmul: NodeProto,
|
|
attn_mask: str = "",
|
|
add_qk: str = "",
|
|
past_k: str = "",
|
|
past_v: str = "",
|
|
present_k: str = "",
|
|
present_v: str = "",
|
|
scale: Optional[float] = None,
|
|
) -> Union[NodeProto, None]:
|
|
assert self.num_heads > 0
|
|
|
|
if self.hidden_size > 0 and (self.hidden_size % self.num_heads) != 0:
|
|
logger.debug(
|
|
f"fuse_rotary_attention: input hidden size {self.hidden_size} is not a multiple of num of heads {self.num_heads}"
|
|
)
|
|
return None
|
|
|
|
mha_node_name = self.model.create_node_name("MultiHeadAttention")
|
|
mha_inputs = [
|
|
q_rotary.output[0],
|
|
k_rotary.output[0],
|
|
v_matmul.output[0],
|
|
"", # bias
|
|
attn_mask, # key_padding_mask
|
|
add_qk, # relative_position_bias
|
|
past_k,
|
|
past_v,
|
|
]
|
|
|
|
mha_outputs = [output]
|
|
if present_k and present_v:
|
|
mha_outputs.extend([present_k, present_v])
|
|
|
|
mha_node = helper.make_node(
|
|
"MultiHeadAttention",
|
|
inputs=mha_inputs,
|
|
outputs=mha_outputs,
|
|
name=mha_node_name,
|
|
)
|
|
|
|
mha_node.domain = "com.microsoft"
|
|
mha_node.attribute.extend([helper.make_attribute("num_heads", self.num_heads)])
|
|
if scale is not None:
|
|
mha_node.attribute.extend([helper.make_attribute("scale", scale)])
|
|
if self.mask_filter_value is not None:
|
|
mha_node.attribute.extend([helper.make_attribute("mask_filter_value", float(self.mask_filter_value))])
|
|
|
|
self.increase_counter("MultiHeadAttention")
|
|
return mha_node
|
|
|
|
def check_runtime_shape_paths_for_function(
|
|
self,
|
|
reshape_qkv_2, # Reshape after Transpose
|
|
reshape_qkv_1, # Reshape before Transpose
|
|
reshape_q_2, # Reshape after RotaryEmbedding
|
|
reshape_k_2, # Reshape after RotaryEmbedding
|
|
reshape_v_2, # Reshape after Transpose
|
|
reshape_v_1, # Reshape before Transpose
|
|
add_qk, # Add before Softmax
|
|
root_input, # Root input to attention subgraph
|
|
):
|
|
# Check #1: check paths for qkv nodes
|
|
concat_qkv_2_path = self.model.match_parent_path(reshape_qkv_2, ["Concat"], [1])
|
|
concat_qkv_1_path = self.model.match_parent_path(reshape_qkv_1, ["Concat"], [1])
|
|
if concat_qkv_2_path is None or concat_qkv_1_path is None:
|
|
return False
|
|
concat_qkv_2, concat_qkv_1 = concat_qkv_2_path[0], concat_qkv_1_path[0]
|
|
|
|
reshape_qkv_2_path_1 = self.model.match_parent_path(concat_qkv_2, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_qkv_2_path_2 = self.model.match_parent_path(concat_qkv_2, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
reshape_qkv_1_path_1 = self.model.match_parent_path(concat_qkv_1, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_qkv_1_path_2 = self.model.match_parent_path(concat_qkv_1, ["Unsqueeze", "Gather", "Shape"], [2, 0, 0])
|
|
if (
|
|
reshape_qkv_2_path_1 is None
|
|
or reshape_qkv_2_path_2 is None
|
|
or reshape_qkv_1_path_1 is None
|
|
or reshape_qkv_1_path_2 is None
|
|
):
|
|
return False
|
|
|
|
_, gather_1, shape_1 = reshape_qkv_2_path_1
|
|
_, gather_2, shape_2 = reshape_qkv_2_path_2
|
|
|
|
# Check root_input --> Shape --> Gather connection
|
|
if shape_1.input[0] != root_input or shape_2.input[0] != root_input:
|
|
return False
|
|
|
|
# Check Gather --> Unsqueeze --> Concat --> Reshape connection for reshape_qkv_1_path_1 and reshape_qkv_1_path_2
|
|
if reshape_qkv_1_path_1[1].name != gather_1.name or reshape_qkv_1_path_2[1].name != gather_2.name:
|
|
return False
|
|
|
|
# Check #2: check paths for v nodes
|
|
concat_v_2_path = self.model.match_parent_path(reshape_v_2, ["Concat"], [1])
|
|
concat_v_1_path = self.model.match_parent_path(reshape_v_1, ["Concat"], [1])
|
|
if concat_v_2_path is None or concat_v_1_path is None:
|
|
return False
|
|
concat_v_2, concat_v_1 = concat_v_2_path[0], concat_v_1_path[0]
|
|
|
|
reshape_v_2_path_1 = self.model.match_parent_path(
|
|
concat_v_2, ["Unsqueeze", "Mul", "Gather", "Shape"], [0, 0, 0, 0]
|
|
)
|
|
reshape_v_2_path_2 = self.model.match_parent_path(
|
|
concat_v_2, ["Unsqueeze", "Add", "Gather", "Shape"], [1, 0, 0, 0]
|
|
)
|
|
reshape_v_1_path_1 = self.model.match_parent_path(concat_v_1, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_v_1_path_2 = self.model.match_parent_path(concat_v_1, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
if (
|
|
reshape_v_2_path_1 is None
|
|
or reshape_v_2_path_2 is None
|
|
or reshape_v_1_path_1 is None
|
|
or reshape_v_1_path_2 is None
|
|
):
|
|
return False
|
|
|
|
# Check Gather --> Mul --> Unsqueeze --> Concat --> Reshape connection for reshape_v_2_path_1
|
|
# Check Gather --> Add --> Unsqueeze --> Concat --> Reshape connection for reshape_v_2_path_2
|
|
# Check Gather --> Unsqueeze --> Concat --> Reshape connection for reshape_v_1_path_1 and reshape_v_1_path_2
|
|
if (
|
|
reshape_v_2_path_1[2].name != gather_1.name
|
|
or reshape_v_2_path_2[2].name != gather_2.name
|
|
or reshape_v_1_path_1[1].name != gather_1.name
|
|
or reshape_v_1_path_2[1].name != gather_2.name
|
|
):
|
|
return False
|
|
|
|
# Check #3: check paths for k nodes
|
|
concat_k_2_path = self.model.match_parent_path(reshape_k_2, ["Concat"], [1])
|
|
if concat_k_2_path is None:
|
|
return False
|
|
concat_k_2 = concat_k_2_path[0]
|
|
|
|
reshape_k_2_path_1 = self.model.match_parent_path(
|
|
concat_k_2, ["Unsqueeze", "Mul", "Gather", "Shape"], [0, 0, 0, 0]
|
|
)
|
|
reshape_k_2_path_2 = self.model.match_parent_path(
|
|
concat_k_2, ["Unsqueeze", "Add", "Gather", "Shape"], [2, 0, 0, 0]
|
|
)
|
|
if reshape_k_2_path_1 is None or reshape_k_2_path_2 is None:
|
|
return False
|
|
|
|
# Check Gather --> Mul --> Unsqueeze --> Concat --> Reshape connection for reshape_k_2_path_1
|
|
# Check Gather --> Add --> Unsqueeze --> Concat --> Reshape connection for reshape_k_2_path_2
|
|
if reshape_k_2_path_1[2].name != gather_1.name or reshape_k_2_path_2[2].name != gather_2.name:
|
|
return False
|
|
|
|
# Check #4: check paths for q nodes
|
|
concat_q_2_path = self.model.match_parent_path(reshape_q_2, ["Concat"], [1])
|
|
if concat_q_2_path is None:
|
|
return False
|
|
concat_q_2 = concat_q_2_path[0]
|
|
|
|
reshape_q_2_path_1 = self.model.match_parent_path(
|
|
concat_q_2, ["Unsqueeze", "Mul", "Gather", "Shape"], [0, 0, 0, 0]
|
|
)
|
|
reshape_q_2_path_2 = self.model.match_parent_path(concat_q_2, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
if reshape_q_2_path_1 is None or reshape_q_2_path_2 is None:
|
|
return False
|
|
|
|
# Check Gather --> Mul --> Unsqueeze --> Concat --> Reshape connection for reshape_q_2_path_1
|
|
# Check Gather --> Unsqueeze --> Concat --> Reshape connection for reshape_q_2_path_2
|
|
if reshape_q_2_path_1[2].name != gather_1.name or reshape_q_2_path_2[1].name != gather_2.name:
|
|
return False
|
|
|
|
# Check #5: check Mul nodes are the same for q, k, v
|
|
mul_q = reshape_q_2_path_1[1]
|
|
mul_k = reshape_k_2_path_1[1]
|
|
mul_v = reshape_v_2_path_1[1]
|
|
gather_1_out = gather_1.output[0]
|
|
if mul_q.input[0] != gather_1_out or mul_k.input[0] != gather_1_out or mul_v.input[0] != gather_1_out:
|
|
return False
|
|
|
|
# Check #6: check paths for attention mask nodes
|
|
attn_mask_path_1 = self.model.match_parent_path(add_qk, ["Concat", "Slice", "Slice"], [1, 0, 0])
|
|
attn_mask_path_2 = self.model.match_parent_path(add_qk, ["Cast", "Concat", "Slice", "Slice"], [1, 0, 0, 0])
|
|
if attn_mask_path_1 is not None:
|
|
_, slice_qk_2, slice_qk_1 = attn_mask_path_1
|
|
elif attn_mask_path_2 is not None:
|
|
_, _, slice_qk_2, slice_qk_1 = attn_mask_path_2
|
|
else:
|
|
return False
|
|
# Check first input to Slice #1 is 3D attention mask of shape (B,S,T)
|
|
if slice_qk_1.input[0] not in {"attn_mask", "attention_mask"}:
|
|
return False
|
|
|
|
slice_qk_2_path = self.model.match_parent_path(
|
|
slice_qk_2, ["Unsqueeze", "Add", "Gather", "Shape"], [2, 0, 1, 0]
|
|
)
|
|
slice_qk_1_path_1 = self.model.match_parent_path(
|
|
slice_qk_1, ["Unsqueeze", "Add", "Gather", "Shape"], [2, 0, 1, 0]
|
|
)
|
|
slice_qk_1_path_2 = self.model.match_parent_path(slice_qk_1, ["Unsqueeze"], [1])
|
|
if slice_qk_2_path is None or slice_qk_1_path_1 is None or slice_qk_1_path_2 is None:
|
|
return False
|
|
|
|
# Check Gather --> Add --> Unsqueeze #3 --> Slice #2 connection for slice_qk_2_path
|
|
# Check Gather --> Add --> Unsqueeze #2 --> Slice #1 connection for slice_qk_1_path_1
|
|
if slice_qk_2_path[1].name != slice_qk_1_path_1[1].name or slice_qk_2_path[2].name != slice_qk_1_path_1[2].name:
|
|
return False
|
|
|
|
# Check Unsqueeze #1 --> Slice #1 connection for slice_qk_1_path_2
|
|
# Check if first input to Add and Unsqueeze #1 is position ids
|
|
if slice_qk_1_path_1[1].input[0] != slice_qk_1_path_2[0].input[0]:
|
|
return False
|
|
|
|
return True
|
|
|
|
def check_runtime_shape_paths_for_nodes(
|
|
self,
|
|
reshape_qkv, # Final reshape before o_proj MatMul
|
|
reshape_q, # Reshape before q_proj MatMul
|
|
reshape_k, # Reshape before k_proj MatMul
|
|
reshape_v, # Reshape before v_proj MatMul
|
|
root_input, # Root input to attention subgraph
|
|
):
|
|
# Check #1: check paths for qkv nodes
|
|
concat_qkv_path = self.model.match_parent_path(reshape_qkv, ["Concat"], [1])
|
|
if concat_qkv_path is None:
|
|
return False
|
|
concat_qkv = concat_qkv_path[0]
|
|
|
|
reshape_qkv_path_1 = self.model.match_parent_path(concat_qkv, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_qkv_path_2 = self.model.match_parent_path(concat_qkv, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
if reshape_qkv_path_1 is None or reshape_qkv_path_2 is None:
|
|
return False
|
|
|
|
_, gather_1, shape_1 = reshape_qkv_path_1
|
|
_, gather_2, shape_2 = reshape_qkv_path_2
|
|
|
|
# Check root_input --> Shape --> Gather connection
|
|
if shape_1.input[0] != root_input or shape_2.input[0] != root_input:
|
|
return False
|
|
|
|
# Check #2: check paths for v nodes
|
|
concat_v_path = self.model.match_parent_path(reshape_v, ["Concat"], [1])
|
|
if concat_v_path is None:
|
|
return False
|
|
concat_v = concat_v_path[0]
|
|
|
|
reshape_v_path_1 = self.model.match_parent_path(concat_v, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_v_path_2 = self.model.match_parent_path(concat_v, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
if reshape_v_path_1 is None or reshape_v_path_2 is None:
|
|
return False
|
|
|
|
# Check Gather --> Unsqueeze --> Concat --> Reshape connection
|
|
if reshape_v_path_1[1].name != gather_1.name or reshape_v_path_2[1].name != gather_2.name:
|
|
return False
|
|
|
|
# Check #3: check paths for k nodes
|
|
concat_k_path = self.model.match_parent_path(reshape_k, ["Concat"], [1])
|
|
if concat_k_path is None:
|
|
return False
|
|
concat_k = concat_k_path[0]
|
|
|
|
reshape_k_path_1 = self.model.match_parent_path(concat_k, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_k_path_2 = self.model.match_parent_path(concat_k, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
if reshape_k_path_1 is None or reshape_k_path_2 is None:
|
|
return False
|
|
|
|
# Check Gather --> Unsqueeze --> Concat --> Reshape connection
|
|
if reshape_k_path_1[1].name != gather_1.name or reshape_k_path_2[1].name != gather_2.name:
|
|
return False
|
|
|
|
# Check #4: check paths for q nodes
|
|
concat_q_path = self.model.match_parent_path(reshape_q, ["Concat"], [1])
|
|
if concat_q_path is None:
|
|
return False
|
|
concat_q = concat_q_path[0]
|
|
|
|
reshape_q_path_1 = self.model.match_parent_path(concat_q, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
|
|
reshape_q_path_2 = self.model.match_parent_path(concat_q, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
|
|
if reshape_q_path_1 is None or reshape_q_path_2 is None:
|
|
return False
|
|
|
|
# Check Gather --> Unsqueeze --> Concat --> Reshape connection
|
|
if reshape_q_path_1[1].name != gather_1.name or reshape_q_path_2[1].name != gather_2.name:
|
|
return False
|
|
|
|
return True
|
|
|
|
def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
|
|
if normalize_node.op_type != "SkipSimplifiedLayerNormalization" and normalize_node.op_type != "Add":
|
|
return
|
|
|
|
# qkv_nodes_1 is for LLaMA-2 Microsoft
|
|
# qkv_nodes_2 is for LLaMA-2 Hugging Face
|
|
# qkv_nodes_3 is for LLaMA-2 distribute Hugging Face model
|
|
qkv_nodes = None
|
|
qkv_nodes_1 = self.model.match_parent_path(
|
|
normalize_node,
|
|
["MatMul", "Reshape", "Transpose", "Reshape", "MatMul"],
|
|
[1, 0, 0, 0, 0],
|
|
)
|
|
qkv_nodes_2 = self.model.match_parent_path(
|
|
normalize_node,
|
|
["MatMul", "Reshape", "Transpose", "MatMul"],
|
|
[1, 0, 0, 0],
|
|
)
|
|
qkv_nodes_3 = self.model.match_parent_path(
|
|
normalize_node,
|
|
["AllReduce", "MatMul", "Reshape", "Transpose", "MatMul"],
|
|
[1, 0, 0, 0, 0],
|
|
)
|
|
if qkv_nodes_1 is not None:
|
|
_, reshape_qkv_2, _, reshape_qkv_1, matmul_qkv = qkv_nodes_1
|
|
qkv_nodes = qkv_nodes_1
|
|
elif qkv_nodes_2 is not None:
|
|
_, reshape_qkv, _, matmul_qkv = qkv_nodes_2
|
|
qkv_nodes = qkv_nodes_2
|
|
elif qkv_nodes_3 is not None:
|
|
_, _, reshape_qkv, _, matmul_qkv = qkv_nodes_3
|
|
qkv_nodes = qkv_nodes_3
|
|
else:
|
|
logger.debug("fuse_rotary_attention: failed to match qkv nodes")
|
|
return
|
|
|
|
# v_nodes_1 is for LLaMA-2 Microsoft
|
|
# v_nodes_3 is for LLaMA-2 Hugging Face
|
|
# v_nodes_4 is for LLaMA-2 70B model
|
|
past_v, present_v, past_seq_len = "", "", ""
|
|
v_nodes = None
|
|
v_nodes_1 = self.model.match_parent_path(
|
|
matmul_qkv,
|
|
["Reshape", "Transpose", "Concat", "Transpose", "Reshape", "MatMul"],
|
|
[1, 0, 0, 1, 0, 0],
|
|
)
|
|
v_nodes_2 = self.model.match_parent_path(
|
|
matmul_qkv,
|
|
["Concat", "Transpose", "Reshape", "MatMul"],
|
|
[1, 1, 0, 0],
|
|
)
|
|
v_nodes_3 = self.model.match_parent_path(
|
|
matmul_qkv,
|
|
["Transpose", "Reshape", "MatMul"],
|
|
[1, 0, 0],
|
|
)
|
|
_, v_nodes_4, _ = self.model.match_parent_paths_all(
|
|
matmul_qkv,
|
|
[
|
|
(
|
|
["Reshape", "Expand", "Unsqueeze", "Concat", "Transpose", "Reshape", "MatMul"],
|
|
[1, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"Equal",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"Equal",
|
|
"Mul",
|
|
"ConstantOfShape",
|
|
"Shape",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"ConstantOfShape",
|
|
"Shape",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 1, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 2, 0, 4, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
["Reshape", "Concat", "Unsqueeze", "Gather", "Shape", "Concat", "Transpose", "Reshape", "MatMul"],
|
|
[1, 1, 0, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Mul",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 1, 1, 0, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
["Reshape", "Concat", "Unsqueeze", "Gather", "Shape", "Concat", "Transpose", "Reshape", "MatMul"],
|
|
[1, 1, 2, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
(
|
|
["Reshape", "Concat", "Unsqueeze", "Gather", "Shape", "Concat", "Transpose", "Reshape", "MatMul"],
|
|
[1, 1, 3, 0, 0, 0, 1, 0, 0],
|
|
),
|
|
],
|
|
output_name_to_node=None,
|
|
)
|
|
if v_nodes_1 is not None:
|
|
reshape_v_2, _, concat_v, _, reshape_v_1, matmul_v = v_nodes_1
|
|
v_nodes = v_nodes_1
|
|
|
|
concat_v_path = self.model.match_parent_path(
|
|
concat_v,
|
|
["Slice", "Unsqueeze"],
|
|
[0, 2],
|
|
)
|
|
if concat_v_path is None:
|
|
logger.debug("fuse_rotary_attention: failed to match past/present concat in v path")
|
|
return
|
|
|
|
past_v = concat_v_path[0].input[0]
|
|
past_seq_len = concat_v_path[-1].input[0]
|
|
present_v = concat_v.output[0]
|
|
elif v_nodes_2 is not None:
|
|
concat_v, transpose_v, reshape_v, matmul_v = v_nodes_2
|
|
v_nodes = v_nodes_2
|
|
past_v = concat_v.input[0]
|
|
present_v = concat_v.output[0]
|
|
elif v_nodes_3 is not None:
|
|
transpose_v, reshape_v, matmul_v = v_nodes_3
|
|
v_nodes = v_nodes_3
|
|
present_v = transpose_v.output[0]
|
|
elif v_nodes_4 is not None and len(v_nodes_4) == 9:
|
|
concat_v, transpose_v, reshape_v, matmul_v = v_nodes_4[0][-4:]
|
|
v_nodes = v_nodes_4
|
|
past_v = concat_v.input[0]
|
|
present_v = concat_v.output[0]
|
|
else:
|
|
logger.debug("fuse_rotary_attention: failed to match v path")
|
|
return
|
|
|
|
qk_nodes = self.model.match_parent_path(
|
|
matmul_qkv,
|
|
["Softmax", "Add", "Div", "MatMul"],
|
|
[0, 0, 0, 0],
|
|
)
|
|
add_qk, matmul_qk = None, None
|
|
if qk_nodes is not None:
|
|
_, add_qk, _, matmul_qk = qk_nodes
|
|
else:
|
|
logger.debug("fuse_rotary_attention: failed to match qk nodes")
|
|
return
|
|
|
|
# attn_mask_nodes_1, attn_mask_nodes_2 are for LLaMA-2 Microsoft's 3D attention mask
|
|
# attn_mask_nodes_3, attn_mask_nodes_4 are for LLaMA-2 Hugging Face's 2D attention mask
|
|
attn_mask, add_qk_str = "", ""
|
|
attn_mask_nodes_1 = self.model.match_parent_path(
|
|
add_qk,
|
|
["Concat", "Slice", "Slice"],
|
|
[1, 0, 0],
|
|
)
|
|
attn_mask_nodes_2 = self.model.match_parent_path(
|
|
add_qk,
|
|
["Cast", "Concat", "Slice", "Slice"],
|
|
[1, 0, 0, 0],
|
|
)
|
|
attn_mask_nodes_3 = self.model.match_parent_path(
|
|
add_qk,
|
|
["Add", "Where", "Sub", "Cast", "Expand", "Unsqueeze", "Unsqueeze"],
|
|
[1, 0, 2, 1, 0, 0, 0],
|
|
)
|
|
attn_mask_nodes_4 = self.model.match_parent_path(
|
|
add_qk,
|
|
["Where", "Sub", "Cast", "Expand", "Unsqueeze", "Unsqueeze"],
|
|
[1, 2, 1, 0, 0, 0],
|
|
)
|
|
attn_mask_nodes_5 = self.model.match_parent_path(
|
|
add_qk,
|
|
["Expand", "Add", "Where", "Sub", "Cast", "Expand", "Unsqueeze", "Unsqueeze"],
|
|
[1, 0, 0, 2, 1, 0, 0, 0],
|
|
)
|
|
attn_mask_nodes_6 = self.model.match_parent_path(
|
|
add_qk,
|
|
["Expand", "Where", "Sub", "Cast", "Expand", "Unsqueeze", "Unsqueeze"],
|
|
[1, 0, 2, 1, 0, 0, 0],
|
|
)
|
|
if attn_mask_nodes_1 is not None:
|
|
_, slice_mask_1, slice_mask_2 = attn_mask_nodes_1
|
|
attn_mask = slice_mask_1.output[0]
|
|
elif attn_mask_nodes_2 is not None:
|
|
_, _, slice_mask_1, slice_mask_2 = attn_mask_nodes_2
|
|
attn_mask = slice_mask_1.output[0]
|
|
elif attn_mask_nodes_3 is not None:
|
|
# Reshape from (B,1,S,T) to (B,N,S,T)
|
|
add_qk_str = self.reshape_add_qk(attn_mask_nodes_3[0].output[0])
|
|
elif attn_mask_nodes_4 is not None:
|
|
# Reshape from (B,1,S,T) to (B,N,S,T)
|
|
add_qk_str = self.reshape_add_qk(attn_mask_nodes_4[0].output[0])
|
|
elif attn_mask_nodes_5 is not None:
|
|
# The mask has already been reshaped to (B,N,S,T)
|
|
add_qk_str = attn_mask_nodes_5[0].output[0]
|
|
elif attn_mask_nodes_6 is not None:
|
|
# The mask has already been reshaped to (B,N,S,T)
|
|
add_qk_str = attn_mask_nodes_6[0].output[0]
|
|
else:
|
|
logger.debug("fuse_rotary_attention: failed to match attention mask nodes")
|
|
return
|
|
|
|
# k_nodes_1 is for LLaMA-2 Microsoft
|
|
# k_nodes_2 is for LLaMA-2 Hugging Face
|
|
# k_nodes_4 is for LLaMA-2 70B Hugging Face
|
|
past_k, present_k = "", ""
|
|
k_nodes = None
|
|
k_nodes_1 = self.model.match_parent_path(
|
|
matmul_qk,
|
|
["Reshape", "Transpose", "Concat", "Transpose", "RotaryEmbedding", "MatMul"],
|
|
[1, 0, 0, 1, 0, 0],
|
|
)
|
|
k_nodes_2 = self.model.match_parent_path(
|
|
matmul_qk,
|
|
["Transpose", "RotaryEmbedding", "Transpose", "Reshape", "MatMul"],
|
|
[1, 0, 0, 0, 0],
|
|
)
|
|
k_nodes_3 = self.model.match_parent_path(
|
|
matmul_qk,
|
|
["Transpose", "Concat", "RotaryEmbedding", "Transpose", "Reshape", "MatMul"],
|
|
[1, 0, 1, 0, 0, 0],
|
|
)
|
|
_, k_nodes_4, _ = self.model.match_parent_paths_all(
|
|
matmul_qk,
|
|
[
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Expand",
|
|
"Unsqueeze",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"Equal",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"Equal",
|
|
"Mul",
|
|
"ConstantOfShape",
|
|
"Shape",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"ConstantOfShape",
|
|
"Shape",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 0, 1, 1, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Expand",
|
|
"Where",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 0, 1, 2, 0, 4, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Mul",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 2, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
(
|
|
[
|
|
"Transpose",
|
|
"Reshape",
|
|
"Concat",
|
|
"Unsqueeze",
|
|
"Gather",
|
|
"Shape",
|
|
"Concat",
|
|
"RotaryEmbedding",
|
|
"Transpose",
|
|
"Reshape",
|
|
"MatMul",
|
|
],
|
|
[1, 0, 1, 3, 0, 0, 0, 1, 0, 0, 0],
|
|
),
|
|
],
|
|
output_name_to_node=None,
|
|
)
|
|
if k_nodes_1 is not None:
|
|
reshape_k_2, _, concat_k, _, rotary_k, matmul_k = k_nodes_1
|
|
k_nodes = k_nodes_1
|
|
|
|
concat_k_path = self.model.match_parent_path(
|
|
concat_k,
|
|
["Slice", "Unsqueeze"],
|
|
[0, 2],
|
|
)
|
|
if concat_k_path is None:
|
|
logger.debug("fuse_rotary_attention: failed to match past/present concat in k path")
|
|
return
|
|
|
|
past_k = concat_k_path[0].input[0]
|
|
shared_past_seq_len = concat_k_path[-1].input[0]
|
|
present_k = concat_k.output[0]
|
|
|
|
assert past_seq_len == shared_past_seq_len
|
|
elif k_nodes_2 is not None:
|
|
_, rotary_k, _, reshape_k, matmul_k = k_nodes_2
|
|
k_nodes = k_nodes_2
|
|
present_k = rotary_k.output[0]
|
|
elif k_nodes_3 is not None:
|
|
_, concat_k, rotary_k, _, reshape_k, matmul_k = k_nodes_3
|
|
k_nodes = k_nodes_3
|
|
past_k = concat_k.input[0]
|
|
present_k = concat_k.output[0]
|
|
elif k_nodes_4 is not None and len(k_nodes_4) == 9:
|
|
reshape_k, matmul_k = k_nodes_4[0][-2:]
|
|
concat_k, rotary_k = k_nodes_4[0][-5:-3]
|
|
k_nodes = k_nodes_4
|
|
past_k = concat_k.input[0]
|
|
present_k = concat_k.output[0]
|
|
else:
|
|
logger.debug("fuse_rotary_attention: failed to match k nodes")
|
|
return
|
|
|
|
# q_nodes_1 is for LLaMA-2 Microsoft
|
|
# q_nodes_2 is for LLaMA-2 Hugging Face
|
|
q_nodes = None
|
|
q_nodes_1 = self.model.match_parent_path(
|
|
matmul_qk,
|
|
["Reshape", "Transpose", "RotaryEmbedding", "MatMul"],
|
|
[0, 0, 0, 0],
|
|
)
|
|
q_nodes_2 = self.model.match_parent_path(
|
|
matmul_qk,
|
|
["RotaryEmbedding", "Transpose", "Reshape", "MatMul"],
|
|
[0, 0, 0, 0],
|
|
)
|
|
if q_nodes_1 is not None:
|
|
reshape_q_2, _, rotary_q, matmul_q = q_nodes_1
|
|
q_nodes = q_nodes_1
|
|
elif q_nodes_2 is not None:
|
|
rotary_q, _, reshape_q, matmul_q = q_nodes_2
|
|
q_nodes = q_nodes_2
|
|
else:
|
|
logger.debug("fuse_rotary_attention: failed to match q nodes")
|
|
return
|
|
|
|
if matmul_q.input[0] != matmul_k.input[0] and matmul_k.input[0] != matmul_v.input[0]:
|
|
logger.debug("fuse_rotary_attention: failed to find the same root_input for q, k, v paths")
|
|
return
|
|
|
|
root_output = ""
|
|
if qkv_nodes == qkv_nodes_1:
|
|
if not self.check_runtime_shape_paths_for_function(
|
|
reshape_qkv_2,
|
|
reshape_qkv_1,
|
|
reshape_q_2,
|
|
reshape_k_2,
|
|
reshape_v_2,
|
|
reshape_v_1,
|
|
add_qk,
|
|
matmul_q.input[0],
|
|
):
|
|
logger.debug("fuse_rotary_attention: failed to verify runtime shape paths")
|
|
return
|
|
root_output = reshape_qkv_2.output[0]
|
|
|
|
elif qkv_nodes in (qkv_nodes_2, qkv_nodes_3):
|
|
if not self.check_runtime_shape_paths_for_nodes(
|
|
reshape_qkv,
|
|
reshape_q,
|
|
reshape_k,
|
|
reshape_v,
|
|
matmul_q.input[0],
|
|
):
|
|
logger.debug("fuse_rotary_attention: failed to verify runtime shape paths")
|
|
return
|
|
root_output = reshape_qkv.output[0]
|
|
|
|
# Rename inputs of rotary_q/k so it connects with output of matmul_q/k
|
|
# Before: MatMul --> Reshape --> Transpose --> RotaryEmbedding
|
|
# After: MatMul --> RotaryEmbedding
|
|
rotary_q.input[0] = matmul_q.output[0]
|
|
rotary_k.input[0] = matmul_k.output[0]
|
|
|
|
# Rename current output of rotary_k (present_key) so it doesn't match output of MHA (present_key)
|
|
rotary_k.output[0] = rotary_k.name + "_output_0"
|
|
|
|
if qkv_nodes == qkv_nodes_3:
|
|
qkv_nodes = qkv_nodes[1:]
|
|
|
|
new_node = self.create_mha_node(
|
|
matmul_q.input[0],
|
|
root_output,
|
|
rotary_q,
|
|
rotary_k,
|
|
matmul_v,
|
|
attn_mask,
|
|
add_qk_str,
|
|
past_k,
|
|
past_v,
|
|
present_k,
|
|
present_v,
|
|
)
|
|
if new_node is None:
|
|
logger.debug("fuse_rotary_attention: failed to create multi-head attention with rotary embeddings")
|
|
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(qkv_nodes[1:])
|
|
|
|
if v_nodes != v_nodes_4:
|
|
self.nodes_to_remove.extend(v_nodes[:-1])
|
|
else:
|
|
nodes_to_keep = [v_nodes[0][-1]]
|
|
for temp_path in v_nodes:
|
|
self.add_nodes_to_remove_with_nodes_to_keep(temp_path, nodes_to_keep)
|
|
|
|
self.nodes_to_remove.extend(qk_nodes)
|
|
|
|
if k_nodes == k_nodes_1:
|
|
self.nodes_to_remove.extend(k_nodes[:-2])
|
|
elif k_nodes == k_nodes_2:
|
|
self.nodes_to_remove.append(k_nodes[0])
|
|
self.nodes_to_remove.append(k_nodes[2])
|
|
self.nodes_to_remove.append(k_nodes[3])
|
|
elif k_nodes == k_nodes_3:
|
|
self.nodes_to_remove.append(k_nodes[0])
|
|
self.nodes_to_remove.append(k_nodes[1])
|
|
self.nodes_to_remove.append(k_nodes[3])
|
|
self.nodes_to_remove.append(k_nodes[4])
|
|
elif k_nodes == k_nodes_4:
|
|
nodes_to_keep = [k_nodes[0][-1], k_nodes[0][-4]]
|
|
for temp_path in k_nodes:
|
|
self.add_nodes_to_remove_with_nodes_to_keep(temp_path, nodes_to_keep)
|
|
|
|
if q_nodes == q_nodes_1:
|
|
self.nodes_to_remove.extend(q_nodes[:-2])
|
|
elif q_nodes == q_nodes_2:
|
|
self.nodes_to_remove.append(q_nodes[1])
|
|
self.nodes_to_remove.append(q_nodes[2])
|
|
|
|
self.prune_graph = True
|
|
|
|
|
|
class FusionRotaryEmbeddings(Fusion):
|
|
def __init__(self, model: OnnxModel):
|
|
self.base_name = "RotaryEmbedding"
|
|
super().__init__(model, self.base_name, [self.base_name, self.base_name + ".1", "Add"])
|
|
|
|
# The RotaryEmbedding function can have multiple extraneous constant outputs even though the function is supposed to produce only one output.
|
|
# This is a byproduct of a potential CSE bug when using `export_modules_as_functions` in the TorchScript exporter.
|
|
# To work around this issue, we set the extraneous constant values from the RotaryEmbedding function as initializers in the locations where they are actually used.
|
|
def reassign_extra_outputs(self, rot_emb_node: NodeProto, function: FunctionProto):
|
|
# Find extra outputs and Constant nodes attached to those outputs
|
|
extra_constants, extra_outputs = [], []
|
|
for fn_node in function.node:
|
|
if fn_node.op_type == "Constant" and fn_node.input == [] and fn_node.output[0] in function.output:
|
|
extra_constants.append(fn_node)
|
|
output_index = list(function.output).index(fn_node.output[0])
|
|
extra_outputs.append(rot_emb_node.output[output_index])
|
|
|
|
# Set extra Constant node outputs as initializers
|
|
extra_initializers = []
|
|
for extra_constant in extra_constants:
|
|
constant_tensorproto = extra_constant.attribute[0].t
|
|
constant_tensorproto.name = self.model.create_node_name("Constant")
|
|
self.model.add_initializer(constant_tensorproto)
|
|
extra_initializers.append(constant_tensorproto.name)
|
|
|
|
# Update references of Constant node outputs to initializer references
|
|
for extra_output, extra_initializer in zip(extra_outputs, extra_initializers):
|
|
nodes_to_update = list(filter(lambda entry: extra_output in entry.input, self.model.model.graph.node))
|
|
for node_to_update in nodes_to_update:
|
|
OnnxModel.replace_node_input(node_to_update, extra_output, extra_initializer)
|
|
|
|
return extra_outputs
|
|
|
|
def create_rotary_embeddings_from_function(self, node: NodeProto):
|
|
rotary_emb_node_name = self.model.create_node_name(self.base_name)
|
|
|
|
matmul_path = self.model.match_parent_path(
|
|
node,
|
|
["Reshape", "MatMul"],
|
|
[0, 0],
|
|
)
|
|
if matmul_path is not None:
|
|
reshape_node, matmul_node = matmul_path
|
|
else:
|
|
logger.debug("fuse_rotary_embeddings: failed to match MatMul")
|
|
return
|
|
|
|
rotary_emb_inputs = [
|
|
matmul_node.output[0], # x is of shape (B,S,D) instead of (B,S,N,H)
|
|
node.input[1], # position_ids
|
|
]
|
|
|
|
# Convert cos_cache and sin_cache from node attributes to model initializers
|
|
cos_cache_node = list(filter(lambda constant: constant.output[0] == node.input[2], self.model.model.graph.node))
|
|
sin_cache_node = list(filter(lambda constant: constant.output[0] == node.input[3], self.model.model.graph.node))
|
|
cos_cache_name, sin_cache_name = "cos_cache", "sin_cache"
|
|
|
|
if (
|
|
len(cos_cache_node) == 1
|
|
and len(sin_cache_node) == 1
|
|
and self.model.get_initializer(cos_cache_name) is None
|
|
and self.model.get_initializer(sin_cache_name) is None
|
|
):
|
|
cos_cache = numpy_helper.to_array(cos_cache_node[0].attribute[0].t).squeeze()
|
|
sin_cache = numpy_helper.to_array(sin_cache_node[0].attribute[0].t).squeeze()
|
|
|
|
cos_cache_tensor = helper.make_tensor(
|
|
name=cos_cache_name,
|
|
data_type=TensorProto.FLOAT,
|
|
dims=list(cos_cache.shape),
|
|
vals=cos_cache.flatten().tolist(),
|
|
)
|
|
self.model.add_initializer(cos_cache_tensor, self.this_graph_name)
|
|
sin_cache_tensor = helper.make_tensor(
|
|
name=sin_cache_name,
|
|
data_type=TensorProto.FLOAT,
|
|
dims=list(sin_cache.shape),
|
|
vals=sin_cache.flatten().tolist(),
|
|
)
|
|
self.model.add_initializer(sin_cache_tensor, self.this_graph_name)
|
|
|
|
self.nodes_to_remove.extend([cos_cache_node[0], sin_cache_node[0]])
|
|
|
|
rotary_emb_inputs.extend([cos_cache_name, sin_cache_name])
|
|
|
|
rotary_emb_outputs = node.output
|
|
if len(rotary_emb_outputs) > 1:
|
|
# Re-assign extraneous constant outputs in RotaryEmbedding functions as initializers
|
|
func = list(filter(lambda fn: fn.name == node.op_type, self.model.model.functions))
|
|
assert len(func) == 1
|
|
extra_outputs = self.reassign_extra_outputs(node, func[0])
|
|
rotary_emb_outputs = list(filter(lambda output_name: output_name not in extra_outputs, rotary_emb_outputs))
|
|
assert len(rotary_emb_outputs) == 1
|
|
|
|
rotary_emb_node = helper.make_node(
|
|
self.base_name,
|
|
inputs=rotary_emb_inputs,
|
|
outputs=rotary_emb_outputs,
|
|
name=rotary_emb_node_name,
|
|
interleaved=1,
|
|
)
|
|
rotary_emb_node.domain = "com.microsoft"
|
|
|
|
self.nodes_to_remove.append(reshape_node)
|
|
|
|
return rotary_emb_node
|
|
|
|
def create_rotary_embeddings_from_nodes(
|
|
self,
|
|
root_input: str,
|
|
position_ids: str,
|
|
cos_slice: str,
|
|
sin_slice: str,
|
|
output: str,
|
|
):
|
|
rotary_emb_node_name = self.model.create_node_name(self.base_name)
|
|
|
|
# Convert cos_cache and sin_cache from node attributes to model initializers
|
|
cos_cache_node = list(filter(lambda constant: constant.output[0] == cos_slice, self.model.model.graph.node))
|
|
sin_cache_node = list(filter(lambda constant: constant.output[0] == sin_slice, self.model.model.graph.node))
|
|
cos_cache_name, sin_cache_name = "cos_cache", "sin_cache"
|
|
|
|
if (
|
|
len(cos_cache_node) == 1
|
|
and len(sin_cache_node) == 1
|
|
and self.model.get_initializer(cos_cache_name) is None
|
|
and self.model.get_initializer(sin_cache_name) is None
|
|
):
|
|
cos_cache = numpy_helper.to_array(cos_cache_node[0].attribute[0].t).squeeze()
|
|
sin_cache = numpy_helper.to_array(sin_cache_node[0].attribute[0].t).squeeze()
|
|
|
|
# Reshape cos/sin cache from (M, H) to (M, H/2)
|
|
head_size = cos_cache.shape[1]
|
|
cos_cache = cos_cache[:, : (head_size // 2)]
|
|
sin_cache = sin_cache[:, : (head_size // 2)]
|
|
|
|
cos_cache_tensor = helper.make_tensor(
|
|
name=cos_cache_name,
|
|
data_type=TensorProto.FLOAT,
|
|
dims=list(cos_cache.shape),
|
|
vals=cos_cache.flatten().tolist(),
|
|
)
|
|
self.model.add_initializer(cos_cache_tensor, self.this_graph_name)
|
|
sin_cache_tensor = helper.make_tensor(
|
|
name=sin_cache_name,
|
|
data_type=TensorProto.FLOAT,
|
|
dims=list(sin_cache.shape),
|
|
vals=sin_cache.flatten().tolist(),
|
|
)
|
|
self.model.add_initializer(sin_cache_tensor, self.this_graph_name)
|
|
|
|
self.nodes_to_remove.extend([cos_cache_node[0], sin_cache_node[0]])
|
|
|
|
rotary_emb_node = helper.make_node(
|
|
self.base_name,
|
|
inputs=[root_input, position_ids, cos_cache_name, sin_cache_name],
|
|
outputs=[output],
|
|
name=rotary_emb_node_name,
|
|
interleaved=0,
|
|
)
|
|
rotary_emb_node.domain = "com.microsoft"
|
|
return rotary_emb_node
|
|
|
|
def fuse(self, node, input_name_to_nodes, output_name_to_node):
|
|
# Node is either RotaryEmbedding function or Add
|
|
if self.base_name not in node.op_type and node.op_type != "Add":
|
|
return
|
|
|
|
# Check if node is "RotaryEmbedding nn.Module" exported as a function
|
|
# (e.g. export_modules_as_functions={RotaryEmbedding} in torch.onnx.export)
|
|
rotary_emb_node = None
|
|
if node.op_type != "Add":
|
|
# Verify that function has the correct inputs
|
|
if len(node.input) not in {4, 5} or node.input[1] not in {
|
|
"pos",
|
|
"pos_id",
|
|
"position_id",
|
|
"pos_ids",
|
|
"position_ids",
|
|
}:
|
|
logger.debug("fuse_rotary_embeddings: failed to verify inputs for RotaryEmbedding function")
|
|
return
|
|
|
|
rotary_emb_node = self.create_rotary_embeddings_from_function(node)
|
|
if rotary_emb_node is None:
|
|
logger.debug("fuse_rotary_embeddings: failed to create RotaryEmbedding node")
|
|
return
|
|
|
|
# Remove RotaryEmbedding function
|
|
self.nodes_to_remove.append(node)
|
|
|
|
# Remove RotaryEmbedding function's shape inference stored in value_info
|
|
# The new shape will be calculated during symbolic shape inference
|
|
old_shape_infer = list(
|
|
filter(lambda node: node.name == rotary_emb_node.output[0], self.model.model.graph.value_info)
|
|
)
|
|
assert len(old_shape_infer) == 1
|
|
self.model.model.graph.value_info.remove(old_shape_infer[0])
|
|
|
|
else:
|
|
# Rotary embeddings are defined using the below functions:
|
|
#
|
|
# def rotate_half(x):
|
|
# """Rotates half the hidden dims of the input."""
|
|
# x1 = x[..., : x.shape[-1] // 2]
|
|
# x2 = x[..., x.shape[-1] // 2 :]
|
|
# return torch.cat((-x2, x1), dim=-1)
|
|
#
|
|
# def apply_rope(x, cos, sin, position_ids):
|
|
# cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
|
|
# sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
|
|
# cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
|
# sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
|
# x_embed = (x * cos) + (rotate_half(x) * sin)
|
|
# return x_embed
|
|
|
|
# Check paths for rotate_half(x)
|
|
rotate_half_x2_path_1 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Concat", "Neg", "Slice", "Transpose"],
|
|
[1, 0, 0, 0, 0],
|
|
)
|
|
rotate_half_x2_path_2 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Concat", "Neg", "Slice", "Unsqueeze", "Div", "Gather", "Shape", "Transpose"],
|
|
[1, 0, 0, 0, 1, 0, 0, 0, 0],
|
|
)
|
|
if rotate_half_x2_path_1 is None or rotate_half_x2_path_2 is None:
|
|
logger.debug("fuse_rotary_embeddings: failed to match x2 in rotate_half")
|
|
return
|
|
|
|
rotate_half_x1_path_1 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Concat", "Slice", "Transpose"],
|
|
[1, 0, 1, 0],
|
|
)
|
|
rotate_half_x1_path_2 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Concat", "Slice", "Unsqueeze", "Div", "Gather", "Shape", "Transpose"],
|
|
[1, 0, 1, 2, 0, 0, 0, 0],
|
|
)
|
|
if rotate_half_x1_path_1 is None or rotate_half_x1_path_2 is None:
|
|
logger.debug("fuse_rotary_embeddings: failed to match x1 in rotate_half")
|
|
return
|
|
|
|
if (
|
|
rotate_half_x1_path_1[-1].name != rotate_half_x1_path_2[-1].name
|
|
or rotate_half_x2_path_1[-1].name != rotate_half_x2_path_2[-1].name
|
|
or rotate_half_x1_path_1[-1].name != rotate_half_x2_path_1[-1].name
|
|
or rotate_half_x1_path_2[-1].name != rotate_half_x2_path_2[-1].name
|
|
):
|
|
logger.debug("fuse_rotary_embeddings: failed to match common input in rotate_half")
|
|
return
|
|
|
|
# Check path for x
|
|
x_path = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Transpose"],
|
|
[0, 0],
|
|
)
|
|
if x_path is None:
|
|
logger.debug("fuse_rotary_embeddings: failed to match x in rotate_half")
|
|
return
|
|
|
|
# Check path for sin
|
|
sin_path, sin_cache, position_ids = None, "", ""
|
|
sin_path_1 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Squeeze", "Squeeze", "Slice", "Unsqueeze", "Gather", "Shape"],
|
|
[1, 1, 0, 0, 0, 0, 2, 0, 0],
|
|
)
|
|
sin_path_2 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Squeeze", "Squeeze", "Slice", "Unsqueeze", "Add"],
|
|
[1, 1, 0, 0, 0, 0, 2, 0],
|
|
)
|
|
sin_path_3 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Slice", "Unsqueeze", "Gather", "Shape"],
|
|
[1, 1, 0, 0, 2, 0, 0],
|
|
)
|
|
sin_path_4 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Slice", "Unsqueeze", "Add"],
|
|
[1, 1, 0, 0, 2, 0],
|
|
)
|
|
if sin_path_1 is not None:
|
|
sin_path = sin_path_1
|
|
sin_cache = sin_path[-4].input[0]
|
|
elif sin_path_2 is not None:
|
|
sin_path = sin_path_2
|
|
sin_cache = sin_path[-3].input[0]
|
|
elif sin_path_3 is not None:
|
|
sin_path = sin_path_3
|
|
sin_cache = sin_path[-4].input[0]
|
|
position_ids = sin_path[2].input[1]
|
|
elif sin_path_4 is not None:
|
|
sin_path = sin_path_4
|
|
sin_cache = sin_path[-3].input[0]
|
|
position_ids = sin_path[2].input[1]
|
|
else:
|
|
logger.debug("fuse_rotary_embeddings: failed to match sin path in apply_rope")
|
|
return
|
|
|
|
# Check path for cos
|
|
cos_path, cos_cache = None, ""
|
|
cos_path_1 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Squeeze", "Squeeze", "Slice", "Unsqueeze", "Gather", "Shape"],
|
|
[0, 1, 0, 0, 0, 0, 2, 0, 0],
|
|
)
|
|
cos_path_2 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Squeeze", "Squeeze", "Slice", "Unsqueeze", "Add"],
|
|
[0, 1, 0, 0, 0, 0, 2, 0],
|
|
)
|
|
cos_path_3 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Slice", "Unsqueeze", "Gather", "Shape"],
|
|
[0, 1, 0, 0, 2, 0, 0],
|
|
)
|
|
cos_path_4 = self.model.match_parent_path(
|
|
node,
|
|
["Mul", "Unsqueeze", "Gather", "Slice", "Unsqueeze", "Add"],
|
|
[0, 1, 0, 0, 2, 0],
|
|
)
|
|
if cos_path_1 is not None:
|
|
cos_path = cos_path_1
|
|
cos_cache = cos_path[-4].input[0]
|
|
elif cos_path_2 is not None:
|
|
cos_path = cos_path_2
|
|
cos_cache = cos_path[-3].input[0]
|
|
elif cos_path_3 is not None:
|
|
cos_path = cos_path_3
|
|
cos_cache = cos_path[-4].input[0]
|
|
position_ids = cos_path[2].input[1]
|
|
elif cos_path_4 is not None:
|
|
cos_path = cos_path_4
|
|
cos_cache = cos_path[-3].input[0]
|
|
position_ids = cos_path[2].input[1]
|
|
else:
|
|
logger.debug("fuse_rotary_embeddings: failed to match sin path in apply_rope")
|
|
return
|
|
|
|
# Check path for position ids
|
|
if position_ids == "":
|
|
position_ids_from_sin_path = self.model.match_parent_path(
|
|
sin_path[2],
|
|
["Reshape"],
|
|
[1],
|
|
)
|
|
position_ids_from_cos_path = self.model.match_parent_path(
|
|
cos_path[2],
|
|
["Reshape"],
|
|
[1],
|
|
)
|
|
if (
|
|
position_ids_from_sin_path is None
|
|
or position_ids_from_cos_path is None
|
|
or position_ids_from_sin_path[0].name != position_ids_from_cos_path[0].name
|
|
):
|
|
logger.debug("fuse_rotary_embeddings: failed to match position ids path in apply_rope")
|
|
return
|
|
position_ids = position_ids_from_cos_path[0].input[0]
|
|
else:
|
|
position_ids_from_sin_path = []
|
|
position_ids_from_cos_path = []
|
|
|
|
past_seq_len_path, curr_seq_len_path = None, None
|
|
if (sin_path == sin_path_1 and cos_path == cos_path_1) or (
|
|
sin_path == sin_path_3 and cos_path == cos_path_3
|
|
):
|
|
if sin_path[-2].name != cos_path[-2].name or sin_path[-1].name != cos_path[-1].name:
|
|
logger.debug(
|
|
"fuse_rotary_embeddings: failed to match common Gather node and Shape node in sin cache and cos cache"
|
|
)
|
|
return
|
|
elif (sin_path == sin_path_2 and cos_path == cos_path_2) or (
|
|
sin_path == sin_path_4 and cos_path == cos_path_4
|
|
):
|
|
if sin_path[-1].name != cos_path[-1].name:
|
|
logger.debug("fuse_rotary_embeddings: failed to match common Add node in sin cache and cos cache")
|
|
return
|
|
# Match past sequence length path: past_key --> Shape --> Gather --> Add
|
|
past_seq_len_path = self.model.match_parent_path(
|
|
sin_path[-1],
|
|
["Gather", "Shape"],
|
|
[1, 0],
|
|
)
|
|
# Match current sequence length path: transpose_k --> Shape --> Gather --> Add
|
|
curr_seq_len_path = self.model.match_parent_path(
|
|
sin_path[-1],
|
|
["Gather", "Shape", "Transpose"],
|
|
[0, 0, 0],
|
|
)
|
|
if (
|
|
past_seq_len_path is None
|
|
or curr_seq_len_path is None
|
|
or self.model.find_graph_input(past_seq_len_path[-1].input[0]) is None
|
|
or curr_seq_len_path[-1].op_type != "Transpose"
|
|
):
|
|
logger.debug("fuse_rotary_embeddings: failed to match past_seq_len and curr_seq_len paths")
|
|
return
|
|
else:
|
|
logger.debug("fuse_rotary_embeddings: failed to match common cache paths")
|
|
|
|
rotary_emb_node = self.create_rotary_embeddings_from_nodes(
|
|
rotate_half_x1_path_1[-1].output[0],
|
|
position_ids,
|
|
cos_cache,
|
|
sin_cache,
|
|
node.output[0],
|
|
)
|
|
if rotary_emb_node is None:
|
|
logger.debug("fuse_rotary_embeddings: failed to create RotaryEmbedding node")
|
|
return
|
|
|
|
# Remove rotary embedding nodes
|
|
self.add_nodes_to_remove([node])
|
|
self.add_nodes_to_remove(rotate_half_x1_path_1[:-1])
|
|
self.add_nodes_to_remove(rotate_half_x1_path_2[:-1])
|
|
self.add_nodes_to_remove(rotate_half_x2_path_1[:-1])
|
|
self.add_nodes_to_remove(rotate_half_x2_path_2[:-1])
|
|
self.add_nodes_to_remove(x_path[:-1])
|
|
self.add_nodes_to_remove(sin_path)
|
|
self.add_nodes_to_remove(cos_path)
|
|
self.add_nodes_to_remove(position_ids_from_sin_path[:-1])
|
|
self.add_nodes_to_remove(position_ids_from_cos_path[:-1])
|
|
|
|
if past_seq_len_path is not None and len(self.model.get_children(past_seq_len_path[0])) == 1:
|
|
# In merged HF model, output of Gather in past_seq_len_path is used twice
|
|
# for past_key_values.0.key and once for other past_key_values
|
|
self.add_nodes_to_remove(past_seq_len_path)
|
|
if curr_seq_len_path is not None:
|
|
self.add_nodes_to_remove(curr_seq_len_path[:-1])
|
|
|
|
self.increase_counter(self.base_name)
|
|
self.node_name_to_graph_name[rotary_emb_node.name] = self.this_graph_name
|
|
self.nodes_to_add.append(rotary_emb_node)
|
|
self.prune_graph = True
|