diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 6fd6084e11..e44cf96021 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3934,8 +3934,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Saved mean used during training to speed up gradient computation
inv_std_var (optional) : U
Saved inverse standard variance used during training to speed up gradient computation.
-
input_skip_sum (optional) : T
-
Sum of the input and skip inputs with shape (batch_size, sequence_length, hidden_size).
+
input_skip_bias_sum (optional) : T
+
Sum of the input and skip inputs (and bias if it exists) with shape (batch_size, sequence_length, hidden_size).
#### Type Constraints @@ -3985,8 +3985,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Saved mean used during training to speed up gradient computation
inv_std_var (optional) : U
Saved inverse standard variance used during training to speed up gradient computation.
-
input_skip_sum (optional) : T
-
Sum of the input and skip inputs with shape (batch_size, sequence_length, hidden_size).
+
input_skip_bias_sum (optional) : T
+
Sum of the input and skip inputs (and bias if it exists) with shape (batch_size, sequence_length, hidden_size).
#### Type Constraints diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 99f52bfb7b..93b8d8fc53 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -439,7 +439,7 @@ Do not modify directly.* |Range|*in* start:**T**
*in* limit:**T**
*in* delta:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64)| |SampleOp|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |Sampling|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*in* attention_mask:**I**
*in* presence_mask:**I**
*out* sequences:**I**
*out* filtered_logits:**T**|1+|**T** = tensor(float)| -|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_sum:**T**|1+|**T** = tensor(double), tensor(float)| +|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(double), tensor(float)| |SparseToDenseMatMul|*in* A:**T**
*in* B:**T1**
*out* Y:**T1**|1+|**T** = sparse_tensor(double), sparse_tensor(float), sparse_tensor(int32), sparse_tensor(int64), sparse_tensor(uint32), sparse_tensor(uint64)
**T1** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |Tokenizer|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(string)| |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -799,8 +799,8 @@ Do not modify directly.* |RestorePadding|*in* input:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |Rfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |Sampling|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*in* attention_mask:**I**
*in* presence_mask:**I**
*out* sequences:**I**
*out* filtered_logits:**T**|1+|**T** = tensor(float), tensor(float16)| -|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_sum:**T**|1+|**T** = tensor(float), tensor(float16)| -|SkipSimplifiedLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_sum:**T**|1+|**T** = tensor(float), tensor(float16)| +|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(float), tensor(float16)| +|SkipSimplifiedLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(float), tensor(float16)| |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*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**
*in* A_scale:**tensor(float)**
*in* A_zero_point:**T**
*in* B:**T**
*in* B_scale:**tensor(float)**
*in* B_zero_point:**T**
*in* C_scale:**tensor(float)**
*in* C_zero_point:**T**
*out* C:**T**|1+|**T** = tensor(int8), tensor(uint8)| |QLinearSigmoid|*in* X:**T**
*in* X_scale:**tensor(float)**
*in* X_zero_point:**T**
*in* Y_scale:**tensor(float)**
*in* Y_zero_point:**T**
*out* Y:**T**|1+|**T** = tensor(int8), tensor(uint8)| |QuantizeLinear|*in* x:**T1**
*in* y_scale:**T1**
*in* y_zero_point:**T2**
*out* y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(uint8)| -|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_sum:**T**|1+|**T** = tensor(float), tensor(float16)| +|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(float), tensor(float16)| | | | | |**Operator Domain:** *com.microsoft.dml*|||| diff --git a/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc b/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc index 991982b855..a394ebb4ce 100644 --- a/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/skip_layer_norm.cc @@ -41,7 +41,7 @@ Status SkipLayerNorm::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::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() : nullptr; + T* skip_input_bias_add_output_data = skip_input_bias_add_output != nullptr ? skip_input_bias_add_output->MutableData() : nullptr; concurrency::ThreadPool::TryBatchParallelFor( p_ctx->GetOperatorThreadPool(), static_cast(task_count), @@ -113,19 +113,22 @@ Status SkipLayerNorm::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; diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc index c914ee8422..53da91420e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc @@ -47,13 +47,13 @@ Status SkipLayerNorm::ComputeInternal(OpKernelContext* ctx) const const Tensor* gamma = ctx->Input(2); const Tensor* beta = Simplified ? nullptr : ctx->Input(3); - const Tensor* bias = Simplified ? ctx->Input(3): ctx->Input(4); + const Tensor* bias = Simplified ? ctx->Input(3) : ctx->Input(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::ComputeInternal(OpKernelContext* ctx) const return LaunchSkipLayerNormKernel( Stream(ctx), reinterpret_cast(output->MutableData()), - skip_input_add_output != nullptr ? reinterpret_cast(skip_input_add_output->MutableData()) : nullptr, + skip_input_bias_add_output != nullptr ? reinterpret_cast(skip_input_bias_add_output->MutableData()) : nullptr, reinterpret_cast(input->Data()), reinterpret_cast(skip->Data()), reinterpret_cast(gamma->Data()), diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index 9e164fbff0..a4ff8aac7f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -52,7 +52,7 @@ template __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(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 __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 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(&input_v); *input_val = *reinterpret_cast(&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(&skip_input_add_output[idx])) = *reinterpret_cast(&skip_input_add_output_v); + if (hasSkipInputBiasAdditionOutput) { + *(reinterpret_cast(&skip_input_bias_add_output[idx])) = *reinterpret_cast(&skip_input_bias_add_output_v); } thread_data = cub::KeyValuePair(rldval_sum, rldvalsq_sum); @@ -136,13 +138,13 @@ __global__ void SkipLayerNormKernelSmall( template 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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(epsilon), output, - skip_input_add_output, hasBias, hasSkipInputAdditionOutput); + skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput); } else { constexpr int block_size = 256; SkipLayerNormKernel <<>>(ld, input, skip, beta, gamma, bias, - maybe2half(epsilon), output, skip_input_add_output); + maybe2half(epsilon), output, skip_input_bias_add_output); } } else { const int grid_size = element_count / ld; @@ -195,41 +197,41 @@ Status LaunchSkipLayerNormKernel( SkipLayerNormKernelSmall <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(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 <<>>(ld, input, skip, beta, gamma, bias, maybe2half(epsilon), output, - skip_input_add_output, hasBias, hasSkipInputAdditionOutput); + skip_input_bias_add_output, hasBias, hasSkipInputBiasAdditionOutput); } else { constexpr int block_size = 256; SkipLayerNormKernel <<>>(ld, input, skip, beta, gamma, bias, - maybe2half(epsilon), output, skip_input_add_output); + maybe2half(epsilon), output, skip_input_bias_add_output); } } return CUDA_CALL(cudaGetLastError()); } -#define SKIPLAYERNORM_IMPL(T, Simplified) \ - template Status LaunchSkipLayerNormKernel(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(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); diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h index 6dab8b298f..da2894928c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h @@ -11,16 +11,16 @@ namespace cuda { template 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 diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 1395ef4ba6..af08a6f7e7 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -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)); diff --git a/onnxruntime/python/tools/transformers/convert_generation.py b/onnxruntime/python/tools/transformers/convert_generation.py index 8c3ecceeff..138f7c1f42 100644 --- a/onnxruntime/python/tools/transformers/convert_generation.py +++ b/onnxruntime/python/tools/transformers/convert_generation.py @@ -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() diff --git a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py index f8112013d4..f5de69e8f0 100644 --- a/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/models/gpt2/gpt2_helper.py @@ -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", diff --git a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc index ed5582e95e..4ed141c893 100644 --- a/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc +++ b/onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc @@ -17,7 +17,7 @@ static void RunTest( const std::vector& beta_data, const std::vector& bias_data, const std::vector& output_data, - const std::vector& skip_input_add_output_data, + const std::vector& skip_input_bias_add_output_data, float epsilon, int batch_size, int sequence_length, @@ -61,14 +61,14 @@ static void RunTest( test.AddOutput("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(); test.AddOptionalOutputEdge(); - test.AddOutput("input_skip_input_add_output", + test.AddOutput("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("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(); test.AddOptionalOutputEdge(); - test.AddOutput("input_skip_input_add_output", + test.AddOutput("skip_input_bias_add_output", output_dims, - ToFloat16(skip_input_add_output_data)); + ToFloat16(skip_input_bias_add_output_data)); } std::vector> 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 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 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 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,