saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

T5 skip_layer_norm cuda op (#14093)

Browse source 
### Description

T5 uses a layer_norm which only scales and doesn't shift, which is also
known as Root Mean Square Layer Normalization.
ORT already have the simplified_layer_norm which is the RMS layer_norm.
This PR extends this T5 layer_norm with support of skip/bias and the
residual output.
This new op is named SkipSimplifiedLayerNorm and has similar interface
as SkipLayerNorm but removes the beta as input


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang 2023-01-04 13:31:53 -08:00 committed by GitHub
parent b6ea60436d
commit ae148ebc05
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 255 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
docs/ContribOperators.md
View file

@ -76,6 +76,7 @@ Do not modify directly.*
* <a href="#com.microsoft.SampleOp">com.microsoft.SampleOp</a>
* <a href="#com.microsoft.Sampling">com.microsoft.Sampling</a>
* <a href="#com.microsoft.SkipLayerNormalization">com.microsoft.SkipLayerNormalization</a>
* <a href="#com.microsoft.SkipSimplifiedLayerNormalization">com.microsoft.SkipSimplifiedLayerNormalization</a>
* <a href="#com.microsoft.Snpe">com.microsoft.Snpe</a>
* <a href="#com.microsoft.SparseToDenseMatMul">com.microsoft.SparseToDenseMatMul</a>
* <a href="#com.microsoft.Tokenizer">com.microsoft.Tokenizer</a>
@ -3947,6 +3948,57 @@ This version of the operator has been available since version 1 of the 'com.micr
</dl>
### <a name="com.microsoft.SkipSimplifiedLayerNormalization"></a><a name="com.microsoft.skipsimplifiedlayernormalization">**com.microsoft.SkipSimplifiedLayerNormalization**</a>
Skip and Root Mean Square Layer Normalization
#### Version
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
#### Attributes
<dl>
<dt><tt>epsilon</tt> : float</dt>
<dd>The epsilon value to use to avoid division by zero.</dd>
</dl>
#### Inputs (3 - 4)
<dl>
<dt><tt>input</tt> : T</dt>
<dd>3D input tensor with shape (batch_size, sequence_length, hidden_size)</dd>
<dt><tt>skip</tt> : T</dt>
<dd>3D skip tensor with shape (batch_size, sequence_length, hidden_size)</dd>
<dt><tt>gamma</tt> : T</dt>
<dd>1D input tensor with shape (hidden_size)</dd>
<dt><tt>bias</tt> (optional) : T</dt>
<dd>1D bias tensor with shape (hidden_size</dd>
</dl>
#### Outputs (1 - 4)
<dl>
<dt><tt>output</tt> : T</dt>
<dd>3D output tensor with shape (batch_size, sequence_length, hidden_size)</dd>
<dt><tt>mean</tt> (optional) : U</dt>
<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
<dl>
<dt><tt>T</tt> : tensor(float), tensor(float16)</dt>
<dd>Constrain input and output types to float or half tensors.</dd>
<dt><tt>U</tt> : tensor(float)</dt>
<dd>Constrain mean and inv_std_var to float tensors.</dd>
</dl>
### <a name="com.microsoft.Snpe"></a><a name="com.microsoft.snpe">**com.microsoft.Snpe**</a>
Onnx node for SNPE.

1
docs/OperatorKernels.md
View file

@ -800,6 +800,7 @@ Do not modify directly.*
|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)|
|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)|
| |

57
onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh
View file

@ -109,6 +109,30 @@ __device__ inline void LayerNorm(
}
}
template <typename T, int TPB>
__device__ inline void SimplifiedLayerNorm(
const T& thread_data, const int ld, const int offset, const T* gamma, const T epsilon, T* output) {
// Assuming thread_data is already divided by ld
using BlockReduce = cub::BlockReduce<T, TPB>;
__shared__ typename BlockReduce::TempStorage temp_storage;
__shared__ T rsigma; // 1 / std.dev.
const T sum = BlockReduce(temp_storage).Sum(thread_data);
if (threadIdx.x == 0) {
rsigma = Rsqrt(sum + epsilon);
}
__syncthreads();
for (int i = threadIdx.x; i < ld; i += TPB) {
const int idx = offset + i;
const T val = output[idx];
const T g(gamma[i]);
output[idx] = g * val * rsigma;
}
}
template <typename T, int TPB, int ILP>
__device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair<T, T>& thread_data,
const int ld, const int idx, const T* beta, const T* gamma,
@ -152,6 +176,39 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair<
}
}
template <typename T, int TPB, int ILP>
__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const T& thread_data, const int ld, const int idx,
const T* gamma, const T epsilon, T* output) {
// Assuming thread_data is already divided by ld
// Small settings: the block covers the leading dimension TPB >= ld. The input
// value is available in a register
using VecT = aligned_vector<T, ILP>;
using BlockReduce = cub::BlockReduce<T, TPB>;
__shared__ typename BlockReduce::TempStorage temp_storage;
__shared__ T rsigma; // 1 / std.dev.
T gamma_v[ILP], output_v[ILP];
VecT* gamma_val = reinterpret_cast<VecT*>(&gamma_v);
*gamma_val = *reinterpret_cast<const VecT*>(&gamma[threadIdx.x * ILP]);
VecT* output_val = reinterpret_cast<VecT*>(&output_v);
const T sum = BlockReduce(temp_storage).Sum(thread_data);
if (threadIdx.x == 0) {
rsigma = Rsqrt(sum + epsilon);
}
__syncthreads();
if (ILP * threadIdx.x < ld) {
#pragma unroll
for (int i = 0; i < ILP; i++) {
output_v[i] = gamma_v[i] * input_v[i] * rsigma;
}
*(reinterpret_cast<VecT*>(&output[idx])) = *output_val;
}
}
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

