wehub-resource-sync
72 lines
3.0 KiB
Python
72 lines
3.0 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from ..config import DeepSpeedInferenceConfig
|
|
from .base import BaseOp
|
|
from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext
|
|
|
|
|
|
class SoftmaxOp(BaseOp):
|
|
|
|
def __init__(self, config: DeepSpeedInferenceConfig):
|
|
super(SoftmaxOp, self).__init__(config)
|
|
self.num_attention_heads_per_partition = config.heads // config.mp_size
|
|
try:
|
|
if self.config.dtype in [torch.float16, torch.int8]:
|
|
self.softmax_func = self.inference_module.softmax_fp16
|
|
elif self.config.dtype == torch.bfloat16:
|
|
self.softmax_func = self.inference_module.softmax_bf16
|
|
else:
|
|
self.softmax_func = self.inference_module.softmax_fp32
|
|
except AttributeError:
|
|
self.softmax_func = self.softmax_fallback
|
|
|
|
@staticmethod
|
|
def softmax_fallback(attn_scores, attn_mask, alibi, triangular, recompute, local_attention, window_size, async_op,
|
|
layer_scale, head_offset, mp_size):
|
|
scores_len = len(attn_scores.size())
|
|
heads = 1
|
|
if scores_len > 1:
|
|
heads = attn_scores.size()[1]
|
|
num_attention_heads_per_partition = heads // mp_size
|
|
|
|
if alibi is not None:
|
|
if len(alibi.shape) == 1:
|
|
alibi = None
|
|
else:
|
|
alibi = alibi[head_offset:head_offset + num_attention_heads_per_partition]
|
|
if attn_mask is not None and len(attn_mask.shape) == 1:
|
|
attn_mask = None
|
|
input_dtype = attn_scores.dtype
|
|
attn_scores *= layer_scale
|
|
|
|
if alibi is not None:
|
|
attn_scores += alibi
|
|
if attn_mask is not None:
|
|
# expand atten_mask from two dim into 4 dim, insert two dims in the middle
|
|
if len(attn_mask.shape) == 2:
|
|
attn_mask = attn_mask[:, None, None, :]
|
|
attn_scores += attn_mask
|
|
if triangular:
|
|
if attn_scores.shape[2] == 1: # query using kv cache
|
|
token_idx = InferenceContext.Instance().current_tokens()
|
|
tri = torch.arange(attn_scores.shape[2], device=attn_scores.device).ge(token_idx)
|
|
else:
|
|
tri = ~torch.tril(torch.ones(attn_scores.size(), device=attn_scores.device)).to(bool)
|
|
attn_scores = torch.masked_fill(attn_scores, tri, float('-inf'))
|
|
output = F.softmax(attn_scores, dim=-1, dtype=torch.float32).to(input_dtype)
|
|
|
|
return output
|
|
|
|
def forward(self, attn_scores: torch.Tensor, attn_mask: torch.Tensor, alibi: torch.Tensor, triangular: bool,
|
|
recompute: bool, local_attention: bool, window_size: int, async_op: bool, layer_scale: float,
|
|
head_offset: int):
|
|
output = self.softmax_func(attn_scores, attn_mask, alibi, triangular, recompute, local_attention, window_size,
|
|
async_op, layer_scale, head_offset, self.config.mp_size)
|
|
|
|
return output
|