mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Support (Bias)SkipLayerNormalization fusion in GPT2 (#13988)
This commit is contained in:
Hariharan Seshadri
GitHub
parent
baba312e30
commit
7ed8bd4f95
22 changed files with 545 additions and 168 deletions
|
|
@ -3840,7 +3840,7 @@ This version of the operator has been available since version 1 of the 'com.micr
|
|||
<dd>1D bias tensor with shape (hidden_size</dd>
|
||||
</dl>
|
||||
|
||||
#### Outputs (1 - 3)
|
||||
#### Outputs (1 - 4)
|
||||
|
||||
<dl>
|
||||
<dt><tt>output</tt> : T</dt>
|
||||
|
|
@ -3849,6 +3849,8 @@ This version of the operator has been available since version 1 of the 'com.micr
|
|||
<dd>Saved mean used during training to speed up gradient computation</dd>
|
||||
<dt><tt>inv_std_var</tt> (optional) : U</dt>
|
||||
<dd>Saved inverse standard variance used during training to speed up gradient computation.</dd>
|
||||
<dt><tt>input_skip_sum</tt> (optional) : T</dt>
|
||||
<dd>Sum of the input and skip inputs with shape (batch_size, sequence_length, hidden_size).</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
|
|
|||
|
|
@ -438,7 +438,7 @@ Do not modify directly.*
|
|||
|QuickGelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
|
||||
|Range|*in* start:**T**<br> *in* limit:**T**<br> *in* delta:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64)|
|
||||
|SampleOp|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**|1+|**T** = tensor(double), tensor(float)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**<br> *out* input_skip_sum:**T**|1+|**T** = tensor(double), tensor(float)|
|
||||
|SparseToDenseMatMul|*in* A:**T**<br> *in* B:**T1**<br> *out* Y:**T1**|1+|**T** = sparse_tensor(double), sparse_tensor(float), sparse_tensor(int32), sparse_tensor(int64), sparse_tensor(uint32), sparse_tensor(uint64)<br/> **T1** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|
||||
|Tokenizer|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(string)|
|
||||
|TransposeMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
|
||||
|
|
@ -797,7 +797,7 @@ Do not modify directly.*
|
|||
|RemovePadding|*in* input:**T**<br> *in* sequence_token_count:**M**<br> *out* output:**T**<br> *out* token_offset:**M**<br> *out* cumulated_seq_len:**M**<br> *out* max_seq_len:**M**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|RestorePadding|*in* input:**T**<br> *in* token_offset:**M**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|Rfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**<br> *out* input_skip_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|TransposeMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|
||||
|Trilu|*in* X:**T**<br> *in* k:**tensor(int64)**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
| |
|
||||
|
|
@ -1135,7 +1135,7 @@ Do not modify directly.*
|
|||
|QLinearAdd|*in* A:**T**<br> *in* A_scale:**tensor(float)**<br> *in* A_zero_point:**T**<br> *in* B:**T**<br> *in* B_scale:**tensor(float)**<br> *in* B_zero_point:**T**<br> *in* C_scale:**tensor(float)**<br> *in* C_zero_point:**T**<br> *out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|
||||
|QLinearSigmoid|*in* X:**T**<br> *in* X_scale:**tensor(float)**<br> *in* X_zero_point:**T**<br> *in* Y_scale:**tensor(float)**<br> *in* Y_zero_point:**T**<br> *out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|
||||
|QuantizeLinear|*in* x:**T1**<br> *in* y_scale:**T1**<br> *in* y_zero_point:**T2**<br> *out* y:**T2**|1+|**T1** = tensor(float)<br/> **T2** = tensor(uint8)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**<br> *out* input_skip_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
| |
|
||||
| |
|
||||
|**Operator Domain:** *com.microsoft.dml*||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,9 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
|
|||
const Tensor* beta = p_ctx->Input<Tensor>(3);
|
||||
const Tensor* bias = p_ctx->Input<Tensor>(4);
|
||||
Tensor* output = p_ctx->Output(0, input->Shape());
|
||||
// For inferencing, we support one more optional output which is the sum
|
||||
// of the input and skip tensors
|
||||
Tensor* skip_input_add_output = p_ctx->Output(3, input->Shape());
|
||||
|
||||
const auto& input_dims = input->Shape().GetDims();
|
||||
if (input_dims.size() != 3) {
|
||||
|
|
@ -98,18 +101,28 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
|
|||
|
||||
T* output_data = output->MutableData<T>();
|
||||
|
||||
// For inferencing, we support one more optional output which is the sum
|
||||
// of the input and skip tensors
|
||||
T* skip_input_add_output_data = skip_input_add_output != nullptr ? skip_input_add_output->MutableData<T>() : nullptr;
|
||||
|
||||
concurrency::ThreadPool::TryBatchParallelFor(
|
||||
p_ctx->GetOperatorThreadPool(), static_cast<int32_t>(task_count),
|
||||
[&](ptrdiff_t task_idx) {
|
||||
const T* p_input = input_data + task_idx * hidden_size;
|
||||
const T* p_skip = skip_data + task_idx * hidden_size;
|
||||
T* p_output = output_data + task_idx * hidden_size;
|
||||
auto offset = task_idx * hidden_size;
|
||||
|
||||
const T* p_input = input_data + offset;
|
||||
const T* p_skip = skip_data + offset;
|
||||
T* p_output = output_data + offset;
|
||||
T* p_skip_input_add_output_data = skip_input_add_output_data != nullptr ? skip_input_add_output_data + offset : nullptr;
|
||||
|
||||
T mean = 0;
|
||||
T mean_square = 0;
|
||||
|
||||
for (int64_t h = 0; h < hidden_size; h++) {
|
||||
T value = p_input[h] + p_skip[h];
|
||||
if (nullptr != p_skip_input_add_output_data) {
|
||||
p_skip_input_add_output_data[h] = value;
|
||||
}
|
||||
if (nullptr != bias_data) {
|
||||
value += bias_data[h];
|
||||
}
|
||||
|
|
|
|||
|
|
@ -41,6 +41,10 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
|
|||
|
||||
Tensor* output = ctx->Output(0, input->Shape());
|
||||
|
||||
// For inferencing, we support one more optional output which is the sum
|
||||
// of the input and skip tensors
|
||||
Tensor* skip_input_add_output = ctx->Output(3, input->Shape());
|
||||
|
||||
if (input->Shape() != skip->Shape()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
"skip is expected to have same shape as input");
|
||||
|
|
@ -99,6 +103,7 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
|
|||
return LaunchSkipLayerNormKernel<CudaT>(
|
||||
Stream(ctx),
|
||||
reinterpret_cast<CudaT*>(output->MutableData<T>()),
|
||||
skip_input_add_output != nullptr ? reinterpret_cast<CudaT*>(skip_input_add_output->MutableData<T>()) : nullptr,
|
||||
reinterpret_cast<const CudaT*>(input->Data<T>()),
|
||||
reinterpret_cast<const CudaT*>(skip->Data<T>()),
|
||||
reinterpret_cast<const CudaT*>(gamma->Data<T>()),
|
||||
|
|
|
|||
|
|
@ -50,8 +50,9 @@ half maybe2half(float x) {
|
|||
|
||||
template <typename T, unsigned TPB>
|
||||
__global__ void SkipLayerNormKernel(
|
||||
const int ld, const T* input, const T* skip, const T* beta, const T* gamma, const T* bias,
|
||||
const T epsilon, T* output) {
|
||||
const int ld, const T* input, const T* skip,
|
||||
const T* beta, const T* gamma, const T* bias,
|
||||
const T epsilon, T* output, T* skip_input_add_output) {
|
||||
const T reverse_ld = T(1.f / ld);
|
||||
const int offset = blockIdx.x * ld;
|
||||
|
||||
|
|
@ -61,6 +62,11 @@ __global__ void SkipLayerNormKernel(
|
|||
|
||||
for (int i = threadIdx.x; i < ld; i += TPB) {
|
||||
const int idx = offset + i;
|
||||
|
||||
if (skip_input_add_output != nullptr) {
|
||||
skip_input_add_output[idx] = input[idx] + skip[idx];
|
||||
}
|
||||
|
||||
const T val = (bias == nullptr) ? input[idx] + skip[idx] : input[idx] + skip[idx] + bias[i];
|
||||
const T rldval = reverse_ld * val;
|
||||
thread_data = pair_sum(thread_data, cub::KeyValuePair<T, T>(rldval, rldval * val));
|
||||
|
|
@ -74,13 +80,14 @@ __global__ void SkipLayerNormKernel(
|
|||
template <typename T, unsigned TPB, int ILP>
|
||||
__global__ void SkipLayerNormKernelSmall(
|
||||
const int ld, const T* input, const T* skip, const T* beta, const T* gamma,
|
||||
const T* bias, const T epsilon, T* output, bool hasBias) {
|
||||
const T* bias, const T epsilon, T* output, T* skip_input_add_output,
|
||||
bool hasBias, bool hasSkipInputAdditionOutput) {
|
||||
const T rld = T(1.f / ld);
|
||||
const int idx = blockIdx.x * ld + threadIdx.x * ILP; // grid_size = n / ld
|
||||
|
||||
using VecT = aligned_vector<T, ILP>;
|
||||
|
||||
T input_v[ILP], skip_v[ILP], bias_v[ILP];
|
||||
T input_v[ILP], skip_v[ILP], bias_v[ILP], skip_input_add_output_v[ILP];
|
||||
|
||||
VecT* input_val = reinterpret_cast<VecT*>(&input_v);
|
||||
*input_val = *reinterpret_cast<const VecT*>(&input[idx]);
|
||||
|
|
@ -100,11 +107,20 @@ __global__ void SkipLayerNormKernelSmall(
|
|||
T rldvalsq_sum = T(0.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < ILP; i++) {
|
||||
if (hasSkipInputAdditionOutput) {
|
||||
skip_input_add_output_v[i] = input_v[i] + skip_v[i];
|
||||
}
|
||||
|
||||
input_v[i] += hasBias ? skip_v[i] + bias_v[i] : skip_v[i];
|
||||
const T rldval = rld * input_v[i];
|
||||
rldval_sum += rldval;
|
||||
rldvalsq_sum += rldval * input_v[i];
|
||||
}
|
||||
|
||||
if (hasSkipInputAdditionOutput) {
|
||||
*(reinterpret_cast<VecT*>(&skip_input_add_output[idx])) = *reinterpret_cast<VecT*>(&skip_input_add_output_v);
|
||||
}
|
||||
|
||||
thread_data = cub::KeyValuePair<T, T>(rldval_sum, rldvalsq_sum);
|
||||
}
|
||||
LayerNormSmall<T, TPB, ILP>(input_v, thread_data, ld, idx, beta, gamma, epsilon, output);
|
||||
|
|
@ -112,49 +128,57 @@ __global__ void SkipLayerNormKernelSmall(
|
|||
|
||||
template <typename T>
|
||||
Status LaunchSkipLayerNormKernel(
|
||||
cudaStream_t stream, T* output, const T* input, const T* skip, const T* gamma,
|
||||
cudaStream_t stream, T* output, T* skip_input_add_output, const T* input, const T* skip, const T* gamma,
|
||||
const T* beta, const T* bias, float epsilon, const int ld, const int element_count,
|
||||
size_t element_size) {
|
||||
// this must be true because n is the total size of the tensor
|
||||
assert(element_count % ld == 0);
|
||||
bool hasBias = (bias == nullptr) ? false : true;
|
||||
bool hasSkipInputAdditionOutput = (skip_input_add_output == nullptr) ? false : true;
|
||||
|
||||
if (0 == (ld % 4)) {
|
||||
const int grid_size = element_count / ld;
|
||||
if (ld <= 32) {
|
||||
constexpr int block_size = 32;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 64) {
|
||||
constexpr int block_size = 64 / 2;
|
||||
SkipLayerNormKernelSmall<T, block_size, 2>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 128) {
|
||||
constexpr int block_size = 128 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 384) {
|
||||
constexpr int block_size = 384 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 768) {
|
||||
constexpr int block_size = 768 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 1024) {
|
||||
constexpr int block_size = 1024 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else {
|
||||
constexpr int block_size = 256;
|
||||
SkipLayerNormKernel<T, block_size>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output);
|
||||
maybe2half<T>(epsilon), output, skip_input_add_output);
|
||||
}
|
||||
} else {
|
||||
const int grid_size = element_count / ld;
|
||||
|
|
@ -162,41 +186,47 @@ Status LaunchSkipLayerNormKernel(
|
|||
constexpr int block_size = 32;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 64) {
|
||||
constexpr int block_size = 64;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld <= 128) {
|
||||
constexpr int block_size = 128;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else if (ld == 384) {
|
||||
constexpr int block_size = 384;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, hasBias);
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
} else {
|
||||
constexpr int block_size = 256;
|
||||
SkipLayerNormKernel<T, block_size>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output);
|
||||
maybe2half<T>(epsilon), output, skip_input_add_output);
|
||||
}
|
||||
}
|
||||
return CUDA_CALL(cudaGetLastError());
|
||||
}
|
||||
|
||||
template Status LaunchSkipLayerNormKernel<float>(cudaStream_t stream, float* output, const float* input,
|
||||
const float* skip, const float* gamma, const float* beta,
|
||||
const float* bias, float epsilon, const int ld,
|
||||
const int element_count, size_t element_size);
|
||||
template Status LaunchSkipLayerNormKernel<float>(cudaStream_t stream, float* output, float* skip_input_add_output,
|
||||
const float* input,
|
||||
const float* skip, const float* gamma, const float* beta,
|
||||
const float* bias, float epsilon, const int ld,
|
||||
const int element_count, size_t element_size);
|
||||
|
||||
template Status LaunchSkipLayerNormKernel<half>(cudaStream_t stream, half* output, const half* input,
|
||||
const half* skip, const half* gamma, const half* beta,
|
||||
const half* bias, float epsilon, const int ld,
|
||||
const int element_count, size_t element_size);
|
||||
template Status LaunchSkipLayerNormKernel<half>(cudaStream_t stream, half* output, half* skip_input_add_output,
|
||||
const half* input,
|
||||
const half* skip, const half* gamma, const half* beta,
|
||||
const half* bias, float epsilon, const int ld,
|
||||
const int element_count, size_t element_size);
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
|
|
|
|||
|
|
@ -11,15 +11,16 @@ namespace cuda {
|
|||
template <typename T>
|
||||
Status LaunchSkipLayerNormKernel(
|
||||
cudaStream_t stream,
|
||||
T* output, // output tensor
|
||||
const T* input, // input tensor
|
||||
const T* skip, // skip tensor
|
||||
const T* gamma, // Layer normalization gamma tensor
|
||||
const T* beta, // Layer normalization beta tensor
|
||||
const T* bias, // Layer normalization beta tensor
|
||||
float epsilon, // Layer normalization epsilon
|
||||
int hidden_size, // hidden size, it is the leading dimension (ld)
|
||||
int element_count, // number of elements in input tensor
|
||||
T* output, // normalized output tensor
|
||||
T* skip_input_add_output, // sum of the input and skip tensors output
|
||||
const T* input, // input tensor
|
||||
const T* skip, // skip tensor
|
||||
const T* gamma, // Layer normalization gamma tensor
|
||||
const T* beta, // Layer normalization beta tensor
|
||||
const T* bias, // Layer normalization beta tensor
|
||||
float epsilon, // Layer normalization epsilon
|
||||
int hidden_size, // hidden size, it is the leading dimension (ld)
|
||||
int element_count, // number of elements in input tensor
|
||||
size_t element_size);
|
||||
|
||||
} // namespace cuda
|
||||
|
|
|
|||
|
|
@ -110,12 +110,12 @@ void RestorePaddingTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx)
|
|||
auto& token_offset_shape = getInputShape(ctx, 1);
|
||||
|
||||
if (input_shape.dim().size() != 2) {
|
||||
fail_shape_inference("input shall be 2 dimensions");
|
||||
}
|
||||
fail_shape_inference("input shall be 2 dimensions");
|
||||
}
|
||||
|
||||
if (token_offset_shape.dim().size() != 2) {
|
||||
fail_shape_inference("token_offset shall be 2 dimensions");
|
||||
}
|
||||
fail_shape_inference("token_offset shall be 2 dimensions");
|
||||
}
|
||||
|
||||
ONNX_NAMESPACE::TensorShapeProto output_shape;
|
||||
*output_shape.add_dim() = token_offset_shape.dim(0);
|
||||
|
|
@ -125,7 +125,6 @@ void RestorePaddingTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx)
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
constexpr const char* Attention_ver1_doc = R"DOC(
|
||||
Multi-Head Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT).
|
||||
|
||||
|
|
@ -390,6 +389,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(SkipLayerNormalization, 1,
|
|||
.Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T")
|
||||
.Output(1, "mean", "Saved mean used during training to speed up gradient computation", "U", OpSchema::Optional)
|
||||
.Output(2, "inv_std_var", "Saved inverse standard variance used during training to speed up gradient computation.", "U", OpSchema::Optional)
|
||||
.Output(3, "input_skip_sum", "Sum of the input and skip inputs with shape (batch_size, sequence_length, hidden_size).", "T", OpSchema::Optional)
|
||||
.TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.")
|
||||
.TypeConstraint("U", {"tensor(float)"}, "Constrain mean and inv_std_var to float tensors.")
|
||||
.TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput));
|
||||
|
|
@ -508,7 +508,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
|
|||
RemovePaddingTypeAndShapeInference(ctx);
|
||||
}));
|
||||
|
||||
|
||||
constexpr const char* RestorePadding_ver1_doc = R"DOC(
|
||||
Restore paddings and fill padding with zeros.
|
||||
|
||||
|
|
|
|||
|
|
@ -537,8 +537,6 @@ class FusionEmbedLayerNoMask(Fusion):
|
|||
if two_gather is None:
|
||||
return False
|
||||
|
||||
add_output = add_before_layernorm.output[0]
|
||||
|
||||
word_embedding_gather, position_embedding_gather = two_gather
|
||||
input_ids = word_embedding_gather.input[1]
|
||||
position_ids = position_embedding_gather.input[1]
|
||||
|
|
@ -549,9 +547,22 @@ class FusionEmbedLayerNoMask(Fusion):
|
|||
if not self.check_embedding(word_embedding_gather, None, position_embedding_gather):
|
||||
return False
|
||||
|
||||
# If the add_before_layernorm node is an Add node, then the add_output output is the first index
|
||||
# output of this node.
|
||||
|
||||
# If the add_before_layernorm node is SkipLayerNormalization node, then the add_output output
|
||||
# is the (optional) fourth index output of this node.
|
||||
add_output = None
|
||||
optional_embedding_sum_output = False
|
||||
if self.is_embedding_sum_needed(add_before_layernorm):
|
||||
if (add_before_layernorm.op_type == "Add" and self.is_embedding_sum_needed(add_before_layernorm)) or (
|
||||
add_before_layernorm.op_type == "SkipLayerNormalization" and len(add_before_layernorm.output) >= 4
|
||||
):
|
||||
optional_embedding_sum_output = True
|
||||
add_output = (
|
||||
add_before_layernorm.output[0]
|
||||
if add_before_layernorm.op_type == "Add"
|
||||
else add_before_layernorm.output[3]
|
||||
)
|
||||
|
||||
# make the fused node
|
||||
embed_node = self.create_fused_node(
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ class FusionGptAttentionPastBase(Fusion):
|
|||
"""Base class for GPT Attention Fusion with past state"""
|
||||
|
||||
def __init__(self, model: OnnxModel, num_heads: int):
|
||||
super().__init__(model, "Attention", "LayerNormalization", "with past")
|
||||
super().__init__(model, "Attention", ["LayerNormalization", "SkipLayerNormalization"], "with past")
|
||||
self.num_heads = num_heads
|
||||
self.utils = FusionUtils(model)
|
||||
self.casted_attention_mask = {} # map from name of attention mask to the name that casted to int32
|
||||
|
|
@ -224,26 +224,52 @@ class FusionGptAttention(FusionGptAttentionPastBase):
|
|||
past = None
|
||||
present = None
|
||||
return_indice = []
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
|
||||
[0, None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice,
|
||||
) # yapf: disable
|
||||
|
||||
is_normalize_node_skiplayernorm = normalize_node.op_type == "SkipLayerNormalization"
|
||||
qkv_nodes = None
|
||||
|
||||
if not is_normalize_node_skiplayernorm:
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
|
||||
[0, None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice,
|
||||
) # yapf: disable
|
||||
else:
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
|
||||
[None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice,
|
||||
) # yapf: disable
|
||||
|
||||
if qkv_nodes is None:
|
||||
return
|
||||
(
|
||||
add_qkv,
|
||||
reshape_qkv,
|
||||
gemm_qkv,
|
||||
reshape_1,
|
||||
reshape_2,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
another_input = add_qkv.input[1 - return_indice[0]]
|
||||
another_input = None
|
||||
if not is_normalize_node_skiplayernorm:
|
||||
(
|
||||
add_qkv,
|
||||
reshape_qkv,
|
||||
gemm_qkv,
|
||||
reshape_1,
|
||||
reshape_2,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
another_input = add_qkv.input[1 - return_indice[0]]
|
||||
else:
|
||||
(
|
||||
reshape_qkv,
|
||||
gemm_qkv,
|
||||
reshape_1,
|
||||
reshape_2,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
v_nodes = self.model.match_parent_path(matmul_qkv, ["Concat", "Transpose", "Reshape", "Split"], [1, 1, 0, 0])
|
||||
if v_nodes is None:
|
||||
|
|
@ -251,22 +277,46 @@ class FusionGptAttention(FusionGptAttentionPastBase):
|
|||
return
|
||||
(concat_v, transpose_v, reshape_v, split_fc) = v_nodes
|
||||
|
||||
# Try match pattern using Gemm + LayerNormalization
|
||||
fc_nodes = self.model.match_parent_path(
|
||||
split_fc,
|
||||
["Reshape", "Gemm", "Reshape", "LayerNormalization"],
|
||||
[0, 0, 0, 0],
|
||||
output_name_to_node,
|
||||
)
|
||||
|
||||
# Try match pattern using Gemm + SkipLayerNormalization
|
||||
if fc_nodes is None:
|
||||
fc_nodes = self.model.match_parent_path(
|
||||
split_fc,
|
||||
["Reshape", "Gemm", "Reshape", "SkipLayerNormalization"],
|
||||
[0, 0, 0, 0],
|
||||
output_name_to_node,
|
||||
)
|
||||
|
||||
# Try match pattern using MatMul
|
||||
if fc_nodes is None:
|
||||
# LayerNormalization
|
||||
fc_nodes = self.model.match_parent_path(
|
||||
split_fc,
|
||||
["Add", "MatMul", "LayerNormalization"],
|
||||
[0, None, 0],
|
||||
output_name_to_node,
|
||||
)
|
||||
|
||||
# SkipLayerNormalization
|
||||
if fc_nodes is None:
|
||||
fc_nodes = self.model.match_parent_path(
|
||||
split_fc,
|
||||
["Add", "MatMul", "SkipLayerNormalization"],
|
||||
[0, None, 0],
|
||||
output_name_to_node,
|
||||
)
|
||||
|
||||
if fc_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match fc path")
|
||||
return
|
||||
|
||||
fc_weight = fc_nodes[1].input[1]
|
||||
i, _ = self.model.get_constant_input(fc_nodes[0])
|
||||
fc_bias = fc_nodes[0].input[i]
|
||||
|
|
@ -276,8 +326,13 @@ class FusionGptAttention(FusionGptAttentionPastBase):
|
|||
|
||||
layernorm_before_attention = fc_nodes[-1]
|
||||
|
||||
if not another_input in layernorm_before_attention.input:
|
||||
logger.debug("Add and LayerNormalization shall have one same input")
|
||||
# `another_input` will be non-None only if
|
||||
# (1) SkipLayerNorm fusion wasn't turned ON
|
||||
# (2) SkipLayerNorm fusion was turned ON but upstream layer's LayerNorm + Add was not
|
||||
# fused into a SkipLayerNorm. This can happen if the shapes to the Add node are different.
|
||||
# So, keep the following check if SkipLayerNorm fusion is turned ON or OFF.
|
||||
if another_input is not None and not another_input in layernorm_before_attention.input:
|
||||
logger.debug("Upstream Add and (Skip)LayerNormalization shall have one same input")
|
||||
return
|
||||
|
||||
is_unidirectional = True
|
||||
|
|
|
|||
|
|
@ -144,6 +144,14 @@ class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
|
|||
["Unsqueeze", "Sub", "Gather", "Shape", "LayerNormalization"],
|
||||
[1, 0, 1, 0, 0],
|
||||
)
|
||||
|
||||
if first_slice_sub_1 is None:
|
||||
first_slice_sub_1 = self.model.match_parent_path(
|
||||
slice_mask,
|
||||
["Unsqueeze", "Sub", "Gather", "Shape", "SkipLayerNormalization"],
|
||||
[1, 0, 1, 0, 0],
|
||||
)
|
||||
|
||||
if first_slice_sub_1 is None or first_slice_sub_1[-1] != layernorm_before_attention:
|
||||
logger.debug("fuse_attention: failed to match last slice sub path 1")
|
||||
return None
|
||||
|
|
@ -154,24 +162,49 @@ class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
|
|||
past = None
|
||||
present = None
|
||||
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "Add", "MatMul", "Reshape", "Transpose", "MatMul"],
|
||||
[0, 1, None, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
) # yapf: disable
|
||||
is_normalize_node_skiplayernorm = normalize_node.op_type == "SkipLayerNormalization"
|
||||
qkv_nodes = None
|
||||
|
||||
if not is_normalize_node_skiplayernorm:
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "Add", "MatMul", "Reshape", "Transpose", "MatMul"],
|
||||
[0, 1, None, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
) # yapf: disable
|
||||
else:
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "MatMul", "Reshape", "Transpose", "MatMul"],
|
||||
[1, None, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
) # yapf: disable
|
||||
|
||||
if qkv_nodes is None:
|
||||
return
|
||||
(
|
||||
add_skip,
|
||||
add_after_attention,
|
||||
matmul_after_attention,
|
||||
reshape_qkv,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
skip_input = add_skip.input[0]
|
||||
skip_input = None
|
||||
if not is_normalize_node_skiplayernorm:
|
||||
(
|
||||
add_skip,
|
||||
add_after_attention,
|
||||
matmul_after_attention,
|
||||
reshape_qkv,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
skip_input = add_skip.input[0]
|
||||
else:
|
||||
(
|
||||
add_after_attention,
|
||||
matmul_after_attention,
|
||||
reshape_qkv,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
skip_input = normalize_node.input[0]
|
||||
|
||||
v_nodes = self.model.match_parent_path(
|
||||
matmul_qkv,
|
||||
|
|
@ -186,6 +219,22 @@ class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
|
|||
],
|
||||
[1, 1, 0, 0, 0, None, 0],
|
||||
) # yapf: disable
|
||||
|
||||
if v_nodes is None:
|
||||
v_nodes = self.model.match_parent_path(
|
||||
matmul_qkv,
|
||||
[
|
||||
"Concat",
|
||||
"Transpose",
|
||||
"Reshape",
|
||||
"Split",
|
||||
"Add",
|
||||
"MatMul",
|
||||
"SkipLayerNormalization",
|
||||
],
|
||||
[1, 1, 0, 0, 0, None, 0],
|
||||
) # yapf: disable
|
||||
|
||||
if v_nodes is None:
|
||||
logger.debug("fuse_attention: failed to match v path")
|
||||
return
|
||||
|
|
@ -198,7 +247,18 @@ class FusionGptAttentionMegatron(FusionGptAttentionPastBase):
|
|||
matmul_before_split,
|
||||
layernorm_before_attention,
|
||||
) = v_nodes
|
||||
if skip_input != layernorm_before_attention.input[0]:
|
||||
|
||||
if (
|
||||
layernorm_before_attention.op_type == "LayerNormalization"
|
||||
and skip_input != layernorm_before_attention.input[0]
|
||||
):
|
||||
logger.debug("fuse_attention: skip_input != layernorm_before_attention.input[0]")
|
||||
return
|
||||
|
||||
if (
|
||||
layernorm_before_attention.op_type == "SkipLayerNormalization"
|
||||
and skip_input != layernorm_before_attention.output[3]
|
||||
):
|
||||
logger.debug("fuse_attention: skip_input != layernorm_before_attention.input[0]")
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ class FusionGptAttentionNoPast(Fusion):
|
|||
"""
|
||||
|
||||
def __init__(self, model: OnnxModel, num_heads: int):
|
||||
super().__init__(model, "Attention", "LayerNormalization", "without past")
|
||||
super().__init__(model, "Attention", ["LayerNormalization", "SkipLayerNormalization"], "without past")
|
||||
# TODO: detect num_heads from graph like FusionAttention
|
||||
self.num_heads = num_heads
|
||||
|
||||
|
|
@ -60,27 +60,56 @@ class FusionGptAttentionNoPast(Fusion):
|
|||
self.node_name_to_graph_name[add_node.name] = self.this_graph_name
|
||||
|
||||
def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
|
||||
# (TODO) hasesh/tlwu: Investigate what fixes the following logic needs in order
|
||||
# to fuse the Attention sub-graph. With some changes to other fusions, this stopped
|
||||
# working.
|
||||
return_indice = []
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
|
||||
[0, None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice,
|
||||
) # yapf: disable
|
||||
|
||||
is_normalize_node_skiplayernorm = normalize_node.op_type == "SkipLayerNormalization"
|
||||
qkv_nodes = None
|
||||
|
||||
if not is_normalize_node_skiplayernorm:
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Add", "Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
|
||||
[0, None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice,
|
||||
) # yapf: disable
|
||||
else:
|
||||
qkv_nodes = self.model.match_parent_path(
|
||||
normalize_node,
|
||||
["Reshape", "Gemm", "Reshape", "Reshape", "Transpose", "MatMul"],
|
||||
[None, 0, 0, 0, 0, 0],
|
||||
output_name_to_node=output_name_to_node,
|
||||
return_indice=return_indice,
|
||||
) # yapf: disable
|
||||
|
||||
if qkv_nodes is None:
|
||||
return
|
||||
(
|
||||
add_qkv,
|
||||
reshape_qkv,
|
||||
gemm_qkv,
|
||||
reshape_1,
|
||||
reshape_2,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
another_input = add_qkv.input[1 - return_indice[0]]
|
||||
another_input = None
|
||||
if not is_normalize_node_skiplayernorm:
|
||||
(
|
||||
add_qkv,
|
||||
reshape_qkv,
|
||||
gemm_qkv,
|
||||
reshape_1,
|
||||
reshape_2,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
another_input = add_qkv.input[1 - return_indice[0]]
|
||||
else:
|
||||
(
|
||||
reshape_qkv,
|
||||
gemm_qkv,
|
||||
reshape_1,
|
||||
reshape_2,
|
||||
transpose_qkv,
|
||||
matmul_qkv,
|
||||
) = qkv_nodes
|
||||
|
||||
v_nodes = self.model.match_parent_path(
|
||||
matmul_qkv,
|
||||
|
|
@ -100,16 +129,25 @@ class FusionGptAttentionNoPast(Fusion):
|
|||
) = v_nodes
|
||||
|
||||
layernorm_before_attention = self.model.get_parent(reshape_before_gemm, 0, output_name_to_node)
|
||||
if layernorm_before_attention is None or layernorm_before_attention.op_type != "LayerNormalization":
|
||||
if layernorm_before_attention is None or (
|
||||
layernorm_before_attention.op_type != "LayerNormalization"
|
||||
and layernorm_before_attention.op_type != "SkipLayerNormalization"
|
||||
):
|
||||
if layernorm_before_attention.op_type != "Add":
|
||||
logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}")
|
||||
logger.debug(f"failed to get (skip)layernorm before gemm. Got {layernorm_before_attention.op_type}")
|
||||
return
|
||||
|
||||
if not another_input in layernorm_before_attention.input:
|
||||
# match openai-gpt
|
||||
if not another_input in layernorm_before_attention.output:
|
||||
logger.debug("Add and LayerNormalization shall have one same input")
|
||||
return
|
||||
# `another_input` will be non-None only if
|
||||
# (1) SkipLayerNorm fusion wasn't turned ON
|
||||
# (2) SkipLayerNorm fusion was turned ON but upstream layer's LayerNorm + Add was not
|
||||
# fused into a SkipLayerNorm. This can happen if the shapes to the Add node are different.
|
||||
# So, keep the following check if SkipLayerNorm fusion is turned ON or OFF.
|
||||
if another_input is not None:
|
||||
if not another_input in layernorm_before_attention.input:
|
||||
# match openai-gpt
|
||||
if not another_input in layernorm_before_attention.output:
|
||||
logger.debug("Add and (Skip)LayerNormalization shall have one same input")
|
||||
return
|
||||
|
||||
qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Sub", "Mul", "Div", "MatMul"], [0, 0, 0, 0, 0])
|
||||
if qk_nodes is not None:
|
||||
|
|
|
|||
|
|
@ -30,9 +30,6 @@ class FusionOptions:
|
|||
self.enable_gemm_fast_gelu = False
|
||||
self.attention_mask_format = AttentionMaskFormat.AttentionMask
|
||||
|
||||
if model_type == "gpt2":
|
||||
self.enable_skip_layer_norm = False
|
||||
|
||||
def use_raw_attention_mask(self, use_raw_mask=True):
|
||||
if use_raw_mask:
|
||||
self.attention_mask_format = AttentionMaskFormat.AttentionMask
|
||||
|
|
|
|||
|
|
@ -61,10 +61,31 @@ class FusionSkipLayerNormalization(Fusion):
|
|||
if self.model.match_parent_path(gather_path[0], ["ConstantOfShape"], [1]) is None:
|
||||
return
|
||||
|
||||
residual_add_has_multiple_consumers = False
|
||||
add_children = self.model.get_children(add, input_name_to_nodes)
|
||||
|
||||
# This means that the residual Add before the LayerNormalization produces an output
|
||||
# that is consumed by some other nodes other than the LayerNormalization itself
|
||||
# We can still go ahead with the SkipLayerNormalization fusion but we need to
|
||||
# preserve the output of Add and that needs to be produced by SkipLayerNormalization.
|
||||
if len(add_children) != 1:
|
||||
residual_add_has_multiple_consumers = True
|
||||
|
||||
outputs_to_keep = node.output
|
||||
|
||||
if residual_add_has_multiple_consumers:
|
||||
outputs_to_keep.extend([add.output[0]])
|
||||
|
||||
outputs = [node.output[0]]
|
||||
|
||||
# Skip the other optional outputs of SkipLayerNormalization before adding the Add's output
|
||||
if residual_add_has_multiple_consumers:
|
||||
outputs.extend(["", "", add.output[0]])
|
||||
|
||||
if (
|
||||
add is not None
|
||||
and add.op_type == "Add"
|
||||
and self.model.is_safe_to_fuse_nodes([add, node], node.output, input_name_to_nodes, output_name_to_node)
|
||||
and self.model.is_safe_to_fuse_nodes([add, node], outputs_to_keep, input_name_to_nodes, output_name_to_node)
|
||||
):
|
||||
self.nodes_to_remove.extend([add, node])
|
||||
|
||||
|
|
@ -72,7 +93,7 @@ class FusionSkipLayerNormalization(Fusion):
|
|||
normalize_node = helper.make_node(
|
||||
"SkipLayerNormalization",
|
||||
inputs=inputs,
|
||||
outputs=[node.output[0]],
|
||||
outputs=outputs,
|
||||
name=self.model.create_node_name("SkipLayerNormalization", name_prefix="SkipLayerNorm"),
|
||||
)
|
||||
normalize_node.domain = "com.microsoft"
|
||||
|
|
@ -101,8 +122,14 @@ class FusionBiasSkipLayerNormalization(Fusion):
|
|||
return_indice = []
|
||||
nodes = self.model.match_parent_path(node, ["Add", "MatMul"], [None, None], None, return_indice)
|
||||
if nodes is None:
|
||||
return
|
||||
assert len(return_indice) == 2
|
||||
# In case of fp16, we could have a Cast between the MatMul and the bias Add
|
||||
nodes = self.model.match_parent_path(
|
||||
node, ["Add", "Cast", "MatMul"], [None, None, None], None, return_indice
|
||||
)
|
||||
if nodes is None:
|
||||
return
|
||||
|
||||
assert len(return_indice) == 2 or len(return_indice) == 3
|
||||
add_input_index = return_indice[0]
|
||||
if add_input_index >= 2:
|
||||
return
|
||||
|
|
@ -126,9 +153,7 @@ class FusionBiasSkipLayerNormalization(Fusion):
|
|||
return
|
||||
|
||||
subgraph_nodes = [node, add]
|
||||
if not self.model.is_safe_to_fuse_nodes(
|
||||
subgraph_nodes, [node.output[0]], input_name_to_nodes, output_name_to_node
|
||||
):
|
||||
if not self.model.is_safe_to_fuse_nodes(subgraph_nodes, node.output, input_name_to_nodes, output_name_to_node):
|
||||
logger.debug(f"Skip fusing SkipLayerNormalization with Bias since it is not safe")
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ static void RunTest(
|
|||
const std::vector<float>& beta_data,
|
||||
const std::vector<float>& bias_data,
|
||||
const std::vector<float>& output_data,
|
||||
const std::vector<float>& skip_input_add_output_data,
|
||||
float epsilon,
|
||||
int batch_size,
|
||||
int sequence_length,
|
||||
|
|
@ -54,6 +55,17 @@ static void RunTest(
|
|||
}
|
||||
|
||||
test.AddOutput<float>("output", output_dims, output_data);
|
||||
|
||||
if (skip_input_add_output_data.size() != 0) {
|
||||
// The second and third outputs are reserved for something else
|
||||
test.AddOptionalOutputEdge<float>();
|
||||
test.AddOptionalOutputEdge<float>();
|
||||
|
||||
test.AddOutput<float>("input_skip_input_add_output",
|
||||
output_dims,
|
||||
skip_input_add_output_data);
|
||||
}
|
||||
|
||||
test.Run();
|
||||
} else if (HasCudaEnvironment(530 /*min_cuda_architecture*/) ||
|
||||
dml_ep != nullptr ||
|
||||
|
|
@ -74,9 +86,19 @@ static void RunTest(
|
|||
|
||||
test.AddOutput<MLFloat16>("output", output_dims, ToFloat16(output_data));
|
||||
|
||||
if (skip_input_add_output_data.size() != 0) {
|
||||
// The second and third outputs are reserved for something else
|
||||
test.AddOptionalOutputEdge<MLFloat16>();
|
||||
test.AddOptionalOutputEdge<MLFloat16>();
|
||||
|
||||
test.AddOutput<MLFloat16>("input_skip_input_add_output",
|
||||
output_dims,
|
||||
ToFloat16(skip_input_add_output_data));
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
if (dml_ep != nullptr) {
|
||||
execution_providers.push_back(DefaultDmlExecutionProvider());
|
||||
execution_providers.push_back(DefaultDmlExecutionProvider());
|
||||
} else if (rocm_ep != nullptr) {
|
||||
execution_providers.push_back(DefaultRocmExecutionProvider());
|
||||
} else {
|
||||
|
|
@ -110,6 +132,7 @@ TEST(SkipLayerNormTest, SkipLayerNormNullInput) {
|
|||
beta_data,
|
||||
std::vector<float>(),
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -145,6 +168,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1) {
|
|||
beta_data,
|
||||
std::vector<float>(),
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -180,6 +204,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1_Float16) {
|
|||
beta_data,
|
||||
std::vector<float>(),
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -272,6 +297,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1_Float16_vec) {
|
|||
beta_data,
|
||||
std::vector<float>(),
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -307,6 +333,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1_NoBeta) {
|
|||
beta_data,
|
||||
std::vector<float>(),
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -350,6 +377,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2) {
|
|||
beta_data,
|
||||
std::vector<float>(),
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -394,11 +422,63 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2_Bias) {
|
|||
beta_data,
|
||||
bias_data,
|
||||
output_data,
|
||||
{},
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
hidden_size);
|
||||
}
|
||||
|
||||
// Don't enable this test for DML/ROCM builds as these EP doesn't produce the new optional output yet
|
||||
#if !defined(USE_ROCM) && !defined(USE_DML)
|
||||
TEST(SkipLayerNormTest, SkipLayerNormBatch2_Bias_ProducingOptionalOutput) {
|
||||
int batch_size = 2;
|
||||
int sequence_length = 2;
|
||||
int hidden_size = 4;
|
||||
|
||||
std::vector<float> input_data = {
|
||||
0.7f, -0.4f, -0.2f, 1.2f,
|
||||
0.4f, 0.3f, 0.1f, -0.4f,
|
||||
0.7f, -0.4f, -0.2f, 1.2f,
|
||||
0.4f, 0.3f, 0.1f, -0.4f};
|
||||
|
||||
std::vector<float> skip_data = {
|
||||
0.1f, -0.2f, 0.3f, 1.0f,
|
||||
0.5f, 0.1f, 0.4f, 1.6f,
|
||||
1.8f, -0.3f, 0.0f, 1.f,
|
||||
-0.5f, 0.4f, 0.8f, -0.6f};
|
||||
|
||||
std::vector<float> gamma_data = {
|
||||
0.3f, 0.2f, 4.0f, 2.2f};
|
||||
|
||||
std::vector<float> beta_data = {
|
||||
0.2f, 0.1f, 0.4f, 1.6f};
|
||||
|
||||
std::vector<float> bias_data = {
|
||||
0.1f, -0.1f, 0.2f, -0.2f};
|
||||
|
||||
std::vector<float> output_data = {
|
||||
0.28433859348297119, -0.17090578377246857, -0.92897164821624756, 4.6924152374267578,
|
||||
0.46111652255058289, -0.21333980560302734, -0.29631003737449646, 3.5148544311523438,
|
||||
0.55470430850982666, -0.15080101788043976, -2.3229825496673584, 3.255286693572998,
|
||||
0.15631480515003204, 0.21066918969154358, 4.9432611465454102, -1.7957965135574341};
|
||||
|
||||
std::vector<float> input_skip_add_output_data(batch_size * sequence_length * hidden_size, 0.f);
|
||||
std::transform(input_data.cbegin(), input_data.cend(), skip_data.cbegin(), input_skip_add_output_data.begin(), [](const float& i, const float& s) { return i + s; });
|
||||
|
||||
RunTest(input_data,
|
||||
skip_data,
|
||||
gamma_data,
|
||||
beta_data,
|
||||
bias_data,
|
||||
output_data,
|
||||
input_skip_add_output_data,
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
hidden_size);
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -14,9 +14,11 @@ from model_loader import get_test_data_path
|
|||
from parity_utilities import find_transformers_source
|
||||
|
||||
if find_transformers_source():
|
||||
from fusion_options import FusionOptions
|
||||
from onnx_model import OnnxModel
|
||||
from optimizer import optimize_by_fusion, optimize_model
|
||||
else:
|
||||
from onnxruntime.transformers.fusion_options import FusionOptions
|
||||
from onnxruntime.transformers.onnx_model import OnnxModel
|
||||
from onnxruntime.transformers.optimizer import optimize_by_fusion, optimize_model
|
||||
|
||||
|
|
@ -120,32 +122,67 @@ class TestFusion(unittest.TestCase):
|
|||
hidden_size = 64
|
||||
num_heads = 4
|
||||
for add_order in [False, True]:
|
||||
model = create_gpt2_attention(
|
||||
hidden_size=hidden_size,
|
||||
num_heads=num_heads,
|
||||
switch_add_inputs=add_order,
|
||||
)
|
||||
dir = "."
|
||||
model_path = os.path.join(dir, "gpt2_attention.onnx")
|
||||
onnx.save(model, model_path)
|
||||
optimized_model = optimize_model(
|
||||
model_path,
|
||||
model_type="gpt2",
|
||||
num_heads=num_heads,
|
||||
hidden_size=hidden_size,
|
||||
)
|
||||
optimized_model.topological_sort()
|
||||
os.remove(model_path)
|
||||
for enable_skip_layer_norm_fusion in [False, True]:
|
||||
model = create_gpt2_attention(
|
||||
hidden_size=hidden_size,
|
||||
num_heads=num_heads,
|
||||
switch_add_inputs=add_order,
|
||||
)
|
||||
dir = "."
|
||||
model_path = os.path.join(dir, "gpt2_attention.onnx")
|
||||
onnx.save(model, model_path)
|
||||
|
||||
model_name = "gpt2_attention_{}.onnx".format("add_opt" if add_order else "opt")
|
||||
self.verify_fusion(optimized_model, model_name)
|
||||
options = FusionOptions("gpt2")
|
||||
options.enable_skip_layer_norm = enable_skip_layer_norm_fusion
|
||||
|
||||
optimized_model = optimize_model(
|
||||
model_path,
|
||||
model_type="gpt2",
|
||||
num_heads=num_heads,
|
||||
hidden_size=hidden_size,
|
||||
optimization_options=options,
|
||||
)
|
||||
|
||||
optimized_model.topological_sort()
|
||||
os.remove(model_path)
|
||||
|
||||
model_suffix = ""
|
||||
if add_order and enable_skip_layer_norm_fusion:
|
||||
model_suffix = "add_opt_skiplayernorm"
|
||||
elif add_order and not enable_skip_layer_norm_fusion:
|
||||
model_suffix = "add_opt_no_skiplayernorm"
|
||||
elif not add_order and enable_skip_layer_norm_fusion:
|
||||
model_suffix = "opt_skiplayernorm"
|
||||
else:
|
||||
model_suffix = "opt_no_skiplayernorm"
|
||||
|
||||
model_name = "gpt2_attention_{}.onnx".format(model_suffix)
|
||||
self.verify_fusion(optimized_model, model_name)
|
||||
|
||||
def test_megatron_gpt2_attention_fusion(self):
|
||||
path = get_test_data_path("models", "gpt2_megatron.onnx")
|
||||
model = onnx.load(path)
|
||||
optimized_model = optimize_by_fusion(model, model_type="gpt2")
|
||||
for enable_skip_layer_norm_fusion in [False, True]:
|
||||
path = get_test_data_path("models", "gpt2_megatron.onnx")
|
||||
model = onnx.load(path)
|
||||
|
||||
self.verify_fusion(optimized_model, "gpt2_megatron_opt.onnx")
|
||||
options = FusionOptions("gpt2")
|
||||
options.enable_skip_layer_norm = enable_skip_layer_norm_fusion
|
||||
|
||||
optimized_model = optimize_by_fusion(
|
||||
model,
|
||||
model_type="gpt2",
|
||||
num_heads=0,
|
||||
hidden_size=0,
|
||||
optimization_options=options,
|
||||
)
|
||||
|
||||
model_suffix = ""
|
||||
if enable_skip_layer_norm_fusion:
|
||||
model_suffix = "opt_skiplayernorm"
|
||||
else:
|
||||
model_suffix = "opt_no_skiplayernorm"
|
||||
|
||||
model_name = "gpt2_megatron_{}.onnx".format(model_suffix)
|
||||
self.verify_fusion(optimized_model, model_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -132,19 +132,31 @@ class TestModelOptimization(unittest.TestCase):
|
|||
self.assertTrue(False)
|
||||
|
||||
def test_gpt2_past(self):
|
||||
input = _get_test_model_path("gpt2_past")
|
||||
model = optimize_model(input, "gpt2", num_heads=2, hidden_size=4)
|
||||
for enable_skip_layer_norm_fusion in [False, True]:
|
||||
input_path = _get_test_model_path("gpt2_past")
|
||||
|
||||
expected_node_count = {
|
||||
"EmbedLayerNormalization": 0,
|
||||
"Attention": 12,
|
||||
"Gelu": 0,
|
||||
"FastGelu": 12,
|
||||
"BiasGelu": 0,
|
||||
"LayerNormalization": 25,
|
||||
"SkipLayerNormalization": 0,
|
||||
}
|
||||
self.verify_node_count(model, expected_node_count, "test_gpt2_past")
|
||||
options = FusionOptions("gpt2")
|
||||
options.enable_skip_layer_norm = enable_skip_layer_norm_fusion
|
||||
|
||||
model = optimize_model(
|
||||
input_path,
|
||||
"gpt2",
|
||||
num_heads=2,
|
||||
hidden_size=4,
|
||||
optimization_options=options,
|
||||
)
|
||||
|
||||
expected_node_count = {
|
||||
"EmbedLayerNormalization": 0,
|
||||
"Attention": 12,
|
||||
"Gelu": 0,
|
||||
"FastGelu": 12,
|
||||
"BiasGelu": 0,
|
||||
# First LayerNorm is never fused to SkipLayerNorm as it doesn't meet the requirements
|
||||
"LayerNormalization": 25 if not enable_skip_layer_norm_fusion else 1,
|
||||
"SkipLayerNormalization": 0 if not enable_skip_layer_norm_fusion else 24,
|
||||
}
|
||||
self.verify_node_count(model, expected_node_count, "test_gpt2_past")
|
||||
|
||||
def test_gpt2_past_fp16(self):
|
||||
input_model_path = _get_test_model_path("gpt2_past")
|
||||
|
|
@ -156,18 +168,30 @@ class TestModelOptimization(unittest.TestCase):
|
|||
self.assertEqual(output.type.tensor_type.elem_type, TensorProto.FLOAT16)
|
||||
|
||||
def test_gpt2_past_mask(self):
|
||||
input = _get_test_model_path("gpt2_past_mask")
|
||||
model = optimize_model(input, "gpt2", num_heads=2, hidden_size=4)
|
||||
expected_node_count = {
|
||||
"EmbedLayerNormalization": 1,
|
||||
"Attention": 1,
|
||||
"Gelu": 0,
|
||||
"FastGelu": 1,
|
||||
"BiasGelu": 0,
|
||||
"LayerNormalization": 1,
|
||||
"SkipLayerNormalization": 0,
|
||||
}
|
||||
self.verify_node_count(model, expected_node_count, "test_gpt2_past_mask")
|
||||
for enable_skip_layer_norm_fusion in [False, True]:
|
||||
input_path = _get_test_model_path("gpt2_past_mask")
|
||||
|
||||
options = FusionOptions("gpt2")
|
||||
options.enable_skip_layer_norm = enable_skip_layer_norm_fusion
|
||||
|
||||
model = optimize_model(
|
||||
input_path,
|
||||
"gpt2",
|
||||
num_heads=2,
|
||||
hidden_size=4,
|
||||
optimization_options=options,
|
||||
)
|
||||
|
||||
expected_node_count = {
|
||||
"EmbedLayerNormalization": 1,
|
||||
"Attention": 1,
|
||||
"Gelu": 0,
|
||||
"FastGelu": 1,
|
||||
"BiasGelu": 0,
|
||||
"LayerNormalization": 1 if not enable_skip_layer_norm_fusion else 0,
|
||||
"SkipLayerNormalization": 0 if not enable_skip_layer_norm_fusion else 1,
|
||||
}
|
||||
self.verify_node_count(model, expected_node_count, "test_gpt2_past_mask")
|
||||
|
||||
def test_multiple_embed(self):
|
||||
input_model_path = _get_test_model_path("multiple_embed")
|
||||
|
|
|
|||
Loading…
Reference in a new issue