46
onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc
View file

@ -18,26 +18,36 @@ namespace cuda {
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
SkipLayerNorm<T>);
SkipLayerNorm<T, false>); \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
SkipSimplifiedLayerNormalization, \
kMSDomain, \
1, \
T, \
kCudaExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
SkipLayerNorm<T, true>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(MLFloat16)
using namespace ONNX_NAMESPACE;
template <typename T>
SkipLayerNorm<T>::SkipLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {
template <typename T, bool Simplified>
SkipLayerNorm<T, Simplified>::SkipLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {
ORT_ENFORCE(op_kernel_info.GetAttr<float>("epsilon", &epsilon_).IsOK());
ORT_ENFORCE(epsilon_ >= 0);
}
template <typename T>
Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
template <typename T, bool Simplified>
Status SkipLayerNorm<T, Simplified>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor* input = ctx->Input<Tensor>(0);
const Tensor* skip = ctx->Input<Tensor>(1);
const Tensor* gamma = ctx->Input<Tensor>(2);
const Tensor* beta = ctx->Input<Tensor>(3);
const Tensor* bias = ctx->Input<Tensor>(4);
const Tensor* beta = Simplified ? nullptr : ctx->Input<Tensor>(3);
const Tensor* bias = Simplified ? ctx->Input<Tensor>(3): ctx->Input<Tensor>(4);
Tensor* output = ctx->Output(0, input->Shape());
@ -70,15 +80,17 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
"Last dimension of gamma and input does not match");
}
if (nullptr != beta) {
const auto& beta_dims = beta->Shape().GetDims();
if (beta_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have 1 dimension, got ", beta_dims.size());
}
if (beta_dims[0] != input_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Last dimension of beta and input does not match");
if (!Simplified) {
if (nullptr != beta) {
const auto& beta_dims = beta->Shape().GetDims();
if (beta_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have 1 dimension, got ", beta_dims.size());
}
if (beta_dims[0] != input_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Last dimension of beta and input does not match");
}
}
}
@ -100,7 +112,7 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
size_t element_size = sizeof(T);
typedef typename ToCudaType<T>::MappedType CudaT;
return LaunchSkipLayerNormKernel<CudaT>(
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,

2
onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.h
View file

@ -11,7 +11,7 @@ namespace cuda {
using namespace onnxruntime::cuda;
template <typename T>
template <typename T, bool Simplified>
class SkipLayerNorm final : public CudaKernel {
public:
SkipLayerNorm(const OpKernelInfo& op_kernel_info);

61
onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu
View file

@ -48,7 +48,7 @@ half maybe2half(float x) {
return __float2half_rn(x);
}
template <typename T, unsigned TPB>
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,
@ -72,12 +72,15 @@ __global__ void SkipLayerNormKernel(
thread_data = pair_sum(thread_data, cub::KeyValuePair<T, T>(rldval, rldval * val));
output[idx] = val;
}
if (Simplified) {
SimplifiedLayerNorm<T, TPB>(thread_data.value, ld, offset, gamma, epsilon, output);
return;
}
LayerNorm<T, TPB>(thread_data, ld, offset, beta, gamma, epsilon, output);
}
// Vectorized kernel
template <typename T, unsigned TPB, int ILP>
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,
@ -123,10 +126,15 @@ __global__ void SkipLayerNormKernelSmall(
thread_data = cub::KeyValuePair<T, T>(rldval_sum, rldvalsq_sum);
}
if (Simplified) {
SimplifiedLayerNormSmall<T, TPB, ILP>(input_v, thread_data.value, ld, idx, gamma, epsilon, output);
return;
}
LayerNormSmall<T, TPB, ILP>(input_v, thread_data, ld, idx, beta, gamma, epsilon, output);
}
template <typename T>
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,
const T* beta, const T* bias, float epsilon, const int ld, const int element_count,
@ -140,43 +148,43 @@ Status LaunchSkipLayerNormKernel(
const int grid_size = element_count / ld;
if (ld <= 32) {
constexpr int block_size = 32;
SkipLayerNormKernelSmall<T, block_size, 1>
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);
} else if (ld <= 64) {
constexpr int block_size = 64 / 2;
SkipLayerNormKernelSmall<T, block_size, 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);
} else if (ld <= 128) {
constexpr int block_size = 128 / 4;
SkipLayerNormKernelSmall<T, block_size, 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);
} else if (ld <= 384) {
constexpr int block_size = 384 / 4;
SkipLayerNormKernelSmall<T, block_size, 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);
} else if (ld <= 768) {
constexpr int block_size = 768 / 4;
SkipLayerNormKernelSmall<T, block_size, 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);
} else if (ld <= 1024) {
constexpr int block_size = 1024 / 4;
SkipLayerNormKernelSmall<T, block_size, 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);
} else {
constexpr int block_size = 256;
SkipLayerNormKernel<T, block_size>
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);
}
@ -184,31 +192,31 @@ Status LaunchSkipLayerNormKernel(
const int grid_size = element_count / ld;
if (ld <= 32) {
constexpr int block_size = 32;
SkipLayerNormKernelSmall<T, block_size, 1>
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);
} else if (ld <= 64) {
constexpr int block_size = 64;
SkipLayerNormKernelSmall<T, block_size, 1>
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);
} else if (ld <= 128) {
constexpr int block_size = 128;
SkipLayerNormKernelSmall<T, block_size, 1>
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);
} else if (ld == 384) {
constexpr int block_size = 384;
SkipLayerNormKernelSmall<T, block_size, 1>
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);
} else {
constexpr int block_size = 256;
SkipLayerNormKernel<T, block_size>
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);
}
@ -216,17 +224,18 @@ Status LaunchSkipLayerNormKernel(
return CUDA_CALL(cudaGetLastError());
}
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);
#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, \
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);
SKIPLAYERNORM_IMPL(float, true);
SKIPLAYERNORM_IMPL(float, false);
SKIPLAYERNORM_IMPL(half, true);
SKIPLAYERNORM_IMPL(half, false);
} // namespace cuda
} // namespace contrib

