wehub-resource-sync
145 lines
7.4 KiB
Python
145 lines
7.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
from deepspeed import comm as dist
|
|
from ..config import DeepSpeedInferenceConfig
|
|
from .base import BaseOp
|
|
from .softmax import SoftmaxOp
|
|
from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext
|
|
|
|
|
|
class SoftmaxContextOp(BaseOp):
|
|
|
|
def __init__(self, config: DeepSpeedInferenceConfig):
|
|
super(SoftmaxContextOp, self).__init__(config)
|
|
try:
|
|
if self.config.dtype in [torch.float16, torch.int8]:
|
|
self.softmax_context_func = self.inference_module.softmax_context_fp16
|
|
elif self.config.dtype == torch.bfloat16:
|
|
self.softmax_context_func = self.inference_module.softmax_context_bf16
|
|
else:
|
|
self.softmax_context_func = self.inference_module.softmax_context_fp32
|
|
except AttributeError:
|
|
self.softmax_context_func = self.softmax_context_fallback
|
|
|
|
@staticmethod
|
|
def transform4d_0213(x, seq_length):
|
|
assert x.dim() == 3, F"Dim {x.dim()} is not supported"
|
|
batch_size, num_heads, seq_length_head_dim = x.shape
|
|
head_dim = seq_length_head_dim // seq_length
|
|
x = x.view(batch_size, num_heads, seq_length, head_dim)
|
|
x = x.permute(0, 2, 1, 3)
|
|
|
|
return x
|
|
|
|
@staticmethod
|
|
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
|
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
|
if n_rep <= 1 or num_key_value_heads == 1:
|
|
return hidden_states
|
|
|
|
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
|
|
|
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
|
|
|
@staticmethod
|
|
def bias_add_transform_0213(input, bias, num_heads, trans_count, perform_bias=False):
|
|
assert trans_count == 1 or trans_count == 3, F"Trans count {trans_count} is not supported"
|
|
assert input.dim() == 3, F"Dim {input.dim()} is not supported"
|
|
input_biased = torch.add(input, bias) if perform_bias else input
|
|
batch_size, seq_length, value_size = input_biased.shape
|
|
hid_dim = value_size // trans_count
|
|
head_dim = hid_dim // num_heads
|
|
|
|
if trans_count == 1:
|
|
query_layer = input.view(batch_size, seq_length, num_heads, head_dim)
|
|
query_layer = query_layer.permute(0, 2, 1, 3)
|
|
key_layer = torch.zeros_like(query_layer)
|
|
value_layer = torch.zeros_like(query_layer)
|
|
return query_layer, key_layer, value_layer
|
|
|
|
qkv_layers = input.view(batch_size, seq_length, 3, num_heads, head_dim)
|
|
query_layer, key_layer, value_layer = qkv_layers[..., 0, :, :], qkv_layers[..., 1, :, :], qkv_layers[...,
|
|
2, :, :]
|
|
query_layer = query_layer.transpose(1, 2)
|
|
key_layer = key_layer.transpose(1, 2)
|
|
value_layer = value_layer.transpose(1, 2)
|
|
|
|
return query_layer, key_layer, value_layer
|
|
|
|
def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads,
|
|
num_kv, norm_factor, triangular_masking, local_attention, window_size, no_masking,
|
|
layer_id, num_layers, alibi, rope_theta, is_prompt, token_idx, position_ids):
|
|
bat_0213_query, bat_0213_key, bat_0213_value = self.bias_add_transform_0213(
|
|
query_key_value, None, heads, 3, False)
|
|
|
|
if rotary_dim > 0 and rotate_half:
|
|
from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
|
|
|
|
rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device)
|
|
cos, sin = rotary(bat_0213_value, InferenceContext.Instance().get_max_tokens_num())
|
|
bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin, position_ids)
|
|
|
|
bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt,
|
|
bat_0213_key, bat_0213_value)
|
|
|
|
bat_0213_key = self.repeat_kv(bat_0213_key, num_kv)
|
|
bat_0213_value = self.repeat_kv(bat_0213_value, num_kv)
|
|
|
|
bsz = query_key_value.shape[0]
|
|
head_dim = query_key_value.shape[2] // (heads * 3)
|
|
|
|
bmm_output = torch.bmm(bat_0213_query.reshape(bsz * heads, bat_0213_query.shape[2], head_dim),
|
|
bat_0213_key.reshape(bsz * heads, bat_0213_key.shape[2], head_dim).transpose(1, 2))
|
|
|
|
layer_scale = 1.0
|
|
if alibi is not None and len(alibi.shape) > 1:
|
|
layer_scale = max(1, layer_id).to(float)
|
|
|
|
alpha = norm_factor * norm_factor / layer_scale
|
|
bmm_output *= alpha
|
|
bmm_output_reshape = bmm_output.reshape(bsz, heads, bmm_output.shape[1], bmm_output.shape[2])
|
|
|
|
recompute = is_prompt
|
|
if attn_mask is not None and len(attn_mask.shape) > 1 and attn_mask.shape[-1] < bmm_output_reshape.shape[3]:
|
|
attn_mask = torch.nn.functional.pad(attn_mask, (0, bmm_output_reshape.shape[3] - attn_mask.shape[-1]),
|
|
value=torch.finfo(attn_mask.dtype).min)
|
|
softmax_output = SoftmaxOp.softmax_fallback(bmm_output_reshape, attn_mask, alibi, triangular_masking,
|
|
recompute, local_attention, window_size, None, layer_scale, 0, 1)
|
|
|
|
output = torch.bmm(softmax_output.reshape(bsz * heads, softmax_output.shape[2], softmax_output.shape[3]),
|
|
bat_0213_value.reshape(bsz * heads, bat_0213_value.shape[2], head_dim))
|
|
|
|
output = output.reshape(bsz, heads, output.shape[1], head_dim)
|
|
output = output.reshape(bsz, heads, output.shape[2] * head_dim)
|
|
input_seq_len = query_key_value.shape[1]
|
|
t4d_0123_output = self.transform4d_0213(output, input_seq_len)
|
|
t4d_0123_output = t4d_0123_output.reshape(bsz, t4d_0123_output.shape[1], heads * head_dim)
|
|
|
|
if layer_id == num_layers - 1:
|
|
InferenceContext.Instance().advance_tokens()
|
|
|
|
return t4d_0123_output, bat_0213_key, bat_0213_value
|
|
|
|
def forward(self, query_key_value: torch.Tensor, attn_mask: torch.Tensor, heads: int, num_kv: int,
|
|
norm_factor: float, no_masking: bool, layer_id: int, num_layers: int, alibi: torch.Tensor,
|
|
is_prompt: bool, token_idx: torch.Tensor, position_ids: torch.Tensor):
|
|
|
|
if alibi is not None:
|
|
batch_heads = query_key_value.shape[0] * heads
|
|
offset = dist.get_rank() * batch_heads if dist.is_initialized() else 0
|
|
alibi = alibi[offset:batch_heads + offset, :, :]
|
|
else:
|
|
alibi = torch.empty(1)
|
|
|
|
output = self.softmax_context_func(query_key_value, attn_mask, self.config.rotary_dim, self.config.rotate_half,
|
|
self.config.rotate_every_two, heads, num_kv, norm_factor,
|
|
self.config.triangular_masking, self.config.local_attention,
|
|
self.config.window_size, no_masking, layer_id, num_layers, alibi,
|
|
self.config.rope_theta, is_prompt, token_idx, position_ids)
|
|
|
|
return output
|