mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-27 20:02:15 +00:00
Misc transformer fixes - 2 (#14156)
### Description 1. The graph pattern search introduced in https://github.com/microsoft/onnxruntime/pull/13914/ needs to be enhanced so that SkipLayerNormalization is supported 2. Fix fp32 parity for GPT-2 while using `SkipLayerNormalization` fusion. The optional output of SLN needs to also include the bias (if present) and the added output should be a sum of `input + skip + (bias)` ### Motivation and Context Fix some breaking tests
This commit is contained in:
Hariharan Seshadri
GitHub
parent
3702806653
commit
d0c5ffd5f7
10 changed files with 201 additions and 98 deletions
|
|
@ -3934,8 +3934,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>
|
||||
<dt><tt>input_skip_bias_sum</tt> (optional) : T</dt>
|
||||
<dd>Sum of the input and skip inputs (and bias if it exists) with shape (batch_size, sequence_length, hidden_size).</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
|
@ -3985,8 +3985,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>
|
||||
<dt><tt>input_skip_bias_sum</tt> (optional) : T</dt>
|
||||
<dd>Sum of the input and skip inputs (and bias if it exists) with shape (batch_size, sequence_length, hidden_size).</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
|
|
|||
|
|
@ -439,7 +439,7 @@ Do not modify directly.*
|
|||
|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)|
|
||||
|Sampling|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *in* attention_mask:**I**<br> *in* presence_mask:**I**<br> *out* sequences:**I**<br> *out* filtered_logits:**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**<br> *out* input_skip_sum:**T**|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_bias_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)|
|
||||
|
|
@ -799,8 +799,8 @@ Do not modify directly.*
|
|||
|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)|
|
||||
|Sampling|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *in* attention_mask:**I**<br> *in* presence_mask:**I**<br> *out* sequences:**I**<br> *out* filtered_logits:**T**|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)|
|
||||
|SkipSimplifiedLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**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)|
|
||||
|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_bias_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|SkipSimplifiedLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**<br> *out* input_skip_bias_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)|
|
||||
| |
|
||||
|
|
@ -1138,7 +1138,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**<br> *out* input_skip_sum:**T**|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_bias_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
| |
|
||||
| |
|
||||
|**Operator Domain:** *com.microsoft.dml*||||
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
|
|||
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());
|
||||
Tensor* skip_input_bias_add_output = p_ctx->Output(3, input->Shape());
|
||||
|
||||
const auto& input_dims = input->Shape().GetDims();
|
||||
if (input_dims.size() != 3) {
|
||||
|
|
@ -103,7 +103,7 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
|
|||
|
||||
// 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;
|
||||
T* skip_input_bias_add_output_data = skip_input_bias_add_output != nullptr ? skip_input_bias_add_output->MutableData<T>() : nullptr;
|
||||
|
||||
concurrency::ThreadPool::TryBatchParallelFor(
|
||||
p_ctx->GetOperatorThreadPool(), static_cast<int32_t>(task_count),
|
||||
|
|
@ -113,19 +113,22 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
|
|||
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* p_skip_input_bias_add_output_data = skip_input_bias_add_output_data != nullptr ? skip_input_bias_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];
|
||||
}
|
||||
|
||||
if (nullptr != p_skip_input_bias_add_output_data) {
|
||||
p_skip_input_bias_add_output_data[h] = value;
|
||||
}
|
||||
|
||||
p_output[h] = value;
|
||||
mean += value;
|
||||
mean_square += value * value;
|
||||
|
|
|
|||
|
|
@ -47,13 +47,13 @@ Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const
|
|||
const Tensor* gamma = ctx->Input<Tensor>(2);
|
||||
|
||||
const Tensor* beta = Simplified ? nullptr : ctx->Input<Tensor>(3);
|
||||
const Tensor* bias = Simplified ? ctx->Input<Tensor>(3): ctx->Input<Tensor>(4);
|
||||
const Tensor* bias = Simplified ? ctx->Input<Tensor>(3) : ctx->Input<Tensor>(4);
|
||||
|
||||
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());
|
||||
Tensor* skip_input_bias_add_output = ctx->Output(3, input->Shape());
|
||||
|
||||
if (input->Shape() != skip->Shape()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
|
||||
|
|
@ -115,7 +115,7 @@ Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const
|
|||
return LaunchSkipLayerNormKernel<CudaT, Simplified>(
|
||||
Stream(ctx),
|
||||
reinterpret_cast<CudaT*>(output->MutableData<T>()),
|
||||
skip_input_add_output != nullptr ? reinterpret_cast<CudaT*>(skip_input_add_output->MutableData<T>()) : nullptr,
|
||||
skip_input_bias_add_output != nullptr ? reinterpret_cast<CudaT*>(skip_input_bias_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>()),
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ template <typename T, unsigned TPB, bool Simplified>
|
|||
__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, T* skip_input_add_output) {
|
||||
const T epsilon, T* output, T* skip_input_bias_add_output) {
|
||||
const T reverse_ld = T(1.f / ld);
|
||||
const int offset = blockIdx.x * ld;
|
||||
|
||||
|
|
@ -63,13 +63,14 @@ __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));
|
||||
|
||||
if (skip_input_bias_add_output != nullptr) {
|
||||
skip_input_bias_add_output[idx] = val;
|
||||
}
|
||||
|
||||
output[idx] = val;
|
||||
}
|
||||
if (Simplified) {
|
||||
|
|
@ -83,14 +84,14 @@ __global__ void SkipLayerNormKernel(
|
|||
template <typename T, unsigned TPB, int ILP, bool Simplified>
|
||||
__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, T* skip_input_add_output,
|
||||
bool hasBias, bool hasSkipInputAdditionOutput) {
|
||||
const T* bias, const T epsilon, T* output, T* skip_input_bias_add_output,
|
||||
bool hasBias, bool hasSkipInputBiasAdditionOutput) {
|
||||
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], skip_input_add_output_v[ILP];
|
||||
T input_v[ILP], skip_v[ILP], bias_v[ILP], skip_input_bias_add_output_v[ILP];
|
||||
|
||||
VecT* input_val = reinterpret_cast<VecT*>(&input_v);
|
||||
*input_val = *reinterpret_cast<const VecT*>(&input[idx]);
|
||||
|
|
@ -110,18 +111,19 @@ __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];
|
||||
|
||||
if (hasSkipInputBiasAdditionOutput) {
|
||||
skip_input_bias_add_output_v[i] = input_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);
|
||||
if (hasSkipInputBiasAdditionOutput) {
|
||||
*(reinterpret_cast<VecT*>(&skip_input_bias_add_output[idx])) = *reinterpret_cast<VecT*>(&skip_input_bias_add_output_v);
|
||||
}
|
||||
|
||||
thread_data = cub::KeyValuePair<T, T>(rldval_sum, rldvalsq_sum);
|
||||
|
|
@ -136,13 +138,13 @@ __global__ void SkipLayerNormKernelSmall(
|
|||
|
||||
template <typename T, bool Simplified>
|
||||
Status LaunchSkipLayerNormKernel(
|
||||
cudaStream_t stream, T* output, T* skip_input_add_output, const T* input, const T* skip, const T* gamma,
|
||||
cudaStream_t stream, T* output, T* skip_input_bias_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;
|
||||
bool hasSkipInputBiasAdditionOutput = (skip_input_bias_add_output == nullptr) ? false : true;
|
||||
|
||||
if (0 == (ld % 4)) {
|
||||
const int grid_size = element_count / ld;
|
||||
|
|
@ -151,42 +153,42 @@ Status LaunchSkipLayerNormKernel(
|
|||
SkipLayerNormKernelSmall<T, block_size, 1, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 64) {
|
||||
constexpr int block_size = 64 / 2;
|
||||
SkipLayerNormKernelSmall<T, block_size, 2, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 128) {
|
||||
constexpr int block_size = 128 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 384) {
|
||||
constexpr int block_size = 384 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 768) {
|
||||
constexpr int block_size = 768 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 1024) {
|
||||
constexpr int block_size = 1024 / 4;
|
||||
SkipLayerNormKernelSmall<T, block_size, 4, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else {
|
||||
constexpr int block_size = 256;
|
||||
SkipLayerNormKernel<T, block_size, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, skip_input_add_output);
|
||||
maybe2half<T>(epsilon), output, skip_input_bias_add_output);
|
||||
}
|
||||
} else {
|
||||
const int grid_size = element_count / ld;
|
||||
|
|
@ -195,41 +197,41 @@ Status LaunchSkipLayerNormKernel(
|
|||
SkipLayerNormKernelSmall<T, block_size, 1, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 64) {
|
||||
constexpr int block_size = 64;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld <= 128) {
|
||||
constexpr int block_size = 128;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else if (ld == 384) {
|
||||
constexpr int block_size = 384;
|
||||
SkipLayerNormKernelSmall<T, block_size, 1, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output,
|
||||
skip_input_add_output, hasBias, hasSkipInputAdditionOutput);
|
||||
skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput);
|
||||
} else {
|
||||
constexpr int block_size = 256;
|
||||
SkipLayerNormKernel<T, block_size, Simplified>
|
||||
<<<grid_size, block_size, 0, stream>>>(ld, input, skip, beta, gamma, bias,
|
||||
maybe2half<T>(epsilon), output, skip_input_add_output);
|
||||
maybe2half<T>(epsilon), output, skip_input_bias_add_output);
|
||||
}
|
||||
}
|
||||
return CUDA_CALL(cudaGetLastError());
|
||||
}
|
||||
|
||||
#define SKIPLAYERNORM_IMPL(T, Simplified) \
|
||||
template Status LaunchSkipLayerNormKernel<T, Simplified>(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, \
|
||||
#define SKIPLAYERNORM_IMPL(T, Simplified) \
|
||||
template Status LaunchSkipLayerNormKernel<T, Simplified>(cudaStream_t stream, T * output, \
|
||||
T * skip_input_bias_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);
|
||||
|
||||
SKIPLAYERNORM_IMPL(float, true);
|
||||
|
|
|
|||
|
|
@ -11,16 +11,16 @@ namespace cuda {
|
|||
template <typename T, bool Simplified>
|
||||
Status LaunchSkipLayerNormKernel(
|
||||
cudaStream_t stream,
|
||||
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
|
||||
T* output, // normalized output tensor
|
||||
T* skip_input_bias_add_output, // sum of the input and skip (and bias if it exists) 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
|
||||
|
|
|
|||
|
|
@ -389,7 +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)
|
||||
.Output(3, "input_skip_bias_sum", "Sum of the input and skip inputs (and bias if it exists) 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));
|
||||
|
|
@ -405,7 +405,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(SkipSimplifiedLayerNormalization, 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)
|
||||
.Output(3, "input_skip_bias_sum", "Sum of the input and skip inputs (and bias if it exists) 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));
|
||||
|
|
|
|||
|
|
@ -455,7 +455,7 @@ def gpt2_to_onnx(args: argparse.Namespace):
|
|||
# TODO(tianleiwu): Use auto mixed precision for fp16 conversion: arguments.append('--auto_mixed_precision')
|
||||
# Need change cuda kernel to support a combination of fp32 logits and fp16 past state.
|
||||
# Currently logits and past state shall be same data type.
|
||||
arguments.extend(["--op_block_list", "Add", "LayerNormalization", "FastGelu"])
|
||||
arguments.extend(["--op_block_list", "Add", "LayerNormalization", "SkipLayerNormalization", "FastGelu"])
|
||||
|
||||
if args.verbose:
|
||||
logger.info(f"arguments for convert_to_onnx:{arguments}")
|
||||
|
|
@ -1114,6 +1114,8 @@ def generate_gpt2_init_decoder(
|
|||
|
||||
# Try to find the last residual Add
|
||||
# For fp16, there are Casts along the way
|
||||
|
||||
# Normalization Node is : LayerNormalization
|
||||
logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
|
||||
logits_matmul_node,
|
||||
[
|
||||
|
|
@ -1134,54 +1136,133 @@ def generate_gpt2_init_decoder(
|
|||
[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
)
|
||||
|
||||
# Normalization Node is : SkipLayerNormalization
|
||||
if logits_matmul_to_residual_add_path is None:
|
||||
logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
|
||||
logits_matmul_node,
|
||||
[
|
||||
"Cast",
|
||||
"SkipLayerNormalization",
|
||||
"Cast",
|
||||
"MatMul",
|
||||
"Cast",
|
||||
"FastGelu",
|
||||
"Cast",
|
||||
"MatMul",
|
||||
"Cast",
|
||||
"SkipLayerNormalization",
|
||||
],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls
|
||||
if logits_matmul_to_residual_add_path is None:
|
||||
|
||||
# Normalization Node is : LayerNormalization
|
||||
logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
|
||||
logits_matmul_node,
|
||||
["LayerNormalization", "Add", "Add", "MatMul", "FastGelu", "MatMul", "LayerNormalization", "Add"],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
)
|
||||
|
||||
# Normalization Node is : SkipLayerNormalization
|
||||
if logits_matmul_to_residual_add_path is None:
|
||||
logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
|
||||
logits_matmul_node,
|
||||
[
|
||||
"SkipLayerNormalization",
|
||||
"MatMul",
|
||||
"FastGelu",
|
||||
"MatMul",
|
||||
"SkipLayerNormalization",
|
||||
],
|
||||
[0, 1, 0, 0, 0],
|
||||
)
|
||||
|
||||
# TODO(hasesh): Are there more permutations to try before returning ?
|
||||
if logits_matmul_to_residual_add_path is None:
|
||||
return False
|
||||
|
||||
residual_add_node = logits_matmul_to_residual_add_path[-1]
|
||||
|
||||
residual_add_to_attention_parent_index = 0
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Add", "Cast", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0, 0]
|
||||
)
|
||||
# If the last node in the pattern is SkipLayerNormalization, we need to adjust our pattern searches accordingly
|
||||
is_skiplayernorm_path = True if residual_add_node.op_type == "SkipLayerNormalization" else False
|
||||
|
||||
# Try other parent index of the residual Add node
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 1
|
||||
# Regular LayerNormalization path
|
||||
if not is_skiplayernorm_path:
|
||||
residual_add_to_attention_parent_index = 0
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Add", "Cast", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0, 0]
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls
|
||||
if residual_add_to_attention_path is None:
|
||||
# Try other parent index of the residual Add node
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 1
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node,
|
||||
["Add", "Cast", "MatMul", "Attention"],
|
||||
[residual_add_to_attention_parent_index, 0, 0, 0],
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 0
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Add", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0]
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls and other parent index of the residual Add node
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 1
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Add", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0]
|
||||
)
|
||||
|
||||
# SkipLayerNormalization path
|
||||
else:
|
||||
residual_add_to_attention_parent_index = 0
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Add", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0]
|
||||
residual_add_node, ["Cast", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0]
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls and other parent index of the residual Add node
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 1
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Add", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0]
|
||||
)
|
||||
# Try other parent index of the residual Add node
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 1
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["Cast", "MatMul", "Attention"], [residual_add_to_attention_parent_index, 0, 0]
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 0
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["MatMul", "Attention"], [residual_add_to_attention_parent_index, 0]
|
||||
)
|
||||
|
||||
# Try without the Casts before and after the MatMuls and other parent index of the residual Add node
|
||||
if residual_add_to_attention_path is None:
|
||||
residual_add_to_attention_parent_index = 1
|
||||
residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
|
||||
residual_add_node, ["MatMul", "Attention"], [residual_add_to_attention_parent_index, 0]
|
||||
)
|
||||
|
||||
# TODO(hasesh): Are there more permutations to try before returning ?
|
||||
if residual_add_to_attention_path is None:
|
||||
return False
|
||||
|
||||
residual_add_to_add_parent_index = 0 if residual_add_to_attention_parent_index == 1 else 1
|
||||
add_before_residual_add = gpt2_init_decoder_model.match_parent(
|
||||
residual_add_node, "Add", residual_add_to_add_parent_index
|
||||
)
|
||||
|
||||
# Regular LayerNormalization path
|
||||
if not is_skiplayernorm_path:
|
||||
add_before_residual_add = gpt2_init_decoder_model.match_parent(
|
||||
residual_add_node, "Add", residual_add_to_add_parent_index
|
||||
)
|
||||
|
||||
# SkipLayerNormalization path
|
||||
else:
|
||||
add_before_residual_add = gpt2_init_decoder_model.match_parent(
|
||||
residual_add_node, "SkipLayerNormalization", residual_add_to_add_parent_index
|
||||
)
|
||||
|
||||
if add_before_residual_add is None:
|
||||
return False
|
||||
|
|
@ -1238,11 +1319,17 @@ def generate_gpt2_init_decoder(
|
|||
)
|
||||
|
||||
# Add Slice node to the graph such that it consumes the output of Add before the residual Add
|
||||
# If the 'Add' output is produced by a SkipLayerNormalization node, then adjust its output
|
||||
# index appropriately
|
||||
add_before_residual_add_output = (
|
||||
add_before_residual_add.output[0] if not is_skiplayernorm_path else add_before_residual_add.output[3]
|
||||
)
|
||||
|
||||
slice_1_output_name = "edge_modified_" + add_before_residual_add.output[0]
|
||||
slice_node_1 = onnx.helper.make_node(
|
||||
"Slice",
|
||||
inputs=[
|
||||
add_before_residual_add.output[0],
|
||||
add_before_residual_add_output,
|
||||
"SliceLastTokenStarts",
|
||||
"SliceLastTokenEnds",
|
||||
"SliceLastTokenAxes",
|
||||
|
|
@ -1258,9 +1345,7 @@ def generate_gpt2_init_decoder(
|
|||
|
||||
# Adjust the input(s) to the nodes consuming the outputs of the added Slice nodes
|
||||
gpt2_init_decoder_model.replace_node_input(matmul_after_attention, attention.output[0], slice_0_output_name)
|
||||
gpt2_init_decoder_model.replace_node_input(
|
||||
residual_add_node, add_before_residual_add.output[0], slice_1_output_name
|
||||
)
|
||||
gpt2_init_decoder_model.replace_node_input(residual_add_node, add_before_residual_add_output, slice_1_output_name)
|
||||
|
||||
# Topologically sort the updated graph
|
||||
gpt2_init_decoder_model.topological_sort()
|
||||
|
|
|
|||
|
|
@ -522,6 +522,11 @@ class Gpt2Helper:
|
|||
# Potentially, we can add other optimization like unpad for effective transformer
|
||||
optimization_options.attention_mask_format = AttentionMaskFormat.MaskIndexEnd
|
||||
|
||||
# TODO(hasesh): Investigate parity issue for GPT-2 fp16 when SkipLayerNormalization
|
||||
# is enabled
|
||||
if is_float16:
|
||||
optimization_options.enable_skip_layer_norm = False
|
||||
|
||||
m = optimize_model(
|
||||
onnx_model_path,
|
||||
model_type="gpt2",
|
||||
|
|
|
|||
|
|
@ -17,7 +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,
|
||||
const std::vector<float>& skip_input_bias_add_output_data,
|
||||
float epsilon,
|
||||
int batch_size,
|
||||
int sequence_length,
|
||||
|
|
@ -61,14 +61,14 @@ static void RunTest(
|
|||
|
||||
test.AddOutput<float>("output", output_dims, output_data);
|
||||
|
||||
if (skip_input_add_output_data.size() != 0) {
|
||||
if (skip_input_bias_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",
|
||||
test.AddOutput<float>("skip_input_bias_add_output",
|
||||
output_dims,
|
||||
skip_input_add_output_data);
|
||||
skip_input_bias_add_output_data);
|
||||
}
|
||||
|
||||
test.Run();
|
||||
|
|
@ -93,14 +93,14 @@ static void RunTest(
|
|||
|
||||
test.AddOutput<MLFloat16>("output", output_dims, ToFloat16(output_data));
|
||||
|
||||
if (skip_input_add_output_data.size() != 0) {
|
||||
if (skip_input_bias_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",
|
||||
test.AddOutput<MLFloat16>("skip_input_bias_add_output",
|
||||
output_dims,
|
||||
ToFloat16(skip_input_add_output_data));
|
||||
ToFloat16(skip_input_bias_add_output_data));
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
|
|
@ -470,8 +470,16 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2_Bias_ProducingOptionalOutput) {
|
|||
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; });
|
||||
std::vector<float> input_skip_bias_add_output_data(batch_size * sequence_length * hidden_size, 0.f);
|
||||
std::transform(input_data.cbegin(), input_data.cend(), skip_data.cbegin(), input_skip_bias_add_output_data.begin(), [](const float& i, const float& s) { return i + s; });
|
||||
// Add bias
|
||||
for (int i = 0; i < batch_size * sequence_length; ++i) {
|
||||
int offset = i * hidden_size;
|
||||
|
||||
for (int j = 0; j < hidden_size; ++j) {
|
||||
input_skip_bias_add_output_data[offset + j] += bias_data[j];
|
||||
}
|
||||
}
|
||||
|
||||
RunTest(input_data,
|
||||
skip_data,
|
||||
|
|
@ -479,7 +487,7 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2_Bias_ProducingOptionalOutput) {
|
|||
beta_data,
|
||||
bias_data,
|
||||
output_data,
|
||||
input_skip_add_output_data,
|
||||
input_skip_bias_add_output_data,
|
||||
epsilon_,
|
||||
batch_size,
|
||||
sequence_length,
|
||||
|
|
@ -503,8 +511,8 @@ TEST(SkipLayerNormTest, SkipSimplifiedLayerNormBatch1_Float16) {
|
|||
0.3f, 0.2f, 4.0f, 2.2f};
|
||||
|
||||
std::vector<float> output_data = {
|
||||
0.3491f, -0.1455f, 0.0000f, 3.2005f,
|
||||
0.3487f, 0.0930f, 2.7899f, -3.0689f};
|
||||
0.3491f, -0.1455f, 0.0000f, 3.2005f,
|
||||
0.3487f, 0.0930f, 2.7899f, -3.0689f};
|
||||
|
||||
RunTest(input_data,
|
||||
skip_data,
|
||||
|
|
|
|||
Loading…
Reference in a new issue