deepspeedai--deepspeed
wehub-resource-sync
124 行
3.7 KiB
Python
124 行
3.7 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import deepspeed.ops.transformer.inference.triton.matmul_ext as matmul_ext
|
|
from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp
|
|
from deepspeed.ops.transformer.inference.triton.layer_norm import layer_norm, layer_norm_residual
|
|
from deepspeed.utils.types import ActivationFuncType
|
|
|
|
|
|
def vector_matmul_func(input, weight, async_op, q_scale, q_int8, transposed_mode):
|
|
assert not transposed_mode and not async_op and not q_int8
|
|
return matmul_ext.matmul(input, weight, bias=None, activation="", use_triton=True)
|
|
|
|
|
|
def fused_gemm_gelu(input,
|
|
weight,
|
|
weight_scale,
|
|
bias,
|
|
weight_out,
|
|
weight_out_scale,
|
|
epsilon,
|
|
pre_layer_norm,
|
|
q_int8,
|
|
async_op,
|
|
transposed_mode,
|
|
use_triton_ln=True):
|
|
assert not transposed_mode
|
|
|
|
# activation
|
|
activation = "gelu"
|
|
|
|
# intermediate fc in FF
|
|
intm_out = matmul_ext.matmul(input, weight, bias=bias, activation=activation, use_triton=True)
|
|
|
|
# output fc in FF
|
|
ff_out = matmul_ext.matmul(
|
|
intm_out,
|
|
weight_out,
|
|
bias=None,
|
|
activation="", # bias added layer with residual_add + bias + layerNorm layer
|
|
use_triton=True)
|
|
return ff_out
|
|
|
|
|
|
def linear_func(input, weight, bias, add_bias, do_flash_attn, num_heads, transposed_mode=False):
|
|
assert not transposed_mode and not do_flash_attn
|
|
qkv_out = matmul_ext.matmul(input, weight, bias=(bias if add_bias else None), activation="", use_triton=True)
|
|
|
|
return qkv_out
|
|
|
|
|
|
def mlp_gemm_func(input,
|
|
residual,
|
|
input_bias,
|
|
weight_interm,
|
|
weight_out,
|
|
bias,
|
|
gamma,
|
|
beta,
|
|
epsilon,
|
|
pre_layer_norm,
|
|
mlp_after_attn,
|
|
weight_interm_scale,
|
|
weight_out_scale,
|
|
q_int8,
|
|
mlp_act_func_type,
|
|
transposed_mode,
|
|
use_triton_ln=True):
|
|
assert not transposed_mode
|
|
|
|
# residual add and layerNorm after attention
|
|
if use_triton_ln:
|
|
mlp_input = layer_norm_residual(input, input_bias, residual, gamma, beta, epsilon)
|
|
else:
|
|
mlp_input = LayerNormOp.layer_norm_residual(input, input_bias, residual, gamma, beta, epsilon)
|
|
|
|
# activation
|
|
if ActivationFuncType(mlp_act_func_type) == ActivationFuncType.GELU:
|
|
activation = "gelu"
|
|
elif ActivationFuncType(mlp_act_func_type) == ActivationFuncType.ReLU:
|
|
activation = "relu"
|
|
else:
|
|
activation = ""
|
|
|
|
# intermediate fc in FF
|
|
intm_out = matmul_ext.matmul(mlp_input, weight_interm, bias=bias, activation=activation, use_triton=True)
|
|
# output fc in FF
|
|
ff_out = matmul_ext.matmul(
|
|
intm_out,
|
|
weight_out,
|
|
bias=None,
|
|
activation="", # bias added layer with residual_add + bias + layerNorm layer
|
|
use_triton=True)
|
|
|
|
return ff_out, mlp_input
|
|
|
|
|
|
def qkv_gemm_func(
|
|
input,
|
|
weight,
|
|
q_scale,
|
|
bias,
|
|
gamma,
|
|
beta,
|
|
epsilon,
|
|
add_bias,
|
|
q_int8,
|
|
transposed_mode=False,
|
|
use_triton_ln=True,
|
|
):
|
|
|
|
assert not transposed_mode
|
|
# residual add and layerNorm after attention
|
|
if use_triton_ln:
|
|
qkv_input = layer_norm(input, gamma, beta, epsilon)
|
|
else:
|
|
qkv_input = LayerNormOp()(input, gamma, beta, epsilon)
|
|
|
|
qkv_out = matmul_ext.matmul(qkv_input, weight, bias=(bias if add_bias else None), activation="", use_triton=True)
|
|
|
|
return qkv_out, qkv_input
|