diff --git a/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h index 182727ab58..00cc451073 100644 --- a/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h @@ -25,5 +25,10 @@ class LongformerAttentionBase { int window_; // Attention windows length (W). It is half (one-sided) of total window size. }; +namespace longformer { +// Environment variable to give a hint about choosing kernels for less memory or latency. +constexpr const char* kUseCompactMemory = "ORT_LONGFORMER_COMPACT_MEMORY"; +} // namespace longformer + } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc index 89156634a8..19e838971a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc @@ -5,6 +5,7 @@ #include "core/framework/tensorprotoutils.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" +#include "core/platform/env_var_utils.h" #include "longformer_global_impl.h" #include "longformer_attention_impl.h" @@ -49,7 +50,9 @@ class AutoDestoryCudaEvent { }; template -LongformerAttention::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {} +LongformerAttention::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) { + use_compact_memory_ = ParseEnvironmentVariableWithDefault(longformer::kUseCompactMemory, false); +} template Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { @@ -80,6 +83,7 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { constexpr size_t element_size = sizeof(T); + // TODO: only calculate once per model. // Build Global Index auto global_index_buffer = GetScratchBuffer(batch_size * sequence_length); auto batch_global_num_buffer = GetScratchBuffer(batch_size); @@ -148,10 +152,11 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { } } - // Cuda kernel implementation has a limitation of number of global tokens. - if (max_num_global > window_) { - ORT_THROW("LongformerAttention CUDA operator does not support number of global tokens > attention window."); - } + // Force to use fast kernel in two situations: + // (1) global tokens > windows size. In that case, compact memory kernel cannot be used. + // (2) sequence_length == 2 * attention_window. Use fast kernel to walk around parity issue of compact memory kernel. + // In other case, we will choose according to user's environment variable setting (default is fast kernel). + bool use_fast_kernel = (max_num_global > window_ || sequence_length == 2 * window_ || !use_compact_memory_); // Fully connection for global projection. // Note that Q only need handle global query tokens if we split GEMM to global Q/K/V separately. @@ -172,7 +177,7 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { &one, reinterpret_cast(global_gemm_buffer.get()), n, device_prop)); } - size_t workSpaceSize = GetLongformerAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, max_num_global, window_); + size_t workSpaceSize = GetLongformerAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, max_num_global, window_, use_fast_kernel); auto workspace_buffer = GetScratchBuffer(workSpaceSize); if (!LaunchLongformerAttentionKernel( device_prop, @@ -193,7 +198,8 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { head_size, window_, max_num_global, - element_size)) { + element_size, + use_fast_kernel)) { // Get last error to reset it to cudaSuccess. CUDA_CALL(cudaGetLastError()); return Status(common::ONNXRUNTIME, common::FAIL); diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h index 151e542c4e..fa9dfd8c9b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.h @@ -18,6 +18,9 @@ class LongformerAttention final : public CudaKernel, public LongformerAttentionB public: LongformerAttention(const OpKernelInfo& info); Status ComputeInternal(OpKernelContext* context) const override; + + private: + bool use_compact_memory_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu index d43b524997..4793d37b57 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu @@ -28,7 +28,7 @@ limitations under the License. #include "core/providers/cuda/cuda_common.h" #include "longformer_attention_impl.h" #include "attention_impl.h" -#include "attention_softmax.h" +#include "longformer_attention_softmax.h" using namespace onnxruntime::cuda; using namespace cub; @@ -53,11 +53,8 @@ namespace cuda { // [SoftmaxSpace: see below] [Q:BxNxSxH] [K:BxNxSxH] [V:BxNxSxH] [Global_Q:BxNxSxH] [Global_K:BxNxSxH] [Global_V:BxNxSxH] // where Global_Q, Global_K and Global_V are optional. They are not allocated when there is no global token. // -// It is feasible to use compact format for Global_Q with shape BxNxGxH to save space. We do not use compact format for now. -// // SoftmaxSpace layout: // [scratch1: (5S-3W)*W*N*B][scratch2: size_t 20] -// // Scratch1 has 5 buffers for local and global attention calculation. // Scratch2 has 5 input pointers, 5 output pointers, 5 buffer sizes and 5 strides related to scratch1. @@ -74,10 +71,17 @@ size_t GetLongformerSoftmaxWorkspaceSize( int batch_size, int num_heads, int sequence_length, - int window) { - size_t scratch1_size = GetScratch1Size(element_size, batch_size, num_heads, sequence_length, window); - size_t scratch2_size = 10 * (sizeof(void*) + sizeof(size_t)); - return scratch1_size + scratch2_size; + int window, + bool use_fast_kernel) { + if (!use_fast_kernel) { + size_t scratch1_size = GetScratch1Size(element_size, batch_size, num_heads, sequence_length, window); + size_t scratch2_size = 10 * (sizeof(void*) + sizeof(size_t)); + return scratch1_size + scratch2_size; + } else { + // Non-compact layout when environment variable ORT_LONGFORMER_COMPACT_MEMORY=0 is set. + // [scratch1: BxNxSxS] [scratch2: BxNxSxS] + return 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length); + } } size_t GetLongformerAttentionWorkspaceSize( @@ -87,8 +91,9 @@ size_t GetLongformerAttentionWorkspaceSize( int head_size, int sequence_length, int max_num_global, - int window) { - size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window); + int window, + bool use_fast_kernel) { + size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window, use_fast_kernel); size_t qkv_size = 3 * batch_size * sequence_length * num_heads * head_size * element_size; size_t global_qkv_size = max_num_global > 0 ? qkv_size : 0; return softmax_size + qkv_size + global_qkv_size; @@ -100,6 +105,7 @@ size_t GetPinnedBufferSize(int batch_size) { return sizeof(int) * batch_size + GetScratch2Size(); } +// Softmax kernel for compact format template __launch_bounds__(blockSize) __global__ void LongformerSoftmaxKernel(const int* global_attention, @@ -354,7 +360,6 @@ bool launchSoftmaxKernel( cudaStream_t stream, cublasHandle_t cublas, void* workspace, - size_t softmax_workspace_size, const void* q, // transposed Q with shape (B, N, S, H) const void* k, // transposed K with shape (B, N, S, H) const void* v, // transposed V with shape (B, N, S, H) @@ -373,10 +378,7 @@ bool launchSoftmaxKernel( int num_heads, // number of heads int head_size, // hidden size per head int window, // one sided window size - int max_num_global, // maximum number of global tokens (G) in all batches size_t element_size) { // size of element: 2 for half, and 4 for float - assert(max_num_global <= window); - const int* global_count = reinterpret_cast(pinned_buffer); bool is_fp16 = (element_size == 2); @@ -605,7 +607,10 @@ bool launchSoftmaxKernel( resultType, algo)); - void* global_q_batch = (char*)global_q + (i * elements_per_batch) * element_size; // For compact format: replace elements_per_batch by num_heads * max_num_global * head_size + // It is feasible to use compact format for Global_Q with shape BxNxGxH to save space. + // In that case, elements_per_batch is num_heads * max_num_global * head_size, and stride_per_head is max_num_global * head_size. + + void* global_q_batch = (char*)global_q + (i * elements_per_batch) * element_size; void* global_k_batch = (char*)global_k + (i * elements_per_batch) * element_size; qk_batch = (char*)input_pointers[4] + (i * buffer_sizes[4] * num_heads) * element_size; @@ -625,7 +630,7 @@ bool launchSoftmaxKernel( global_q_batch, // B Btype, // B type head_size, // ldb - stride_per_head, // strideB. For compact format: max_num_global * head_size. + stride_per_head, // strideB. beta_0, // beta qk_batch, // C Ctype, // C type @@ -827,8 +832,9 @@ bool LongformerQkvToContext( const T* global_input, const int* global_attention, const int* global_index, const int* batch_global_num, const int max_num_global, void* pinned_buffer, T* workspace, - T* output) { - size_t softmax_workspace_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window); + T* output, + size_t softmax_workspace_size, + bool use_fast_kernel) { T* qkv = reinterpret_cast((char*)workspace + softmax_workspace_size); // Number of elements in Q, K, V, Global_Q, Global_K or Global_V are same: BxNxSxH @@ -862,33 +868,62 @@ bool LongformerQkvToContext( const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); T* temp_output = qkv; // Q will be overwritten - if (!launchSoftmaxKernel( - stream, - cublas, - workspace, - softmax_workspace_size, - q, // Transposed Q with shape B x N x S x H - k, // Transposed K with shape B x N x S x H - v, // Transposed V with shape B x N x S x H - attention_mask, // Attention mask flags with shape B x S - global_q, // Transposed global Q with shape B x N x S x H. - global_k, // Transposed global K with shape B x N x S x H - global_v, // Transposed global V with shape B x N x S x H - global_attention, // Global attention flags with shape B x S - global_index, // Global index with shape B x S - batch_global_num, // Number of global token per batch with shape B x 1 - pinned_buffer, // Pinned Memory Buffer - temp_output, // Output with shape B x N x S x H - rsqrt_head_size, // Scaler - batch_size, // Batch size - sequence_length, // Sequence length - num_heads, // Number of attention heads - head_size, // Hidden size per head - window, // Half (one-sided) window size - max_num_global, // Maximum number of global tokens (G) - element_size)) { - return false; + + if (use_fast_kernel) { + if (!launchSoftmaxFastKernel( + stream, + cublas, + workspace, // softmax space + q, // transposed Q with shape (B, N, S, H) + k, // transposed K with shape (B, N, S, H) + v, // transposed V with shape (B, N, S, H) + attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked. + global_q, // Q for global tokens with shape (B, N, S, H) + global_k, // K for global tokens with shape (B, N, S, H) + global_v, // V for global tokens with shape (B, N, S, H) + global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention. + global_index, // Global index with shape (B, S) + batch_global_num, // Number of global tokens per batch with shape (B, 1) + pinned_buffer, // Pinned memory in CPU. Number of global tokens per batch with shape (B, 1) + temp_output, // output with shape (B, N, S, H) + rsqrt_head_size, // scalar + batch_size, // batch size + sequence_length, // sequence length + num_heads, // number of heads + head_size, // hidden size per head + window, // Half (one-sided) window size + element_size)) { + return false; + } + } else { + assert(max_num_global <= window); + if (!launchSoftmaxKernel( + stream, + cublas, + workspace, // softmax space + q, // Transposed Q with shape B x N x S x H + k, // Transposed K with shape B x N x S x H + v, // Transposed V with shape B x N x S x H + attention_mask, // Attention mask flags with shape B x S. Value -10000.0 means masked, and 0.0 not mased. + global_q, // Transposed global Q with shape B x N x S x H. + global_k, // Transposed global K with shape B x N x S x H + global_v, // Transposed global V with shape B x N x S x H + global_attention, // Global attention flags with shape B x S + global_index, // Global index with shape B x S + batch_global_num, // Number of global token per batch with shape B x 1 + pinned_buffer, // Pinned Memory Buffer + temp_output, // Output with shape B x N x S x H + rsqrt_head_size, // Scaler + batch_size, // Batch size + sequence_length, // Sequence length + num_heads, // Number of attention heads + head_size, // Hidden size per head + window, // Half (one-sided) window size + element_size)) { + return false; + } } + // The temp_output is BxNxSxH, transpose it to final output BxSxNxH return LaunchTransCtx(stream, sequence_length, batch_size, head_size, num_heads, temp_output, output); @@ -913,8 +948,10 @@ bool LaunchLongformerAttentionKernel( int head_size, int window, int max_num_global, - const size_t element_size) { + const size_t element_size, + bool use_fast_kernel) { CublasMathModeSetter helper(device_prop, cublas, CUBLAS_TENSOR_OP_MATH); + size_t softmax_workspace_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window, use_fast_kernel); if (element_size == 2) { return LongformerQkvToContext(cublas, stream, batch_size, sequence_length, num_heads, head_size, window, element_size, @@ -927,7 +964,9 @@ bool LaunchLongformerAttentionKernel( max_num_global, pinned_buffer, reinterpret_cast(workspace), - reinterpret_cast(output)); + reinterpret_cast(output), + softmax_workspace_size, + use_fast_kernel); } else { return LongformerQkvToContext(cublas, stream, batch_size, sequence_length, num_heads, head_size, window, element_size, @@ -940,7 +979,9 @@ bool LaunchLongformerAttentionKernel( max_num_global, pinned_buffer, reinterpret_cast(workspace), - reinterpret_cast(output)); + reinterpret_cast(output), + softmax_workspace_size, + use_fast_kernel); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h index 1cc0d80313..6530061b83 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h @@ -18,9 +18,10 @@ size_t GetLongformerAttentionWorkspaceSize( int head_size, int sequence_length, int max_num_global, - int window); + int window, + bool use_fast_kernel); - bool LaunchLongformerAttentionKernel( +bool LaunchLongformerAttentionKernel( const cudaDeviceProp& device_prop, // Device Properties cublasHandle_t& cublas, // Cublas handle cudaStream_t stream, // CUDA stream @@ -39,7 +40,8 @@ size_t GetLongformerAttentionWorkspaceSize( int head_size, // Hidden layer size per head (H) int window, // One sided attention window (W) int max_num_global, // Maximum number of global tokens (G) - const size_t element_size // Element size of input tensor + const size_t element_size, // Element size of input tensor, + bool use_fast_kernel // Use compact memory ); } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu new file mode 100644 index 0000000000..0f95df570f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu @@ -0,0 +1,657 @@ +/* +Copyright (c) NVIDIA Corporation and Microsoft Corporation + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// This is fast cuda kernels for longformer attention softmax. +// It uses two temporary matrix of BxNxSxS, and consumes more memory when sequence length is large. + +#include +#include +#include +#include +#include +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cuda_common.h" +#include "longformer_attention_softmax.h" +#include "attention_impl.h" + +using namespace onnxruntime::cuda; +using namespace cub; + +#define CHECK(expr) \ + if (!CUBLAS_CALL(expr)) { \ + return false; \ + } + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +__launch_bounds__(blockSize) + __global__ void LongformerSoftmaxFastKernel(const int* global_attention, + const int* global_index, + const int* batch_global_num, + const T* input, + const T* attention_mask, + T* output, + float scaler, + int dim0, + int sequence_length, + int attention_window) { + typedef cub::BlockReduce BlockReduce; + __shared__ typename BlockReduce::TempStorage block_reduce_temp; + __shared__ float max_shared; + __shared__ float sum_shared; + + const T* input_block = input + sequence_length * blockIdx.x; + T* output_block = output + sequence_length * blockIdx.x; + const int batch_index = blockIdx.x / dim0; + const int row_index = blockIdx.x % sequence_length; + const int global_num = batch_global_num[batch_index]; + + // To be consistent with Huggingface Longformer, the row of maksed word are set as zero. + if ((float)attention_mask[batch_index * sequence_length + row_index] < 0.0f) { + for (int i = threadIdx.x; i < sequence_length; i += blockSize) { + output_block[i] = (T)(0); + } + return; + } + + // local attention token + int col_start = 0; + int col_end = sequence_length; + bool is_local_row = (global_attention[batch_index * sequence_length + row_index] == (int)0); + if (is_local_row) { + col_start = row_index - attention_window; + if (col_start < 0) { + col_start = 0; + } + + col_end = row_index + attention_window + 1; + if (col_end > sequence_length) { + col_end = sequence_length; + } + } + + const T* mask_block = attention_mask + sequence_length * batch_index; + int tid = threadIdx.x; + + // calculate max input + float max_input = -CUDART_INF_F; + // #pragma unroll 16 + for (int i = tid + col_start; i < col_end; i += blockSize) { + float x = input_block[i]; + x = x * scaler + (float)mask_block[i]; + if (max_input < x) { + max_input = x; + } + } + + if (is_local_row) { + for (int g = tid; g < global_num; g += blockSize) { + int i = global_index[g]; + if (i < col_start || i > col_end) { + float x = input_block[i]; + x = x * scaler + (float)mask_block[i]; + if (max_input < x) { + max_input = x; + } + } + } + } + + float max_block = BlockReduce(block_reduce_temp).Reduce(max_input, cub::Max()); + if (tid == 0) { + max_shared = max_block; + } + __syncthreads(); + + float sum_input = 0.f; + // #pragma unroll 16 + for (int i = tid + col_start; i < col_end; i += blockSize) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + sum_input += x; + } + + if (is_local_row) { + for (int g = tid; g < global_num; g += blockSize) { + int i = global_index[g]; + if (i < col_start || i > col_end) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + sum_input += x; + } + } + } + + float sum_block = BlockReduce(block_reduce_temp).Reduce(sum_input, cub::Sum()); + if (tid == 0) { + sum_shared = sum_block; + } + __syncthreads(); + float recip_sum = 1.f / sum_shared; + + if (is_local_row) { + // We only need to fill in zeros for blocks that will be used in the matrix multiplication + // following the Softmax. + // + // For now zero-out only [row_index - 2*attention_window, row_index + 2*attention_window], + // we can even be more agressive and reduce the zeroing out window size since + // each row has entries in 3 blocks (3*attention_window size instead of 4*attention_window) + int zero_start = row_index - 2 * attention_window; + if (zero_start < 0) { + zero_start = 0; + } + + int zero_end = row_index + 2 * attention_window; + if (zero_end > sequence_length) { + zero_end = sequence_length; + } + + for (int i = tid + zero_start; i < zero_end; i += blockSize) { + output_block[i] = (T)(0.); + } + + for (int g = tid; g < global_num; g += blockSize) { + int i = global_index[g]; + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + output_block[i] = (T)(recip_sum * x); + } + } + + // #pragma unroll 16 + for (int i = tid + col_start; i < col_end; i += blockSize) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[i] - max_shared); + output_block[i] = (T)(recip_sum * x); + } +} + +// Launch the softmax kernel for non compact memory. +bool launchSoftmaxFastKernel( + cudaStream_t stream, + cublasHandle_t cublas, + void* workspace, // softmax space + const void* q, // transposed Q with shape (B, N, S, H) + const void* k, // transposed K with shape (B, N, S, H) + const void* v, // transposed V with shape (B, N, S, H) + const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked. + const void* global_q, // Q for global tokens with shape (B, N, S, H) + const void* global_k, // K for global tokens with shape (B, N, S, H) + const void* global_v, // V for global tokens with shape (B, N, S, H) + const int* global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention. + const int* global_index, // Global index with shape (B, S) + const int* batch_global_num, // Number of global tokens per batch with shape (B, 1) + void* pinned_buffer, // Pinned memory in CPU. Number of global tokens per batch with shape (B, 1) + void* output, // output with shape (B, N, S, H) + float scaler, // scalar + int batch_size, // batch size + int sequence_length, // sequence length + int num_heads, // number of heads + int head_size, // hidden size per head + int attention_window, // one sided windows size + size_t element_size) { // size of element: 2 for half, and 4 for float + + bool is_fp16 = (element_size == 2); + void* scratch1 = reinterpret_cast(workspace); + void* scratch2 = reinterpret_cast(scratch1) + GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length); + + // setup shared parameters for two strided batched matrix multiplies + cudaDataType_t Atype; + cudaDataType_t Btype; + cudaDataType_t Ctype; + cudaDataType_t resultType; + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT; + + __half one_fp16, zero_fp16; + float one_fp32, zero_fp32; + void *alpha, *beta_0, *beta_1; + + if (is_fp16) { + one_fp16 = __float2half(1.f); + zero_fp16 = __float2half(0.f); + alpha = static_cast(&one_fp16); + beta_0 = static_cast(&zero_fp16); + beta_1 = static_cast(&one_fp16); + Atype = CUDA_R_16F; + Btype = CUDA_R_16F; + Ctype = CUDA_R_16F; + resultType = CUDA_R_16F; + algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP; + } else { + one_fp32 = 1.f; + zero_fp32 = 0.f; + alpha = static_cast(&one_fp32); + beta_0 = static_cast(&zero_fp32); + beta_1 = static_cast(&one_fp32); + Atype = CUDA_R_32F; + Btype = CUDA_R_32F; + Ctype = CUDA_R_32F; + resultType = CUDA_R_32F; + } + + // Strided batch matrix multiply + // qk = q * k^T + // Shapes: q and k = B x N x S x H, qk = B x N x S x S + // Convert col-major to row-major by swapping q and k in Gemm + + // Local attention part + // S x S is calculated using sliding block WxW (W is one sided window size) like the following: + // [W][W] + // [W][W][W] + // [W][W][W] + // [W][W] + // The first and last rows have 2 blocks, and the remaining has 3 blocks per row. + // The calculation are splited into 3 parts. Firstly, fill the middle rows, then the first row and finally the last row. + // The results are stored in scratch1. + + int w = attention_window; + int x_offset = num_heads * sequence_length * head_size; + int y_offset = num_heads * sequence_length * sequence_length; + int last_block = (sequence_length / w) - 1; + int strideA = sequence_length * head_size; + int strideB = sequence_length * head_size; + int strideC = sequence_length * sequence_length; + + // When S == 2W, there is no middle rows of blocks: + // [W][W] + // [W][W] + // We can use normal matrix multiplication in this case. + if (sequence_length == 2 * w) { + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + sequence_length, + sequence_length, + head_size, + alpha, + k, + Atype, + head_size, + sequence_length * head_size, + q, + Btype, + head_size, + sequence_length * head_size, + beta_0, + scratch1, + Ctype, + sequence_length, + sequence_length * sequence_length, + batch_size * num_heads, + resultType, + algo)); + } else { // sequence_length > 2 * w + for (int i = 0; i < batch_size; ++i) { + for (int j = 0; j < num_heads; ++j) { + void* q_head = (char*)q + (i * x_offset + j * sequence_length * head_size + w * head_size) * element_size; + void* k_head = (char*)k + (i * x_offset + j * sequence_length * head_size) * element_size; + void* qk_head = (char*)scratch1 + (i * y_offset + j * sequence_length * sequence_length + w * sequence_length) * element_size; + int count = (sequence_length - 2 * w) / w; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + 3 * w, // m + w, // n + head_size, // k + alpha, // alpha + k_head, // A + Atype, // A type + head_size, // lda + w * head_size, // strideA + q_head, // B + Btype, // B type + head_size, // ldb + w * head_size, // strideB + beta_0, // beta + qk_head, // C + Ctype, // C type + sequence_length, // ldc + sequence_length * w + w, // strideC + count, // batch count + resultType, + algo)); + } + } + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + 2 * w, // m + w, // n + head_size, // k + alpha, // alpha + k, // A + Atype, // A type + head_size, // lda + strideA, // strideA + q, // B + Btype, // B type + head_size, // ldb + strideB, // strideB + beta_0, // beta + scratch1, // C + Ctype, // C type + sequence_length, // ldc + strideC, // strideC + batch_size * num_heads, // batch count + resultType, + algo)); + + void* q_head = (char*)q + (last_block * w * head_size) * element_size; + void* k_head = (char*)k + ((last_block - 1) * w * head_size) * element_size; + void* qk_head = (char*)scratch1 + (last_block * w * sequence_length + (last_block - 1) * w) * element_size; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + 2 * w, + w, + head_size, + alpha, + k_head, + Atype, + head_size, + strideA, + q_head, + Btype, + head_size, + strideB, + beta_0, + qk_head, + Ctype, + sequence_length, + strideC, + batch_size * num_heads, + resultType, + algo)); + } + + const int* batch_global_count = reinterpret_cast(pinned_buffer); + // Global attention part + for (int i = 0; i < batch_size; ++i) { + if (batch_global_count[i] > 0) { + void* q_batch = (char*)q + (i * x_offset) * element_size; + void* k_batch = (char*)k + (i * x_offset) * element_size; + void* qk_batch = (char*)scratch1 + (i * y_offset) * element_size; + // Local tokens attending global tokens + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + batch_global_count[i], + sequence_length, + head_size, + alpha, + k_batch, + Atype, + head_size, + strideA, + q_batch, + Btype, + head_size, + strideB, + beta_0, + qk_batch, + Ctype, + sequence_length, + strideC, + num_heads, + resultType, + algo)); + + void* global_q_batch = (char*)global_q + (i * num_heads * sequence_length * head_size) * element_size; + void* global_k_batch = (char*)global_k + (i * x_offset) * element_size; + int strideB_global = sequence_length * head_size; + + // Global tokens attending everything + // This GEMMs need to be last to make sure all global token entries are re-written. + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_T, + CUBLAS_OP_N, + sequence_length, + batch_global_count[i], + head_size, + alpha, + global_k_batch, + Atype, + head_size, + strideA, + global_q_batch, + Btype, + head_size, + strideB_global, + beta_0, + qk_batch, + Ctype, + sequence_length, + strideC, + num_heads, + resultType, + algo)); + } + } + + int dim0 = sequence_length * num_heads; + int dim1 = sequence_length; + void* softmax_out = scratch2; + + const int blockSize = 64; + const int gridSize = batch_size * num_heads * sequence_length; + if (is_fp16) { + LongformerSoftmaxFastKernel<__half, blockSize><<>>( + global_attention, + global_index, + batch_global_num, + static_cast(scratch1), + static_cast(attention_mask), + static_cast<__half*>(softmax_out), scaler, dim0, dim1, attention_window); + } else { + LongformerSoftmaxFastKernel<<>>( + global_attention, + global_index, + batch_global_num, + static_cast(scratch1), + static_cast(attention_mask), + static_cast(softmax_out), scaler, dim0, dim1, attention_window); + } + + // Run the matrix multiply: output = softmax_out * v + // softmax_out: B x N x S x S + // v: B x N x S x H + // attn_out: B x N x S x H + // Calculation uses full Gemm (S == 2W) or sliding blocks (S > 2W) in a way similar to local attention part. + + if (sequence_length == 2 * w) { + // convert col-major to row-major by swapping softmax_out and v + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + sequence_length, + sequence_length, + alpha, + v, + Atype, + head_size, + sequence_length * head_size, + softmax_out, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, + output, + Ctype, + head_size, + sequence_length * head_size, + batch_size * num_heads, + resultType, + algo)); + } else { // sequence_length > 2 * w + for (int i = 0; i < batch_size; ++i) { + for (int j = 0; j < num_heads; ++j) { + void* v_head = (char*)v + (i * x_offset + j * head_size * sequence_length) * element_size; + void* prob_head = (char*)softmax_out + (i * y_offset + j * sequence_length * sequence_length + w * sequence_length) * element_size; + void* out_head = (char*)output + (i * x_offset + j * head_size * sequence_length + w * head_size) * element_size; + int count = (sequence_length - 2 * w) / w; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + w, + 3 * w, + alpha, + v_head, + Atype, + head_size, + w * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * w + w, + beta_0, + out_head, + Ctype, + head_size, + w * head_size, + count, + resultType, + algo)); + } + } + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + w, + 2 * w, + alpha, + v, + Atype, + head_size, + sequence_length * head_size, + softmax_out, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, + output, + Ctype, + head_size, + sequence_length * head_size, + batch_size * num_heads, + resultType, + algo)); + + void* v_head = (char*)v + (last_block - 1) * w * head_size * element_size; + void* prob_head = (char*)softmax_out + (sequence_length * last_block * w + (last_block - 1) * w) * element_size; + void* out_head = (char*)output + last_block * w * head_size * element_size; + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + w, + 2 * w, + alpha, + v_head, + Atype, + head_size, + sequence_length * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, + out_head, + Ctype, + head_size, + sequence_length * head_size, + batch_size * num_heads, + resultType, + algo)); + } + + for (int i = 0; i < batch_size; ++i) { + if (batch_global_count[i] > 0) { + int glob_longdim_mm = (last_block - 1) * w; + + void* v_head = (char*)v + (i * x_offset) * element_size; + void* prob_head = (char*)softmax_out + (i * y_offset + 2 * w * sequence_length) * element_size; + void* out_head = (char*)output + (i * x_offset + 2 * w * head_size) * element_size; + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + glob_longdim_mm, + batch_global_count[i], + alpha, + v_head, + Atype, + head_size, + sequence_length * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_1, + out_head, + Ctype, + head_size, + sequence_length * head_size, + num_heads, + resultType, + algo)); + + // Global tokens + v_head = (char*)global_v + (i * x_offset) * element_size; + prob_head = (char*)softmax_out + (i * y_offset) * element_size; + out_head = (char*)output + (i * x_offset) * element_size; + + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, + batch_global_count[i], + sequence_length, // Re-write entries completely + alpha, + v_head, + Atype, + head_size, + sequence_length * head_size, + prob_head, + Btype, + sequence_length, + sequence_length * sequence_length, + beta_0, // Use beta=0 to overwrite + out_head, // Here assumes global tokens are at the beginning of sequence. + Ctype, + head_size, + sequence_length * head_size, + num_heads, + resultType, + algo)); + } + } + + return true; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h new file mode 100644 index 0000000000..ade76f556e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h @@ -0,0 +1,51 @@ +/* +Copyright (c) NVIDIA Corporation and Microsoft Corporation + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// This is fast cuda kernels for longformer attention softmax. +// It uses two temporary matrix of BxNxSxS, and consumes more memory when sequence length is large. + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Launch the softmax kernel for non compact memory. +bool launchSoftmaxFastKernel( + cudaStream_t stream, + cublasHandle_t cublas, + void* workspace, // softmax space + const void* q, // transposed Q with shape (B, N, S, H) + const void* k, // transposed K with shape (B, N, S, H) + const void* v, // transposed V with shape (B, N, S, H) + const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked. + const void* global_q, // Q for global tokens with shape (B, N, S, H) + const void* global_k, // K for global tokens with shape (B, N, S, H) + const void* global_v, // V for global tokens with shape (B, N, S, H) + const int* global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention. + const int* global_index, // Global index with shape (B, S) + const int* batch_global_num, // Number of global tokens per batch with shape (B, 1) + void* pinned_buffer, // Pinned memory in CPU. Number of global tokens per batch with shape (B, 1) + void* output, // output with shape (B, N, S, H) + float scaler, // scalar + int batch_size, // batch size + int sequence_length, // sequence length + int num_heads, // number of heads + int head_size, // hidden size per head + int attention_window, // one sided windows size + size_t element_size); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/transformers/benchmark_helper.py b/onnxruntime/python/tools/transformers/benchmark_helper.py index 22a2d8fe36..626e7be65b 100644 --- a/onnxruntime/python/tools/transformers/benchmark_helper.py +++ b/onnxruntime/python/tools/transformers/benchmark_helper.py @@ -235,3 +235,85 @@ def allocateOutputBuffers(output_buffers, output_buffer_max_sizes, device): for i in output_buffer_max_sizes: output_buffers.append(torch.empty(i, dtype=torch.float32, device=device)) + + +def set_random_seed(seed=123): + """Set random seed manully to get deterministic results""" + import random + random.seed(seed) + numpy.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + #torch.backends.cudnn.enabled = False + #torch.backends.cudnn.benchmark = False + #torch.backends.cudnn.deterministic = True + + +def measure_memory(is_gpu, func): + import os + import psutil + from time import sleep + + class MemoryMonitor: + def __init__(self, keep_measuring=True): + self.keep_measuring = keep_measuring + + def measure_cpu_usage(self): + max_usage = 0 + while True: + max_usage = max(max_usage, psutil.Process(os.getpid()).memory_info().rss / 1024**2) + sleep(0.005) # 5ms + if not self.keep_measuring: + break + return max_usage + + def measure_gpu_usage(self): + from py3nvml.py3nvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, \ + nvmlDeviceGetMemoryInfo, nvmlDeviceGetName, nvmlShutdown, NVMLError + max_gpu_usage = [] + gpu_name = [] + try: + nvmlInit() + deviceCount = nvmlDeviceGetCount() + max_gpu_usage = [0 for i in range(deviceCount)] + gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(deviceCount)] + while True: + for i in range(deviceCount): + info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i)) + max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2) + sleep(0.005) # 5ms + if not self.keep_measuring: + break + nvmlShutdown() + return [{ + "device_id": i, + "name": gpu_name[i], + "max_used_MB": max_gpu_usage[i] + } for i in range(deviceCount)] + except NVMLError as error: + if not self.silent: + self.logger.error("Error fetching GPU information using nvml: %s", error) + return None + + monitor = MemoryMonitor(False) + if is_gpu: + print(f"GPU memory usage before testing: {monitor.measure_gpu_usage()}") + else: + print(f"Peak CPU memory usage before testing: {monitor.measure_cpu_usage():.2f} MB") + + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor() as executor: + monitor = MemoryMonitor() + mem_thread = executor.submit(monitor.measure_gpu_usage if is_gpu else monitor.measure_cpu_usage) + try: + fn_thread = executor.submit(func) + result = fn_thread.result() + finally: + monitor.keep_measuring = False + max_usage = mem_thread.result() + if is_gpu: + print(f"Peak GPU memory usage: {max_usage}") + else: + print(f"Peak CPU memory usage: {max_usage:.2f} MB") + return max_usage diff --git a/onnxruntime/python/tools/transformers/longformer/__init__.py b/onnxruntime/python/tools/transformers/longformer/__init__.py index ad5632855c..864a0f3b5b 100644 --- a/onnxruntime/python/tools/transformers/longformer/__init__.py +++ b/onnxruntime/python/tools/transformers/longformer/__init__.py @@ -1,3 +1,9 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + import os import sys diff --git a/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py b/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py index 80f9d44a70..30b13d4e8b 100644 --- a/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py +++ b/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py @@ -1,11 +1,29 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +# +# This script run benchmark of latency or peak memory usage of Longformer model inference. +# # Please run convert_longformer_to_onnx.py to get onnx model before running this script. -# Tested with python 3.7, onnxruntime-gpu 1.6.0 (or nightly), PyTorch 1.7.0, transformers 4.0, CUDA 10.2, CUDNN 8.0 -# Example step by step command lines for benchmarking longformer base model (without/with optimizer) in Linux: +# Tested with python 3.6, onnxruntime-gpu 1.7.0, PyTorch 1.7.1, transformers 4.3.2, CUDA 10.2. +# +# Example commands for exporting longformer base model in Linux or WSL: +# cd ../torch_extensions # python setup.py install -# python convert_longformer_to_onnx.py -m longformer-base-4096 -# python benchmark_longformer.py -m longformer-base-4096 -# python convert_longformer_to_onnx.py -m longformer-base-4096 -o -# python benchmark_longformer.py -m longformer-base-4096 +# cd ../longformer +# python convert_longformer_to_onnx.py --model longformer-base-4096 --precision fp32 --optimize_onnx +# +# Benchmark the latency (Exported onnx model is in the current directory): +# python benchmark_longformer.py --models longformer-base-4096 --batch_sizes 1 --sequence_lengths 512 1024 2048 4096 --global_lengths 8 --onnx_dir . --validate_onnx -t 100 +# +# Benchmark GPU peak memory: +# export ORT_LONGFORMER_COMPACT_MEMORY=0 +# python benchmark_longformer.py --models longformer-base-4096 --batch_sizes 1 --sequence_lengths 4096 --global_lengths 8 --onnx_dir . --memory -t 10 +# export ORT_LONGFORMER_COMPACT_MEMORY=1 +# python benchmark_longformer.py --models longformer-base-4096 --batch_sizes 1 --sequence_lengths 4096 --global_lengths 8 --onnx_dir . --memory -t 10 +# By default, compact memory kernel is not enabled since it is slower. You need set an environment variable ORT_LONGFORMER_COMPACT_MEMORY=1 to enable it, which uses less memory in this test. import timeit from datetime import datetime @@ -17,57 +35,30 @@ import torch import onnxruntime import numpy as np import pprint +import math -# Mapping from model name to pretrained model name -MODELS = { - "longformer-base-4096": "allenai/longformer-base-4096", - "longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging -} - -is_debug = False +from longformer_helper import LongformerHelper, PRETRAINED_LONGFORMER_MODELS sys.path.append(os.path.join(os.path.dirname(__file__), '..')) import benchmark_helper -def get_dummy_inputs(batch_size, sequence_length, num_global_tokens, device): - input_ids = torch.randint(low=0, high=100, size=(batch_size, sequence_length), dtype=torch.long, device=device) - attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device) - global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device) - global_token_index = list(range(num_global_tokens)) - global_attention_mask[:, global_token_index] = 1 - return input_ids, attention_mask, global_attention_mask - - -def diff_outputs(ort_outputs, torch_outputs): - max_diff = [] - # Compare the outputs to find max difference - for i in range(2): - print(f"output {i} shape: ORT={ort_outputs[i].shape}, Torch={torch_outputs[i].shape}") - diff = (torch.from_numpy(ort_outputs[i]) - torch_outputs[i].to('cpu')).abs().max() - max_diff.append(diff) - print(f"max diff for output: {max_diff}") - return max_diff - - -def test_torch(device, model, model_name, batch_sizes, sequence_lengths, global_lengths, test_times, num_threads, - verbose): +def test_torch_latency(device, model, model_name, batch_sizes, sequence_lengths, global_lengths, test_times, + num_threads, verbose): if num_threads > 0: torch.set_num_threads(num_threads) results = [] for batch_size in batch_sizes: - for sequence_length in sequence_lengths: # This is total length of . - for global_length in global_lengths: # This is length of . Short query (8) for search keywords, and longer query (16) for question like + for sequence_length in sequence_lengths: + for global_length in global_lengths: print(f"batch_size={batch_size} sequence_length={sequence_length} global_length={global_length}...") - input_ids, attention_mask, global_attention_mask = get_dummy_inputs(batch_size, sequence_length, - global_length, device) + inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, global_length, + device) + input_list = inputs.to_list() - # Run PyTorch - _ = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask) - runtimes = timeit.repeat(lambda: model(input_ids, attention_mask, global_attention_mask), - repeat=test_times, - number=1) + _ = model(*input_list) + runtimes = timeit.repeat(lambda: model(*input_list), repeat=test_times, number=1) result = { "engine": "torch", #TODO: test torchscript "version": torch.__version__, @@ -90,7 +81,24 @@ def test_torch(device, model, model_name, batch_sizes, sequence_lengths, global_ return results -def test_onnxruntime(device, +def test_parity(device, model, ort_session, batch_size, sequence_length, global_length, verbose=True): + print( + f"Comparing Torch and ORT outputs for batch_size={batch_size} sequence_length={sequence_length} global_length={global_length}..." + ) + dummy_inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, global_length, + device) + ort_inputs = dummy_inputs.get_ort_inputs() + ort_outputs = ort_session.run(None, ort_inputs) + input_list = dummy_inputs.to_list() + torch_outputs = model(*input_list) + max_diff = np.amax(torch_outputs[0].cpu().numpy() - ort_outputs[0]) + print(f"last_state max diff = {max_diff}") + if verbose and (math.isnan(max_diff) or max_diff > 0.001): + print("torch last_state:", torch_outputs[0]) + print("ort last_state:", ort_outputs[0]) + + +def test_ort_latency(device, model, model_name, ort_session, @@ -101,43 +109,30 @@ def test_onnxruntime(device, num_threads, optimizer=False, precision='fp32', + validate_onnx=True, + disable_io_binding=False, verbose=True): results = [] for batch_size in batch_sizes: - for sequence_length in sequence_lengths: # This is total length of . - for global_length in global_lengths: # This is length of . Short query (8) for search keywords, and longer query (16) for question like + for sequence_length in sequence_lengths: + for global_length in global_lengths: + assert global_length <= model.config.attention_window[ + 0], "Limitation of current implementation: number of global token <= attention_window" print( - f"Testing batch_size={batch_size} sequence_length={sequence_length} global_length={global_length} optimizer={optimizer}, precision={precision}..." + f"Testing batch_size={batch_size} sequence_length={sequence_length} global_length={global_length} optimizer={optimizer}, precision={precision} io_binding={not disable_io_binding}..." ) - input_ids, attention_mask, global_attention_mask = get_dummy_inputs(batch_size, sequence_length, - global_length, device) + dummy_inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, + global_length, device) # Run OnnxRuntime - ort_inputs = { - "input_ids": input_ids.cpu().numpy(), - "attention_mask": attention_mask.cpu().numpy(), - "global_attention_mask": global_attention_mask.cpu().numpy() - } + ort_inputs = dummy_inputs.get_ort_inputs() if verbose: - pprint.pprint(ort_inputs) + print(ort_inputs) # run one query for warm up ort_outputs = ort_session.run(None, ort_inputs) - if verbose: - # Run PyTorch then compare the results with OnnxRuntime. - torch_outputs = model(input_ids, - attention_mask=attention_mask, - global_attention_mask=global_attention_mask) - - max_diff = diff_outputs(ort_outputs, torch_outputs) - print("max diff for outputs", max_diff) - if max(max_diff) > 0.001: - print("ort_inputs", ort_inputs) - print("ort_outputs", ort_outputs) - - device = input_ids.device result_template = { "model_name": model_name, "inputs": 3, @@ -154,32 +149,59 @@ def test_onnxruntime(device, "datetime": str(datetime.now()), } - max_last_state_size = max(batch_sizes) * max(sequence_lengths) * model.config.hidden_size - max_pooler_size = max(batch_sizes) * max(sequence_lengths) - """ - result = benchmark_helper.inference_ort_with_io_binding( - ort_session, - ort_inputs, - result_template=result_template, - repeat_times=test_times, - ort_output_names=["last_state", "pooler"], - ort_outputs=ort_outputs, - output_buffers=[], - output_buffer_max_sizes=[max_last_state_size, max_pooler_size], - batch_size=batch_size, - device=device) - """ - result = benchmark_helper.inference_ort(ort_session, - ort_inputs, - result_template=result_template, - repeat_times=test_times, - batch_size=batch_size) + if not disable_io_binding: + max_last_state_size = max(batch_sizes) * max(sequence_lengths) * model.config.hidden_size + max_pooler_size = max(batch_sizes) * max(sequence_lengths) + result = benchmark_helper.inference_ort_with_io_binding( + ort_session, + ort_inputs, + result_template=result_template, + repeat_times=test_times, + ort_output_names=["last_state", "pooler"], + ort_outputs=ort_outputs, + output_buffers=[], + output_buffer_max_sizes=[max_last_state_size, max_pooler_size], + batch_size=batch_size, + device=device, + data_type=np.longlong, #input data type + ) + else: + result = benchmark_helper.inference_ort(ort_session, + ort_inputs, + result_template=result_template, + repeat_times=test_times, + batch_size=batch_size) pprint.pprint(result) results.append(result) + + if validate_onnx: + test_parity(device, model, ort_session, batch_size, sequence_length, global_length, verbose) + return results +def test_ort_memory(device, onnx_model_path, batch_size, sequence_length, global_length, test_times, num_threads): + print( + f"Testing memory for model={onnx_model_path}, batch_size={batch_size}, sequence_length={sequence_length}, global_length={global_length}, test_times={test_times}, num_threads={num_threads}" + ) + + def inference(): + session = benchmark_helper.create_onnxruntime_session(onnx_model_path, + use_gpu=True, + enable_all_optimization=True, + num_threads=num_threads) + + dummy_inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, global_length, + device) + ort_inputs = dummy_inputs.get_ort_inputs() + for _ in range(test_times): + ort_outputs = session.run(None, ort_inputs) + + benchmark_helper.measure_memory(is_gpu=True, func=inference) + print("Memory test is done") + + def test_all(args): # Currently, the longformer attention operator could only run in GPU (no CPU implementation yet). device = torch.device('cuda:0') @@ -188,17 +210,18 @@ def test_all(args): for model_name in args.models: # Here we run an example input from transformers import LongformerModel - torch_model_name_or_dir = MODELS[model_name] + torch_model_name_or_dir = PRETRAINED_LONGFORMER_MODELS[model_name] model = LongformerModel.from_pretrained(torch_model_name_or_dir) # pretrained model name or directory model.to(device) # Search onnx model in the following order: optimized fp16 model, optimized fp32 model, raw model + # TODO: call convert_longformer_to_onnx to export onnx instead. + import os.path optimized = False precision = 'fp32' - onnx_model_path = model_name + ".onnx" - optimized_fp32_model = model_name + "_fp32.onnx" - optimized_fp16_model = model_name + "_fp16.onnx" - import os.path + onnx_model_path = os.path.join(args.onnx_dir, model_name + ".onnx") + optimized_fp32_model = os.path.join(args.onnx_dir, model_name + "_fp32.onnx") + optimized_fp16_model = os.path.join(args.onnx_dir, model_name + "_fp16.onnx") if os.path.isfile(optimized_fp16_model): onnx_model_path = optimized_fp16_model optimized = True @@ -206,24 +229,33 @@ def test_all(args): elif os.path.isfile(optimized_fp32_model): onnx_model_path = optimized_fp32_model optimized = True + print("ONNX model path:", onnx_model_path) for num_threads in args.num_threads: if "torch" in args.engines: - results += test_torch(device, model, model_name, args.batch_sizes, args.sequence_lengths, - args.global_lengths, args.test_times, num_threads, args.verbose) + results += test_torch_latency(device, model, model_name, args.batch_sizes, args.sequence_lengths, + args.global_lengths, args.test_times, num_threads, args.verbose) if "onnxruntime" in args.engines: - session = benchmark_helper.create_onnxruntime_session(onnx_model_path, - use_gpu=True, - enable_all_optimization=True, - num_threads=num_threads) - results += test_onnxruntime(device, model, model_name, session, args.batch_sizes, args.sequence_lengths, - args.global_lengths, args.test_times, num_threads, optimized, precision, - args.verbose) + if args.memory: + test_ort_memory(device, onnx_model_path, args.batch_sizes[0], args.sequence_lengths[0], + args.global_lengths[0], args.test_times, num_threads) + else: # test latency + session = benchmark_helper.create_onnxruntime_session(onnx_model_path, + use_gpu=True, + enable_all_optimization=True, + num_threads=num_threads) + if session is None: + raise RuntimeError(f"Failed to create ORT sesssion from ONNX file {onnx_model_path}") + + results += test_ort_latency(device, model, model_name, session, args.batch_sizes, + args.sequence_lengths, args.global_lengths, args.test_times, + num_threads, optimized, precision, args.validate_onnx, + args.disable_io_binding, args.verbose) return results -def parse_arguments(): +def parse_arguments(argv=None): parser = argparse.ArgumentParser() parser.add_argument("-m", @@ -231,9 +263,9 @@ def parse_arguments(): required=False, nargs="+", type=str, - default=["longformer-random-tiny"] if is_debug else ["longformer-base-4096"], - choices=list(MODELS.keys()), - help="Pre-trained models in the list: " + ", ".join(MODELS.keys())) + default=["longformer-base-4096"], + help="Checkpoint directory or pre-trained model names in the list: " + + ", ".join(PRETRAINED_LONGFORMER_MODELS.keys())) parser.add_argument("-e", "--engines", @@ -253,21 +285,52 @@ def parse_arguments(): parser.add_argument("-b", "--batch_sizes", nargs="+", type=int, default=[1]) - # If multiple of window size is used during exporting onnx model, there is no padding in ONNX model so you will need padding by yourself before running onnx model. - # In that case, you can only test sequence length that is multiple of window size (4 or 512 for these two models). - parser.add_argument("-s", - "--sequence_lengths", + # If --export_padding is not used in exporting onnx model, there is no padding in ONNX model so you will need padding inputs by yourself before running onnx model. + # In that case, you can only test sequence length that is multiple of attention window size. + parser.add_argument( + "-s", + "--sequence_lengths", + nargs="+", + type=int, + default=[512, 1024, 2048, 4096], + help= + "Sequence lengths. It could have multiple values in latency test. If --export_padding is not used in exporting onnx model, sequence length shall be multiple of window size." + ) + + parser.add_argument("--onnx_dir", + required=False, + type=str, + default=os.path.join('.', 'onnx_models'), + help="Directory to search onnx models.") + + parser.add_argument("-g", + "--global_lengths", nargs="+", type=int, - default=[4] if is_debug else [512, 1024, 2048, 4096]) + default=[0], + help="Number of global tokens. It could have multiple values in latency test.") - parser.add_argument("-g", "--global_lengths", nargs="+", type=int, default=[1] if is_debug else [8]) + parser.add_argument("-n", + "--num_threads", + required=False, + nargs="+", + type=int, + default=[0], + help="Threads to use. It could have multiple values in latency test.") - parser.add_argument("-n", "--num_threads", required=False, nargs="+", type=int, default=[0], help="Threads to use") + parser.add_argument("-v", + "--validate_onnx", + required=False, + action="store_true", + help="Validate that ONNX model generates same output as PyTorch model.") - parser.add_argument("--verbose", required=False, action="store_true", help="Print more information") + parser.add_argument("--disable_io_binding", required=False, action="store_true", help="Do not use IO Binding.") - args = parser.parse_args() + parser.add_argument("--memory", required=False, action="store_true", help="Test memory usage instead of latency.") + + parser.add_argument("--verbose", required=False, action="store_true", help="Print more information.") + + args = parser.parse_args(argv) return args @@ -327,16 +390,27 @@ def output_details(results, csv_filename): print(f"Detail results are saved to csv file: {csv_filename}") -def main(): - args = parse_arguments() - +def main(args): assert len(args.models) == 1, "run only one model at a time" + if args.memory: + if len(args.batch_sizes) > 1: + raise RuntimeError("For memory test, only one batch_size (-b) is allowed.") + if len(args.sequence_lengths) > 1: + raise RuntimeError("For memory test, only one sequence_length (-s) is allowed.") + if len(args.global_lengths) > 1: + raise RuntimeError("For memory test, only one global_length (-g) is allowed.") + if len(args.num_threads) > 1: + raise RuntimeError("For memory test, only one value of --num_threads is allowed.") + if not torch.cuda.is_available(): raise RuntimeError("Please install PyTorch with Cuda, and use a machine with GPU for testing gpu performance.") torch.set_grad_enabled(False) + # set random seed manully to get deterministic results + #benchmark_helper.set_random_seed(123) + all_results = test_all(args) time_stamp = datetime.now().strftime("%Y%m%d-%H%M%S") @@ -348,4 +422,7 @@ def main(): if __name__ == "__main__": - main() + args = parse_arguments() + #args = parse_arguments("-e onnxruntime -t 1 -b 1 -s 4 -g 2 --onnx_dir . -t 1 -m longformer-random-tiny".split(' ')) + + main(args) diff --git a/onnxruntime/python/tools/transformers/longformer/convert_longformer_to_onnx.py b/onnxruntime/python/tools/transformers/longformer/convert_longformer_to_onnx.py index d0cdea4302..b827d05c3a 100644 --- a/onnxruntime/python/tools/transformers/longformer/convert_longformer_to_onnx.py +++ b/onnxruntime/python/tools/transformers/longformer/convert_longformer_to_onnx.py @@ -1,9 +1,19 @@ -# Before running this script, please run "python setup.py install" in ../torch_extensions to build longformer_attention.cpp -# under a python environment with PyTorch installed. Then you can update the path of longformer_attention.cpython-*.so -# and run this script in same environment. -# Tested in Ubuntu 18.04, python 3.6, PyTorch 1.7.1, transformers 4.3.0. -# GPU is not needed for this script. You can run it in CPU. -# For inference of the onnx model, you will need latest onnxruntime-gpu 1.7.0 or above. +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +# This script converts Longformer model from huggingface transformers 4.0 or later to ONNX. +# Unlike normal ONNX model exporting, it will directly translate LongformerSelfAttention to the LongformerAttention operator in ONNX Runtime. +# +# Before running this script, please run "python setup.py install" in ../torch_extensions under Linux with PyTorch installed. +# Then you can update the path of longformer_attention.cpython-*.so and run this script in same environment. +# +# It is tested in Ubuntu 18.04, python 3.6, PyTorch 1.7.1, transformers 4.3.0 or 4.3.2. +# GPU is not needed for this script. You can run it in CPU. For --optimize_onnx, you can use either onnxruntime or onnxruntime-gpu package. +# +# For inference of the onnx model, you will need onnxruntime-gpu 1.7.0 or above. import torch import numpy as np @@ -13,6 +23,8 @@ from torch.onnx import register_custom_op_symbolic from torch.onnx.symbolic_helper import parse_args from packaging import version +from longformer_helper import LongformerHelper, PRETRAINED_LONGFORMER_MODELS + @parse_args('v', 'v', 'v', 'v', 'v', 'v', 'v', 'i', 'i') def my_longformer_attention(g, input, weight, bias, mask, global_weight, global_bias, global_mask, num_heads, window): @@ -31,18 +43,10 @@ def my_longformer_attention(g, input, weight, bias, mask, global_weight, global_ # namespace is onnxruntime which is registered in longformer_attention.cpp register_custom_op_symbolic('onnxruntime::LongformerAttention', my_longformer_attention, 9) -# TODO: update the path according to output of "python setup.py install" when your python version is not 3.6 +# TODO: search the directory to find correct output filename of "python setup.py install" when python version is not 3.6 torch.ops.load_library( r'../torch_extensions/build/lib.linux-x86_64-3.6/longformer_attention.cpython-36m-x86_64-linux-gnu.so') -# mapping from model name to pretrained model name -MODELS = { - "longformer-base-4096": "allenai/longformer-base-4096", - "longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging -} - -is_debug = False - def parse_arguments(): parser = argparse.ArgumentParser() @@ -51,21 +55,24 @@ def parse_arguments(): "--model", required=False, type=str, - default="longformer-random-tiny" if is_debug else "longformer-base-4096", - choices=list(MODELS.keys()), - help="Pre-trained models in the list: " + ", ".join(MODELS.keys())) + default="longformer-base-4096", + help="Checkpoint directory or pre-trained model names in the list: " + + ", ".join(PRETRAINED_LONGFORMER_MODELS.keys())) - # Sequence length shall choose properly. - # If multiple of windows size is used, there is no padding in ONNX model so you will need padding by yourself before running onnx model. - parser.add_argument("-s", "--sequence_length", type=int, default=4 if is_debug else 512) - - parser.add_argument("-g", "--global_length", type=int, default=1 if is_debug else 8) + parser.add_argument( + '--export_padding', + required=False, + action='store_true', + help= + 'Export padding logic to ONNX graph. If not enabled, user need pad input so that sequence length is multiple of window size.' + ) + parser.set_defaults(export_padding=False) parser.add_argument('-o', '--optimize_onnx', required=False, action='store_true', - help='Use optimizer.py to optimize onnx model') + help='Use optimizer.py to optimize onnx model.') parser.set_defaults(optimize_onnx=False) parser.add_argument("-p", @@ -80,36 +87,77 @@ def parse_arguments(): return args -def get_dummy_inputs(sequence_length, num_global_tokens, device): +# Create a dummy input for ONNX export. +def get_dummy_inputs(config, export_padding, device): + + # When sequence length is multiple of windows size, there is no padding logic in ONNX graph + sequence_length = config.attention_window[0] + 1 if export_padding else config.attention_window[0] + # Create dummy inputs input_ids = torch.arange(sequence_length).unsqueeze(0).to(device) - attention_mask = torch.ones(input_ids.shape, dtype=torch.long, - device=input_ids.device) # TODO: use random word ID. #TODO: simulate masked word - global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=input_ids.device) - if num_global_tokens > 0: - global_token_index = list(range(num_global_tokens)) - global_attention_mask[:, global_token_index] = 1 - # TODO: support more inputs like token_type_ids, position_ids + + attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device) + attention_mask[:, sequence_length - 1] = 0 # last token is masked + + global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device) + global_attention_mask[:, 0] = 1 # first token is global token + return input_ids, attention_mask, global_attention_mask -args = parse_arguments() - -model_name = args.model -onnx_model_path = model_name + ".onnx" - -from transformers import LongformerModel -model = LongformerModel.from_pretrained(MODELS[model_name]) # pretrained model name or directory - -input_ids, attention_mask, global_attention_mask = get_dummy_inputs(sequence_length=args.sequence_length, - num_global_tokens=args.global_length, - device=torch.device('cpu')) - -example_outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask) - - # A new function to replace LongformerSelfAttention.forward -#For transformer 4.3 +# For transformers 4.0.0 +def my_longformer_self_attention_forward_4(self, + hidden_states, + attention_mask=None, + is_index_masked=None, + is_index_global_attn=None, + is_global_attn=None): + global_mask = is_index_global_attn.int() + # The following check is based on the dummy inputs (only the first token is global). + assert len(global_mask.shape) == 2 and global_mask.shape[0] == 1 and global_mask.count_nonzero().item( + ) == 1 and global_mask.tolist()[0][0] == 1 + + input_mask = is_index_masked.float() + input_mask = input_mask.masked_fill(is_index_masked, -10000.0) + # Yet another way to generate input_mask = torch.masked_fill(attention_mask, is_index_global_attn, 0.0) + + # TODO: add postprocess of ONNX model to calculate based on graph input: input_mask = (attention_mask - 1) * 10000.0 + # TODO: add postprocess of ONNX model to use graph input directly: glboal_mask = global_attention_mask + + # The following check is based on the dummy inputs (only the last token is masked). + assert len(input_mask.shape) == 2 and input_mask.shape[0] == 1 and input_mask.count_nonzero().item( + ) == 1 and input_mask.tolist()[0][-1] == -10000.0 + + weight = torch.stack( + (self.query.weight.transpose(0, 1), \ + self.key.weight.transpose(0, 1), \ + self.value.weight.transpose(0, 1)), dim=1) + weight = weight.reshape(self.embed_dim, 3 * self.embed_dim) + + bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0) + bias = bias.reshape(3 * self.embed_dim) + + global_weight = torch.stack((self.query_global.weight.transpose(0, 1), \ + self.key_global.weight.transpose(0, 1), \ + self.value_global.weight.transpose(0, 1)), + dim=1) + global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim) + + global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0) + global_bias = global_bias.reshape(3 * self.embed_dim) + + attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, input_mask, global_weight, + global_bias, global_mask, self.num_heads, + self.one_sided_attn_window_size) + + assert attn_output.size() == hidden_states.size(), "Unexpected size" + + outputs = (attn_output, ) + return outputs + + +# For transformers 4.3.0 def my_longformer_self_attention_forward_4_3(self, hidden_states, attention_mask=None, @@ -117,190 +165,135 @@ def my_longformer_self_attention_forward_4_3(self, is_index_global_attn=None, is_global_attn=None, output_attentions=False): - # TODO: move mask calculation to LongFormerModel class to avoid calculating it again and again in each layer. - global_mask = is_index_global_attn.int() - torch.masked_fill(attention_mask, is_index_global_attn, 0.0) - - weight = torch.stack( - (self.query.weight.transpose(0, 1), self.key.weight.transpose(0, 1), self.value.weight.transpose(0, 1)), dim=1) - weight = weight.reshape(self.embed_dim, 3 * self.embed_dim) - - bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0) - bias = bias.reshape(3 * self.embed_dim) - - global_weight = torch.stack((self.query_global.weight.transpose(0, 1), self.key_global.weight.transpose( - 0, 1), self.value_global.weight.transpose(0, 1)), - dim=1) - global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim) - - global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0) - global_bias = global_bias.reshape(3 * self.embed_dim) - - attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, attention_mask, global_weight, - global_bias, global_mask, self.num_heads, - self.one_sided_attn_window_size) - - assert attn_output.size() == hidden_states.size(), "Unexpected size" - - outputs = (attn_output, ) - return outputs + assert output_attentions == False + return my_longformer_self_attention_forward_4(self, hidden_states, attention_mask, is_index_masked, + is_index_global_attn, is_global_attn) -#For transformers 4.0 -def my_longformer_self_attention_forward_4(self, - hidden_states, - attention_mask=None, - is_index_masked=None, - is_index_global_attn=None, - is_global_attn=None): - # TODO: move mask calculation to LongFormerModel class to avoid calculating it again and again in each layer. - global_mask = is_index_global_attn.int() - torch.masked_fill(attention_mask, is_index_global_attn, 0.0) - - weight = torch.stack( - (self.query.weight.transpose(0, 1), self.key.weight.transpose(0, 1), self.value.weight.transpose(0, 1)), dim=1) - weight = weight.reshape(self.embed_dim, 3 * self.embed_dim) - - bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0) - bias = bias.reshape(3 * self.embed_dim) - - global_weight = torch.stack((self.query_global.weight.transpose(0, 1), self.key_global.weight.transpose( - 0, 1), self.value_global.weight.transpose(0, 1)), - dim=1) - global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim) - - global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0) - global_bias = global_bias.reshape(3 * self.embed_dim) - - attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, attention_mask, global_weight, - global_bias, global_mask, self.num_heads, - self.one_sided_attn_window_size) - - assert attn_output.size() == hidden_states.size(), "Unexpected size" - - outputs = (attn_output, ) - return outputs +# For transformers 4.3.2 +def my_longformer_self_attention_forward_4_3_2(self, + hidden_states, + attention_mask=None, + layer_head_mask=None, + is_index_masked=None, + is_index_global_attn=None, + is_global_attn=None, + output_attentions=False): + assert output_attentions == False + assert layer_head_mask is None + return my_longformer_self_attention_forward_4(self, hidden_states, attention_mask, is_index_masked, + is_index_global_attn, is_global_attn) -# For transformers 3.0 -def my_longformer_attention_forward_3(self, hidden_states, attention_mask, output_attentions=False): +def export_longformer(model, onnx_model_path, export_padding): + input_ids, attention_mask, global_attention_mask = get_dummy_inputs(model.config, + export_padding, + device=torch.device('cpu')) - assert output_attentions is False + example_outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask) - # TODO: attention_mask can directly be passed from inputs of model to avoid these processing. - attention_mask = attention_mask.squeeze(dim=2).squeeze(dim=1) - global_mask = (attention_mask > 0).int() - torch.masked_fill(attention_mask, attention_mask > 0, 0.0) + if version.parse(transformers.__version__) < version.parse("4.0.0"): + raise RuntimeError("This tool requires transformers 4.0.0 or later.") - weight = torch.stack( - (self.query.weight.transpose(0, 1), self.key.weight.transpose(0, 1), self.value.weight.transpose(0, 1)), dim=1) - weight = weight.reshape(self.embed_dim, 3 * self.embed_dim) + # Here we replace LongformerSelfAttention.forward using our implmentation for exporting ONNX model + from transformers import LongformerSelfAttention + import inspect + key = ' '.join(inspect.getfullargspec(LongformerSelfAttention.forward).args) + args_to_func = { + 'self hidden_states attention_mask layer_head_mask is_index_masked is_index_global_attn is_global_attn output_attentions': + my_longformer_self_attention_forward_4_3_2, + 'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn output_attentions': + my_longformer_self_attention_forward_4_3, + 'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn': + my_longformer_self_attention_forward_4, + } - bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0) - bias = bias.reshape(3 * self.embed_dim) + if key not in args_to_func: + print("Current arguments", inspect.getfullargspec(LongformerSelfAttention.forward).args) + raise RuntimeError( + "LongformerSelfAttention.forward arguments are different. Please install supported version (like transformers 4.3.0)." + ) - global_weight = torch.stack((self.query_global.weight.transpose(0, 1), self.key_global.weight.transpose( - 0, 1), self.value_global.weight.transpose(0, 1)), - dim=1) - global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim) + # Store for restoring later + original_forward = LongformerSelfAttention.forward - global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0) - global_bias = global_bias.reshape(3 * self.embed_dim) + LongformerSelfAttention.forward = args_to_func[key] - attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, attention_mask, global_weight, - global_bias, global_mask, self.num_heads, - self.one_sided_attn_window_size) + example_inputs = (input_ids, attention_mask, global_attention_mask) - assert attn_output.size() == hidden_states.size(), "Unexpected size" + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) - outputs = (attn_output, ) - return outputs - - -# Here we replace LongformerSelfAttention.forward using our implmentation for exporting ONNX model -from transformers.modeling_longformer import LongformerSelfAttention -key = ' '.join(inspect.getfullargspec(LongformerSelfAttention.forward).args) -args_to_func = { - 'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn output_attentions': - my_longformer_self_attention_forward_4_3, - 'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn': - my_longformer_self_attention_forward_4, - 'self hidden_states attention_mask output_attentions': my_longformer_self_attention_forward_3, -} - -if key not in args_to_func: - raise RuntimeError( - "LongformerSelfAttention.forward arguments are different. Please install supported version (like 4.3.0) of transformers package." - ) - -LongformerSelfAttention.forward = args_to_func[key] -""" -if version.parse(transformers.__version__) < version.parse("4.0.0"): - from transformers.modeling_longformer import LongformerSelfAttention - #original_forward = LongformerSelfAttention.forward - LongformerSelfAttention.forward = my_longformer_attention_forward_3 -else: - from transformers.models.longformer.modeling_longformer import LongformerSelfAttention - #original_forward = LongformerSelfAttention.forward - LongformerSelfAttention.forward = my_longformer_self_attention_forward_4 -""" - -# TODO: support more inputs like (input_ids, attention_mask, global_attention_mask, token_type_ids, position_ids) -example_inputs = (input_ids, attention_mask, global_attention_mask) - -torch.onnx.export(model, - example_inputs, - onnx_model_path, - opset_version=11, - example_outputs=example_outputs, - input_names=["input_ids", "attention_mask", "global_attention_mask"], - output_names=["last_state", "pooler"], - dynamic_axes={ - 'input_ids': { - 0: 'batch_size', - 1: 'sequence_length' + torch.onnx.export(model, + example_inputs, + onnx_model_path, + opset_version=11, + example_outputs=example_outputs, + input_names=["input_ids", "attention_mask", "global_attention_mask"], + output_names=["last_state", "pooler"], + dynamic_axes={ + 'input_ids': { + 0: 'batch_size', + 1: 'sequence_length' + }, + 'attention_mask': { + 0: 'batch_size', + 1: 'sequence_length' + }, + 'global_attention_mask': { + 0: 'batch_size', + 1: 'sequence_length' + }, + 'last_state': { + 0: 'batch_size', + 1: 'sequence_length' + }, + 'pooler': { + 0: 'batch_size', + 1: 'sequence_length' + } }, - 'attention_mask': { - 0: 'batch_size', - 1: 'sequence_length' - }, - 'global_attention_mask': { - 0: 'batch_size', - 1: 'sequence_length' - }, - 'last_state': { - 0: 'batch_size', - 1: 'sequence_length' - }, - 'pooler': { - 0: 'batch_size', - 1: 'sequence_length' - } - }, - custom_opsets={"com.microsoft": 1}) -print(f"ONNX model exported to {onnx_model_path}") + custom_opsets={"com.microsoft": 1}) + print(f"ONNX model exported to {onnx_model_path}") -if args.sequence_length % model.config.attention_window[0] == 0: - print( - f"*Attention*: You need input padding for inference: input sequece length shall be multiple of {model.config.attention_window[0]}. It is because the example input for export ONNX model does not need padding so padding logic is not in onnx model." - ) + # Restore original implementaiton: + LongformerSelfAttention.forward = original_forward -# Restore Huggingface implementaiton like the following: -# LongformerSelfAttention.forward = original_forward -if args.precision != 'fp32' or args.optimize_onnx: +def optimize_longformer(onnx_model_path, fp32_model_path, fp16_model_path=None): from onnx import load_model from onnxruntime.transformers.onnx_model_bert import BertOnnxModel, BertOptimizationOptions model = load_model(onnx_model_path, format=None, load_external_data=True) optimization_options = BertOptimizationOptions('bert') - optimizer = BertOnnxModel(model, num_heads=16, hidden_size=768) + optimizer = BertOnnxModel(model, num_heads=16, + hidden_size=768) # paramters does not matter since attention fusion is not needed. optimizer.optimize(optimization_options) - optimized_model_path = model_name + "_fp32.onnx" - optimizer.save_model_to_file(optimized_model_path) - print(f"optimized fp32 model saved to {optimized_model_path}") - if args.precision == 'fp16': + use_external_data_format = False + if fp32_model_path: + optimizer.save_model_to_file(fp32_model_path, use_external_data_format) + print(f"optimized fp32 model saved to {fp32_model_path}") + + if fp16_model_path: optimizer.convert_model_float32_to_float16(cast_input_output=True) - optimized_model_path = model_name + "_fp16.onnx" - optimizer.save_model_to_file(optimized_model_path) - print(f"optimized fp16 model saved to {optimized_model_path}") + optimizer.save_model_to_file(fp16_model_path, use_external_data_format) + print(f"optimized fp16 model saved to {fp16_model_path}") + + +def main(args): + model_name = args.model + onnx_model_path = model_name + ".onnx" + + from transformers import LongformerModel + model = LongformerModel.from_pretrained(PRETRAINED_LONGFORMER_MODELS[model_name]) + + export_longformer(model, onnx_model_path, args.export_padding) + + if args.optimize_onnx or args.precision != 'fp32': + fp32_model_path = model_name + "_fp32.onnx" + fp16_model_path = model_name + "_fp16.onnx" if args.precision == 'fp16' else None + optimize_longformer(onnx_model_path, fp32_model_path, fp16_model_path) + + +if __name__ == "__main__": + args = parse_arguments() + main(args) diff --git a/onnxruntime/python/tools/transformers/longformer/longformer_helper.py b/onnxruntime/python/tools/transformers/longformer/longformer_helper.py new file mode 100644 index 0000000000..7a7ea17f54 --- /dev/null +++ b/onnxruntime/python/tools/transformers/longformer/longformer_helper.py @@ -0,0 +1,76 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +# This script helps creating dummy inputs for Longformer model. + +import os +import logging +import torch +import onnx +import random +import numpy +import time +import re +from pathlib import Path +from typing import List, Dict, Tuple, Union + +logger = logging.getLogger(__name__) + +PRETRAINED_LONGFORMER_MODELS = { + "longformer-base-4096": "allenai/longformer-base-4096", + "longformer-large-4096": "allenai/longformer-large-4096", + "longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging +} + + +class LongformerInputs: + def __init__(self, input_ids, attention_mask, global_attention_mask): + self.input_ids: torch.LongTensor = input_ids + self.attention_mask: Union[torch.FloatTensor, torch.HalfTensor] = attention_mask + self.global_attention_mask: Union[torch.FloatTensor, torch.HalfTensor] = global_attention_mask + + def to_list(self) -> List: + return [v for v in [self.input_ids, self.attention_mask, self.global_attention_mask] if v is not None] + + def to_tuple(self) -> Tuple: + return tuple(v for v in self.to_list()) + + def get_ort_inputs(self) -> Dict: + return { + "input_ids": numpy.ascontiguousarray(self.input_ids.cpu().numpy()), + "attention_mask": numpy.ascontiguousarray(self.attention_mask.cpu().numpy()), + "global_attention_mask": numpy.ascontiguousarray(self.global_attention_mask.cpu().numpy()), + } + + +class LongformerHelper: + """ A helper class for Longformer model conversion, inference and verification. + """ + @staticmethod + def get_dummy_inputs(batch_size: int, + sequence_length: int, + num_global_tokens: int, + device: torch.device, + vocab_size: int = 100) -> LongformerInputs: + """ Create random inputs for Longformer model. + Returns torch tensors of input_ids, attention_mask and global_attention_mask tensors. + """ + + input_ids = torch.randint(low=0, + high=vocab_size - 1, + size=(batch_size, sequence_length), + dtype=torch.long, + device=device) + attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device) + global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device) + global_token_index = list(range(num_global_tokens)) + global_attention_mask[:, global_token_index] = 1 + return LongformerInputs(input_ids, attention_mask, global_attention_mask) + + @staticmethod + def get_output_shapes(batch_size: int, sequence_length: int, hidden_size: int) -> Dict[str, List[int]]: + """ Returns a dictionary with output name as key, and shape as value. + """ + return {"last_state": [batch_size, sequence_length, hidden_size], "pooler": [batch_size, sequence_length]} diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index 02a953e713..5f821feea1 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -405,7 +405,10 @@ class OnnxModel: self.model.opset_import[0].version = original_opset_version def convert_model_float32_to_float16(self, cast_input_output=True): - """ Convert a graph to FLOAT16 + """Convert a graph to FLOAT16. By default, we will keep data types of inputs and outputs. + For decoder model with past_key_values, it is recommended to set cast_input_output=False for better performance. + Args: + cast_input_output (bool, optional): keep data type of inputs and outputs, and add Cast nodes to convert float32 inputs to float16, and float16 to float32 for outputs. Defaults to True. """ from packaging.version import Version import onnxconverter_common as oc diff --git a/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc b/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc index 8ad0d2f44c..82f97eede0 100644 --- a/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/longformer_attention_op_test.cc @@ -5,6 +5,8 @@ #include "test/common/tensor_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" +#include "test/util/include/scoped_env_vars.h" +#include "contrib_ops/cpu/bert/longformer_attention_base.h" namespace onnxruntime { namespace test { @@ -148,16 +150,33 @@ static void GetTinyLongformerData( static void RunTinyLongformerBatch1( std::vector& mask_data, std::vector& global_data, + std::vector& input_data, std::vector& output_data, - bool use_float16, - bool window_cover_whole_sequence = false) { + bool use_float16) { int batch_size = 1; int one_sided_attention_window_size = 2; int hidden_size = 8; int number_of_heads = 2; + std::vector weight_data; + std::vector bias_data; + std::vector global_weight_data; + std::vector global_bias_data; + GetTinyLongformerData(weight_data, bias_data, global_weight_data, global_bias_data); + + int sequence_length = static_cast(mask_data.size()) / batch_size; + + RunAttentionTest(input_data, weight_data, bias_data, mask_data, global_weight_data, global_bias_data, global_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, one_sided_attention_window_size, use_float16); +} + +static void RunTinyLongformerBatch1( + std::vector& mask_data, + std::vector& global_data, + std::vector& output_data, + bool use_float16, + bool window_cover_whole_sequence = false) { // Total windows size 4 will cover the whole sequence length 4 - int sequence_length = window_cover_whole_sequence ? 4 : 8; std::vector input_data; if (window_cover_whole_sequence) { input_data = { @@ -176,15 +195,7 @@ static void RunTinyLongformerBatch1( -1.0536f, -0.0425f, -1.1194f, -0.6423f, 2.1825f, 0.2547f, 0.6015f, -0.1809f, 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f}; } - - std::vector weight_data; - std::vector bias_data; - std::vector global_weight_data; - std::vector global_bias_data; - GetTinyLongformerData(weight_data, bias_data, global_weight_data, global_bias_data); - - RunAttentionTest(input_data, weight_data, bias_data, mask_data, global_weight_data, global_bias_data, global_data, output_data, - batch_size, sequence_length, hidden_size, number_of_heads, one_sided_attention_window_size, use_float16); + return RunTinyLongformerBatch1(mask_data, global_data, input_data, output_data, use_float16); } TEST(LongformerAttentionTest, LongformerAttention_NoGlobal) { @@ -238,6 +249,50 @@ TEST(LongformerAttentionTest, LongformerAttention_GlobalStart) { RunTinyLongformerBatch1(mask_data, global_data, output_data, false); } +/* +* This following case is generated from the first self-attention of a tiny longformer model like the following: + import torch + from transformers import AutoModel + model = AutoModel.from_pretrained("patrickvonplaten/longformer-random-tiny") + input_ids = torch.LongTensor([[97, 24, 71, 5, 8]])) + attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device) + global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=input_ids.device) + global_attention_mask[:, [0, 1,]] = 1 + outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask) +*/ +TEST(LongformerAttentionTest, LongformerAttention_UseCompactMemory) { + std::vector mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f, -10000.0f, -10000.0f}; + + std::vector global_data = {1, 1, 0, 0, 0, 0, 0, 0}; + + std::vector input_data = { + -1.6886f, -0.0712f, 0.7017f, -1.5513f, 0.1059f, 0.8119f, 1.1139f, 0.5776f, + -1.0874f, 0.2870f, -0.3767f, -1.8203f, 0.0116f, 1.4285f, 0.6643f, 0.8931f, + -1.0425f, 1.0148f, -0.6387f, -1.6496f, -0.3027f, 0.3839f, 1.0230f, 1.2119f, + 0.3576f, -1.5848f, 0.1713f, -0.9284f, 1.7172f, -0.8574f, 0.8390f, 0.2855f, + -1.7443f, 0.4251f, 1.6560f, -0.5054f, -0.8734f, 0.4769f, 0.8434f, -0.2781f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f, + 0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f}; + + std::vector output_data = { + -0.0330f, 0.0517f, -0.0292f, 0.0194f, -0.0154f, 0.0012f, -0.0041f, -0.0518f, + -0.0331f, 0.0517f, -0.0291f, 0.0194f, -0.0154f, 0.0011f, -0.0041f, -0.0517f, + 0.0421f, 0.0627f, 0.0445f, 0.0328f, 0.0196f, -0.0712f, -0.0233f, -0.0875f, + 0.0419f, 0.0626f, 0.0447f, 0.0328f, 0.0196f, -0.0712f, -0.0233f, -0.0875f, + 0.0421f, 0.0627f, 0.0445f, 0.0328f, 0.0196f, -0.0712f, -0.0233f, -0.0875f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, + 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f}; + + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::longformer::kUseCompactMemory, "1"}, + }}; + + RunTinyLongformerBatch1(mask_data, global_data, input_data, output_data, false); +} + TEST(LongformerAttentionTest, LongformerAttention_Float16) { std::vector mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f}; diff --git a/tools/ci_build/amd_hipify.py b/tools/ci_build/amd_hipify.py index 258297a9e5..4816783f25 100644 --- a/tools/ci_build/amd_hipify.py +++ b/tools/ci_build/amd_hipify.py @@ -26,6 +26,8 @@ contrib_ops_excluded_files = [ 'bert/layer_norm.cuh', 'bert/longformer_attention.cc', 'bert/longformer_attention.h', + 'bert/longformer_attention_softmax.cu', + 'bert/longformer_attention_softmax.h', 'bert/longformer_attention_impl.cu', 'bert/longformer_attention_impl.h', 'bert/longformer_global_impl.cu',