wehub-resource-sync
136 lines
5.4 KiB
Python
136 lines
5.4 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import math
|
|
|
|
import paddle
|
|
import paddle.nn.functional as F
|
|
|
|
import paddlenlp
|
|
from paddlenlp.transformers.llama.modeling import get_triangle_upper_mask
|
|
|
|
ssa_group_size_ratio = 1 / 4
|
|
|
|
|
|
def shift(qkv, bsz, q_len, group_size, num_heads, head_dim):
|
|
assert qkv.shape == [bsz, num_heads, q_len, head_dim], "qkv shape does not match expected shape"
|
|
# Calculate the shift amount for rolling
|
|
shift_amount = -group_size // 2
|
|
# Roll the qkv tensor along the sequence length axis
|
|
qkv[:, num_heads // 2 :] = qkv[:, num_heads // 2 :].roll(shift_amount, axis=2)
|
|
|
|
# Reshape the tensor to the desired shape
|
|
qkv = qkv.reshape([bsz * (q_len // group_size), group_size, num_heads, head_dim])
|
|
return qkv
|
|
|
|
|
|
def ssa_scaled_dot_product_attention(
|
|
query_states,
|
|
config,
|
|
key_states,
|
|
value_states,
|
|
attention_mask,
|
|
output_attentions,
|
|
alibi=None,
|
|
sequence_parallel=False,
|
|
reshard_layer=None,
|
|
**kwargs
|
|
):
|
|
bsz, q_len, num_heads, head_dim = query_states.shape
|
|
if config.context_parallel_degree > 1:
|
|
raise ValueError("Context parallel requires `use_flash_attention=True`")
|
|
# [ bz, seqlen, nhead, head_dim] -> [bs, nhead, seq_len, head_dim]
|
|
query_states = paddle.transpose(query_states, [0, 2, 1, 3])
|
|
# merge with the next transpose
|
|
key_states = paddle.transpose(key_states, [0, 2, 1, 3])
|
|
value_states = paddle.transpose(value_states, [0, 2, 1, 3])
|
|
assert ssa_group_size_ratio is not None, "ssa_group_size_ratio must provide"
|
|
|
|
# Calculate the group size based on the sequence length and the group size ratio
|
|
group_size = q_len if int(q_len * ssa_group_size_ratio) == 0 else int(q_len * ssa_group_size_ratio)
|
|
assert q_len % group_size == 0, f"q_len {q_len} must be divisible by group size {group_size}."
|
|
|
|
num_group = q_len // group_size
|
|
|
|
# Apply shifting to the query, key, and value states
|
|
query_states = shift(query_states, bsz, q_len, group_size, num_heads, head_dim)
|
|
key_states = shift(key_states, bsz, q_len, group_size, num_heads, head_dim)
|
|
value_states = shift(value_states, bsz, q_len, group_size, num_heads, head_dim)
|
|
query_states = paddle.transpose(query_states, [0, 2, 1, 3])
|
|
key_states = paddle.transpose(key_states, [0, 2, 1, 3])
|
|
value_states = paddle.transpose(value_states, [0, 2, 1, 3])
|
|
# matmul and device by sqrt(head_dim)
|
|
attn_weights = paddle.matmul(query_states / math.sqrt(head_dim), key_states.transpose([0, 1, 3, 2]))
|
|
|
|
# then add alibi bias
|
|
if alibi is not None:
|
|
alibi = alibi.reshape([bsz, num_heads, 1, -1])
|
|
attn_weights = attn_weights + alibi
|
|
if paddle.in_dynamic_mode() and attn_weights.shape != [bsz * num_group, num_heads, group_size, group_size]:
|
|
raise ValueError(
|
|
f"Attention weights should be of shape {(bsz * num_group, num_heads, group_size, group_size)}, but is"
|
|
f" {attn_weights.shape}"
|
|
)
|
|
|
|
# In sep mode, the attenion mask should be created in the runtime.
|
|
if reshard_layer is not None:
|
|
attention_mask = None
|
|
|
|
if attention_mask is None:
|
|
attention_mask = get_triangle_upper_mask(attn_weights)
|
|
attention_mask = paddle.tile(
|
|
paddle.cast(attention_mask[:, :, :group_size, :group_size], dtype="float32"), [num_group, 1, 1, 1]
|
|
)
|
|
|
|
if attention_mask.shape != [bsz * num_group, 1, group_size, group_size]:
|
|
attention_mask = attention_mask[: bsz * num_group, :, :, :]
|
|
|
|
attn_weights = attn_weights + attention_mask
|
|
if not paddle.in_dynamic_mode():
|
|
attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(query_states.dtype)
|
|
else:
|
|
with paddle.amp.auto_cast(False):
|
|
attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(query_states.dtype)
|
|
|
|
attn_output = paddle.matmul(attn_weights, value_states)
|
|
attn_output = attn_output.transpose([0, 2, 1, 3])
|
|
|
|
# shift back
|
|
attn_output = attn_output.reshape([bsz, q_len, num_heads, head_dim])
|
|
attn_output[:, :, num_heads // 2 :] = attn_output[:, :, num_heads // 2 :].roll(group_size // 2, axis=1)
|
|
|
|
if reshard_layer is not None:
|
|
attn_output = reshard_layer(
|
|
attn_output,
|
|
split_axis=1,
|
|
concat_axis=2,
|
|
)
|
|
q_len = q_len // config.sep_parallel_degree
|
|
num_heads = num_heads * config.sep_parallel_degree
|
|
|
|
if sequence_parallel:
|
|
attn_output = attn_output.reshape([bsz * q_len, head_dim * num_heads])
|
|
else:
|
|
attn_output = attn_output.reshape([bsz, q_len, head_dim * num_heads])
|
|
return (attn_output, attn_weights) if output_attentions else attn_output
|
|
|
|
|
|
def set_group_size(group_size_ratio):
|
|
global ssa_group_size_ratio
|
|
ssa_group_size_ratio = group_size_ratio
|
|
|
|
|
|
def replace_llama_attn():
|
|
paddlenlp.transformers.llama.modeling.scaled_dot_product_attention = ssa_scaled_dot_product_attention
|