wehub-resource-sync
308 lines
15 KiB
Python
308 lines
15 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from enum import IntEnum
|
|
from .builder import NPUOpBuilder
|
|
|
|
try:
|
|
import torch
|
|
import torch_npu
|
|
except ImportError as e:
|
|
pass
|
|
|
|
|
|
class ActivationFuncType(IntEnum):
|
|
UNKNOWN = 0
|
|
GELU = 1
|
|
ReLU = 2
|
|
GATED_GELU = 3
|
|
GATED_SILU = 4
|
|
|
|
|
|
class InferenceContext:
|
|
_workspace = None
|
|
|
|
_seed = 42
|
|
_curr_offset = 0
|
|
_stream = 0
|
|
_free_memory_size = 0
|
|
_num_tokens = 1
|
|
_attention_unfused_workspace_offset = 0
|
|
_workSpaceSize = 0
|
|
|
|
workSpaceSize = 0
|
|
kv_caches = None
|
|
|
|
@staticmethod
|
|
def reset_tokens(initial_tokens=1):
|
|
InferenceContext._num_tokens = initial_tokens
|
|
|
|
@staticmethod
|
|
def current_tokens():
|
|
return InferenceContext._num_tokens
|
|
|
|
@staticmethod
|
|
def GetWorkSpace():
|
|
return InferenceContext._workspace
|
|
|
|
|
|
class NPUInference:
|
|
|
|
@staticmethod
|
|
def layer_norm(inputs, gamma, beta, epsilon):
|
|
return torch.nn.functional.layer_norm(inputs, [inputs.shape[-1]], gamma, beta, eps=epsilon)
|
|
|
|
@staticmethod
|
|
def _qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose):
|
|
inp_norm = torch.nn.functional.layer_norm(inputs, (inputs.shape[2], ), gamma, beta, eps)
|
|
weight = weight.t() if transpose else weight
|
|
tmp = torch.matmul(inp_norm, weight)
|
|
if add_bias:
|
|
tmp += bias
|
|
output = [tmp, inp_norm]
|
|
return output
|
|
|
|
@staticmethod
|
|
def qkv_gemm_fp16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose):
|
|
return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose)
|
|
|
|
@staticmethod
|
|
def qkv_gemm_bf16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose):
|
|
return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose)
|
|
|
|
@staticmethod
|
|
def qkv_gemm_fp32(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose):
|
|
return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose)
|
|
|
|
@staticmethod
|
|
def _bias_add_transform_0213(vals, bias, hidden_dim, seq_length, seq_offset, heads, num_kv, rotary_dim,
|
|
rotate_half, rotate_every_two, rope_theta):
|
|
bsz, _, _ = vals.shape
|
|
q = vals[..., :hidden_dim].reshape(bsz, seq_length, heads, -1)
|
|
k = vals[..., hidden_dim:hidden_dim + num_kv * (hidden_dim // heads)].reshape(bsz, seq_length, num_kv, -1)
|
|
v = vals[..., hidden_dim + num_kv * (hidden_dim // heads):]
|
|
|
|
if rotary_dim > 0 and rotate_every_two:
|
|
# sin, cos may use cache
|
|
seq_id = torch.arange(0, seq_length).to("npu")
|
|
inv_freq = torch.arange(0, rotary_dim, 2) / rotary_dim
|
|
inv_freq = inv_freq.to("npu")
|
|
inv_freq = 1.0 / torch.pow(rope_theta, inv_freq)
|
|
inv_freq = torch.outer(seq_id, inv_freq)
|
|
sin = inv_freq.sin()
|
|
cos = inv_freq.cos()
|
|
# shape: [bsz=1, seq_len, heads=1, rotary_dim]
|
|
sin = sin.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1)
|
|
cos = cos.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1)
|
|
|
|
q_pos, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
|
|
k_pos, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
|
|
|
|
q_pos = torch_npu.npu_rotary_mul(q_pos, cos, sin)
|
|
q = torch.cat([q_pos, q_pass], dim=-1)
|
|
k_pos = torch_npu.npu_rotary_mul(k_pos, cos, sin)
|
|
k = torch.cat([k_pos, k_pass], dim=-1)
|
|
|
|
output = q.reshape(bsz, seq_length, -1).contiguous() # [b, s, H]
|
|
k_cache = k.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d]
|
|
v_cache = v.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d]
|
|
return output, k_cache, v_cache
|
|
|
|
@staticmethod
|
|
def _softmax_context(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):
|
|
bsz, seq_len, k = query_key_value.size()
|
|
k = k // (heads + 2 * (num_kv if num_kv > 0 else heads))
|
|
hidden_dim = heads * k
|
|
|
|
is_promt = seq_len > 1
|
|
if not InferenceContext.kv_caches:
|
|
InferenceContext.kv_caches = [[None, None] for _ in range(num_layers)]
|
|
if is_promt:
|
|
InferenceContext.reset_tokens(seq_len)
|
|
InferenceContext.kv_caches[layer_id] = [None, None]
|
|
|
|
soft_len = InferenceContext.current_tokens()
|
|
workspace = InferenceContext.GetWorkSpace()
|
|
seq_offset = 0 if is_promt else soft_len - 1
|
|
|
|
q, k, v = NPUInference._bias_add_transform_0213(vals=query_key_value,
|
|
bias=None,
|
|
hidden_dim=hidden_dim,
|
|
seq_length=seq_len,
|
|
seq_offset=seq_offset,
|
|
heads=heads,
|
|
num_kv=num_kv if num_kv > 0 else heads,
|
|
rotary_dim=rotary_dim,
|
|
rotate_half=rotate_half,
|
|
rotate_every_two=rotate_every_two,
|
|
rope_theta=rope_theta)
|
|
|
|
if not is_promt:
|
|
k_cache, v_cache = InferenceContext.kv_caches[layer_id]
|
|
if k_cache is not None:
|
|
k = torch.cat([k_cache, k], dim=2)
|
|
v = torch.cat([v_cache, v], dim=2)
|
|
InferenceContext.kv_caches[layer_id] = [k, v]
|
|
seq_len = k.shape[2]
|
|
|
|
layer_scale = max(1, layer_id) if len(alibi.size()) > 1 else 1.0
|
|
alpha = norm_factor * norm_factor / layer_scale
|
|
|
|
output = torch_npu.npu_fusion_attention(q,
|
|
k.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(),
|
|
v.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(),
|
|
heads,
|
|
"BSH",
|
|
pse=None,
|
|
padding_mask=None,
|
|
atten_mask=attn_mask.bool(),
|
|
scale=alpha,
|
|
pre_tockens=65536,
|
|
next_tockens=65536,
|
|
keep_prob=1,
|
|
inner_precise=0)[0]
|
|
|
|
return output, k, v
|
|
|
|
@staticmethod
|
|
def softmax_context_fp16(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):
|
|
return NPUInference._softmax_context(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)
|
|
|
|
@staticmethod
|
|
def softmax_context_bf16(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):
|
|
return NPUInference._softmax_context(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)
|
|
|
|
@staticmethod
|
|
def softmax_context_fp32(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):
|
|
return NPUInference._softmax_context(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)
|
|
|
|
@staticmethod
|
|
def _vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode):
|
|
if transposed_mode:
|
|
return torch.matmul(input, weight.t())
|
|
return torch.matmul(input, weight)
|
|
|
|
@staticmethod
|
|
def vector_matmul_fp16(input, weight, async_op, q_scale, q_int8, transposed_mode):
|
|
return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode)
|
|
|
|
@staticmethod
|
|
def vector_matmul_bf16(input, weight, async_op, q_scale, q_int8, transposed_mode):
|
|
return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode)
|
|
|
|
@staticmethod
|
|
def vector_matmul_fp32(input, weight, async_op, q_scale, q_int8, transposed_mode):
|
|
return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode)
|
|
|
|
@staticmethod
|
|
def _mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm,
|
|
mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose):
|
|
if mlp_after_attn:
|
|
residual_add = torch.nn.functional.layer_norm(input + residual + input_bias, (input.shape[-1], ), gamma,
|
|
beta, eps)
|
|
else:
|
|
residual_add = torch.nn.functional.layer_norm(input, (input.shape[-1], ), gamma, beta, eps)
|
|
|
|
weight_interm = weight_interm.t() if transpose else weight_interm
|
|
tmp = torch.matmul(residual_add, weight_interm)
|
|
if mlp_act_func_type == ActivationFuncType.GELU:
|
|
tmp = torch.nn.functional.gelu(tmp + bias)
|
|
elif mlp_act_func_type == ActivationFuncType.ReLU:
|
|
tmp = torch.nn.functional.relu(tmp + bias)
|
|
else:
|
|
raise Exception('Unsupported ActivationFuncType {}'.format(mlp_act_func_type))
|
|
output = torch.matmul(tmp, weight_out.t())
|
|
return output, residual_add
|
|
|
|
@staticmethod
|
|
def mlp_gemm_fp16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm,
|
|
mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose):
|
|
return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps,
|
|
pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype,
|
|
mlp_act_func_type, transpose)
|
|
|
|
@staticmethod
|
|
def mlp_gemm_bf16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm,
|
|
mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose):
|
|
return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps,
|
|
pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype,
|
|
mlp_act_func_type, transpose)
|
|
|
|
@staticmethod
|
|
def mlp_gemm_fp32(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm,
|
|
mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose):
|
|
return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps,
|
|
pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype,
|
|
mlp_act_func_type, transpose)
|
|
|
|
@staticmethod
|
|
def _residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size,
|
|
mlp_after_attn, add_bias, pre_layer_norm):
|
|
if mlp_after_attn:
|
|
if pre_layer_norm:
|
|
tmp = (residual.float() + attention_output.float() + attention_bias.float() +
|
|
final_bias.float()) / mp_size + hidden_state.float()
|
|
else:
|
|
tmp = residual.float() + hidden_state.float() + final_bias.float()
|
|
else:
|
|
if add_bias:
|
|
residual += attention_bias.float()
|
|
tmp = hidden_state.float() + attention_output.float() + (residual.float() + final_bias.float()) / mp_size
|
|
|
|
input_dtype = hidden_state.dtype
|
|
residual.set_(tmp.to(input_dtype))
|
|
|
|
@staticmethod
|
|
def residual_add_bias_fp16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size,
|
|
mlp_after_attn, add_bias, pre_layer_norm):
|
|
return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias,
|
|
mp_size, mlp_after_attn, add_bias, pre_layer_norm)
|
|
|
|
@staticmethod
|
|
def residual_add_bias_bf16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size,
|
|
mlp_after_attn, add_bias, pre_layer_norm):
|
|
return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias,
|
|
mp_size, mlp_after_attn, add_bias, pre_layer_norm)
|
|
|
|
@staticmethod
|
|
def residual_add_bias_fp32(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size,
|
|
mlp_after_attn, add_bias, pre_layer_norm):
|
|
return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias,
|
|
mp_size, mlp_after_attn, add_bias, pre_layer_norm)
|
|
|
|
|
|
class InferenceBuilder(NPUOpBuilder):
|
|
BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE"
|
|
NAME = "transformer_inference"
|
|
|
|
def __init__(self):
|
|
super().__init__(name=self.NAME)
|
|
|
|
def absolute_name(self):
|
|
return f'deepspeed.ops.transformer.inference.{self.NAME}_op'
|
|
|
|
def sources(self):
|
|
return []
|
|
|
|
def include_paths(self):
|
|
return []
|
|
|
|
def load(self, verbose=True):
|
|
return NPUInference
|