2
onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.h
View file

@ -8,7 +8,7 @@ namespace onnxruntime {
namespace contrib {
namespace cuda {
template <typename T>
template <typename T, bool Simplified>
Status LaunchSkipLayerNormKernel(
cudaStream_t stream,
T* output, // normalized output tensor

4
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
View file

@ -81,6 +81,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ScaledTanh);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, SkipLayerNormalization);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, SkipLayerNormalization);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, SkipSimplifiedLayerNormalization);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, SkipSimplifiedLayerNormalization);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ThresholdedRelu);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ThresholdedRelu);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ThresholdedRelu);
@ -196,6 +198,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ScaledTanh)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, SkipLayerNormalization)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, SkipLayerNormalization)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, SkipSimplifiedLayerNormalization)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, SkipSimplifiedLayerNormalization)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ThresholdedRelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ThresholdedRelu)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ThresholdedRelu)>,

16
onnxruntime/core/graph/contrib_ops/bert_defs.cc
View file

@ -394,6 +394,22 @@ ONNX_MS_OPERATOR_SET_SCHEMA(SkipLayerNormalization, 1,
.TypeConstraint("U", {"tensor(float)"}, "Constrain mean and inv_std_var to float tensors.")
.TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput));
ONNX_MS_OPERATOR_SET_SCHEMA(SkipSimplifiedLayerNormalization, 1,
OpSchema()
.SetDoc("Skip and Root Mean Square Layer Normalization")
.Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, kDefaultSkipLayerNormEpsilon)
.Input(0, "input", "3D input tensor with shape (batch_size, sequence_length, hidden_size)", "T")
.Input(1, "skip", "3D skip tensor with shape (batch_size, sequence_length, hidden_size)", "T")
.Input(2, "gamma", "1D input tensor with shape (hidden_size)", "T")
.Input(3, "bias", "1D bias tensor with shape (hidden_size", "T", OpSchema::Optional)
.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));
constexpr const char* NGramRepeatBlock_ver1_doc = R"DOC(
Enforce no repetition of n-grams. Scores are set to `-inf` for tokens that form a repeated n-gram if added to the back of the input_ids.
)DOC";

2
onnxruntime/core/graph/contrib_ops/ms_opset.h
View file

@ -83,6 +83,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Rfft);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SampleOp);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Sampling);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipLayerNormalization);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipSimplifiedLayerNormalization);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparseToDenseMatMul);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Tokenizer);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, TorchEmbedding);
@ -168,6 +169,7 @@ class OpSet_Microsoft_ver1 {
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SampleOp)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Sampling)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipLayerNormalization)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SkipSimplifiedLayerNormalization)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SparseToDenseMatMul)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Tokenizer)>());
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, TorchEmbedding)>());

3
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -198,6 +198,7 @@ class SymbolicShapeInference:
"LongformerAttention": self._infer_LongformerAttention,
"PythonOp": self._infer_PythonOp,
"SkipLayerNormalization": self._infer_SkipLayerNormalization,
"SkipSimplifiedLayerNormalization": self._infer_SkipLayerNormalization,
}
self.aten_op_dispatcher_ = {
"embedding": self._infer_Gather,
@ -2029,6 +2030,8 @@ class SymbolicShapeInference:
def _infer_SkipLayerNormalization(self, node):
self._propagate_shape_and_type(node)
if len(node.output) > 3:
self._propagate_shape_and_type(node, 0, 3)
# If the SkipLayerNormalization node contains the optional
# output for inference, infer the shape and type for it too

65
onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc
View file

@ -23,7 +23,8 @@ static void RunTest(
int sequence_length,
int hidden_size,
bool use_float16 = false,
bool no_beta = false) {
bool no_beta = false,
bool simplified = false) {
// Input and output shapes
// Input 0 - input: (batch_size, sequence_length, hidden_size)
// Input 1 - skip : (batch_size, sequence_length, hidden_size)
@ -37,17 +38,21 @@ static void RunTest(
std::vector<int64_t> bias_dims = gamma_dims;
std::vector<int64_t> output_dims = input_dims;
std::string op_type = simplified ? "SkipSimplifiedLayerNormalization" : "SkipLayerNormalization";
auto rocm_ep = DefaultRocmExecutionProvider();
auto dml_ep = DefaultDmlExecutionProvider();
if (!use_float16) {
OpTester test("SkipLayerNormalization", 1, onnxruntime::kMSDomain);
OpTester test(op_type.c_str(), 1, onnxruntime::kMSDomain);
test.AddInput<float>("input", input_dims, input_data);
test.AddInput<float>("skip", skip_dims, skip_data);
test.AddInput<float>("gamma", gamma_dims, gamma_data);
if (!no_beta) {
test.AddInput<float>("beta", beta_dims, beta_data);
} else {
test.AddOptionalInputEdge<float>();
if (!simplified) {
if (!no_beta) {
test.AddInput<float>("beta", beta_dims, beta_data);
} else {
test.AddOptionalInputEdge<float>();
}
}
test.AddAttribute("epsilon", epsilon);
if (!bias_data.empty()) {
@ -70,14 +75,16 @@ static void RunTest(
} else if (HasCudaEnvironment(530 /*min_cuda_architecture*/) ||
dml_ep != nullptr ||
rocm_ep != nullptr) {
OpTester test("SkipLayerNormalization", 1, onnxruntime::kMSDomain);
OpTester test(op_type.c_str(), 1, onnxruntime::kMSDomain);
test.AddInput<MLFloat16>("input", input_dims, ToFloat16(input_data));
test.AddInput<MLFloat16>("skip", skip_dims, ToFloat16(skip_data));
test.AddInput<MLFloat16>("gamma", gamma_dims, ToFloat16(gamma_data));
if (!no_beta) {
test.AddInput<MLFloat16>("beta", beta_dims, ToFloat16(beta_data));
} else {
test.AddOptionalInputEdge<float>();
if (!simplified) {
if (!no_beta) {
test.AddInput<MLFloat16>("beta", beta_dims, ToFloat16(beta_data));
} else {
test.AddOptionalInputEdge<float>();
}
}
test.AddAttribute("epsilon", epsilon);
if (!bias_data.empty()) {
@ -478,6 +485,42 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch2_Bias_ProducingOptionalOutput) {
sequence_length,
hidden_size);
}
TEST(SkipLayerNormTest, SkipSimplifiedLayerNormBatch1_Float16) {
int batch_size = 1;
int sequence_length = 2;
int hidden_size = 4;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> skip_data = {
0.0f, 0.0f, 0.0f, 0.0f,
0.0f, 0.0f, 0.0f, 0.0f};
std::vector<float> gamma_data = {
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};
RunTest(input_data,
skip_data,
gamma_data,
std::vector<float>(),
std::vector<float>(),
output_data,
{},
epsilon_,
batch_size,
sequence_length,
hidden_size,
true,
true,
true);
}
#endif
} // namespace test