From 9b446d5f7e51780619c95b12e12e8ec445948628 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 16 Feb 2021 14:54:48 -0800 Subject: [PATCH] Longformer Attention CUDA kernel memory Improvements (#6646) * Integrate memory improvements from NVidia * compute max_global_num before buffer allocation * update conversion script to support transformers 4.0 * update benchmark script for creating dummy inputs for different batch_size * Use a wrapper of cuda event to avoid memory leak --- .../cuda/bert/longformer_attention.cc | 82 +- .../cuda/bert/longformer_attention_impl.cu | 1002 +++++++++-------- .../cuda/bert/longformer_attention_impl.h | 17 +- .../cuda/bert/longformer_global_impl.cu | 76 ++ .../cuda/bert/longformer_global_impl.h | 26 + .../tools/transformers/longformer/__init__.py | 4 + .../longformer/benchmark_longformer.py | 63 +- .../longformer/convert_longformer_to_onnx.py | 59 +- 8 files changed, 834 insertions(+), 495 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.h create mode 100644 onnxruntime/python/tools/transformers/longformer/__init__.py diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc b/onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc index 9ec5298c2b..89156634a8 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 "longformer_global_impl.h" #include "longformer_attention_impl.h" using namespace onnxruntime::cuda; @@ -29,6 +30,24 @@ namespace cuda { REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(MLFloat16) +// A wrapper class of cudaEvent_t to destroy the event automatically for avoiding memory leak. +class AutoDestoryCudaEvent { + public: + AutoDestoryCudaEvent() : cuda_event_(nullptr) { + } + + ~AutoDestoryCudaEvent() { + if (cuda_event_ != nullptr) + cudaEventDestroy(cuda_event_); + } + + cudaEvent_t& Get() { + return cuda_event_; + } + private: + cudaEvent_t cuda_event_; +}; + template LongformerAttention::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {} @@ -56,8 +75,41 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { Tensor* output = context->Output(0, shape); cublasHandle_t cublas = CublasHandle(); + cudaStream_t stream = Stream(); + CUBLAS_RETURN_IF_ERROR(cublasSetStream(cublas, stream)); + constexpr size_t element_size = sizeof(T); + // Build Global Index + auto global_index_buffer = GetScratchBuffer(batch_size * sequence_length); + auto batch_global_num_buffer = GetScratchBuffer(batch_size); + + size_t global_scratch_bytes = GetGlobalScratchSize(batch_size, sequence_length); + auto global_scratch_buffer = GetScratchBuffer(global_scratch_bytes); + + BuildGlobalIndex( + stream, + global_attention->template Data(), + batch_size, + sequence_length, + global_index_buffer.get(), + batch_global_num_buffer.get(), + global_scratch_buffer.get(), + global_scratch_bytes); + + // Copy batch_global_num to CPU + size_t pinned_buffer_bytes = GetPinnedBufferSize(batch_size); + auto pinned_buffer = AllocateBufferOnCPUPinned(pinned_buffer_bytes); + int* batch_global_num_pinned = reinterpret_cast(pinned_buffer.get()); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(batch_global_num_pinned, batch_global_num_buffer.get(), batch_size * sizeof(int), cudaMemcpyDeviceToHost, stream)); + + // Create an event to make sure the async copy is finished before reading the data. + AutoDestoryCudaEvent new_event; + cudaEvent_t& isCopyDone = new_event.Get(); + + CUDA_RETURN_IF_ERROR(cudaEventCreate(&isCopyDone)); + CUDA_RETURN_IF_ERROR(cudaEventRecord(isCopyDone, stream)); + // Use GEMM for fully connection. int m = batch_size * sequence_length; int n = 3 * hidden_size; @@ -85,8 +137,21 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { reinterpret_cast(input->template Data()), k, &one, reinterpret_cast(gemm_buffer.get()), n, device_prop)); - // TODO: calculate the exact value from global flags. - int max_num_global = sequence_length; + // Wait for async copy of batch_global_num + CUDA_RETURN_IF_ERROR(cudaEventSynchronize(isCopyDone)); + + // Find the maximum number of global tokens in all batches + int max_num_global = 0; + for (int i = 0; i < batch_size; ++i) { + if (max_num_global < batch_global_num_pinned[i]) { + max_num_global = batch_global_num_pinned[i]; + } + } + + // 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."); + } // Fully connection for global projection. // Note that Q only need handle global query tokens if we split GEMM to global Q/K/V separately. @@ -107,15 +172,20 @@ 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); + size_t workSpaceSize = GetLongformerAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, max_num_global, window_); auto workspace_buffer = GetScratchBuffer(workSpaceSize); if (!LaunchLongformerAttentionKernel( device_prop, - Stream(), + cublas, + stream, reinterpret_cast(gemm_buffer.get()), reinterpret_cast(mask->template Data()), reinterpret_cast(global_gemm_buffer.get()), global_attention->template Data(), + global_index_buffer.get(), + batch_global_num_buffer.get(), + pinned_buffer.get(), + workspace_buffer.get(), output->template MutableData(), batch_size, sequence_length, @@ -123,14 +193,14 @@ Status LongformerAttention::ComputeInternal(OpKernelContext* context) const { head_size, window_, max_num_global, - workspace_buffer.get(), - cublas, element_size)) { // Get last error to reset it to cudaSuccess. CUDA_CALL(cudaGetLastError()); return Status(common::ONNXRUNTIME, common::FAIL); } + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); + this->AddDeferredReleaseCPUPtr(pinned_buffer.release()); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu index 191a979fc9..d43b524997 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu @@ -16,7 +16,7 @@ limitations under the License. // Limitations of current Longformer Attention CUDA Kernels: // (1) Does not support global tokens in the middle. All global tokens shall be in the beginning of sequence. -// (2) Batch size <= 128 (defined in MAX_LONGFORMER_BATCH_SIZE) +// (2) Maximum number of global tokens <= one-sided attention window #include #include @@ -33,12 +33,15 @@ limitations under the License. using namespace onnxruntime::cuda; using namespace cub; -#define CHECK(status) \ - if (!CUBLAS_CALL(status)) { \ - return false; \ +#define CHECK(expr) \ + if (!CUBLAS_CALL(expr)) { \ + return false; \ } -constexpr int MAX_LONGFORMER_BATCH_SIZE = 128; +#define CHECK_CUDA(expr) \ + if (!CUDA_CALL(expr)) { \ + return false; \ + } namespace onnxruntime { namespace contrib { @@ -46,24 +49,35 @@ namespace cuda { // Denote: batch size (B), sequence length (S), number of heads (N), dimension per head (H), max number of global tokens (G) // -// Workspace layout (by default, the data type T is float or half): -// [SoftmaxSpace: see below] [Q:BxNxSxH] [K:BxNxSxH] [V:BxNxSxH] [Global_Q:BxNxGxH] [Global_K:BxNxSxH] [Global_V:BxNxSxH] -// where Global_Q, Global_K and Global_V are optional. They are not allocated when there is no any global token. +// Workspace layout (default data type T is float or half): +// [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. // -// SoftmaxSpace layout (tmp_storage could use the space of scratch1, scratch2, Q and K): -// [Global_Idx: int BxS][batch_global_num: int BxS][sequence_index: int BxS][tmp_storage: int 1024x1] -// [scratch1: BxNxSxS ] [scratch2: BxNxSxS ] -// Allocated size could be slightly larger than needed: batch_global_num uses only Bx1 and allocated BxS. -// Scratch size is allocated as multiples of 256. +// 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. + +size_t GetScratch1Size(size_t element_size, int batch_size, int num_heads, int sequence_length, int window) { + return (5 * sequence_length - 3 * window) * window * num_heads * batch_size * element_size; +} + +constexpr size_t GetScratch2Size() { + return 10 * (sizeof(void*) + sizeof(size_t)); +} size_t GetLongformerSoftmaxWorkspaceSize( size_t element_size, int batch_size, int num_heads, - int sequence_length) { - size_t temp_size = sizeof(int) * 1024; - size_t scratch_size = 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length); - return 3 * batch_size * sequence_length * sizeof(int) + std::max(scratch_size, temp_size); + 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; } size_t GetLongformerAttentionWorkspaceSize( @@ -72,187 +86,268 @@ size_t GetLongformerAttentionWorkspaceSize( int num_heads, int head_size, int sequence_length, - int max_num_global) { - size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length); + int max_num_global, + int window) { + size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window); 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; } -__global__ void InitSequenceIndexKernel(int* sequence_index, int sequence_length) { - int batch_index = blockIdx.x; - for (int i = threadIdx.x; i < sequence_length; i += blockDim.x) { - sequence_index[batch_index * sequence_length + i] = i; - } -} - -// TODO: Move this to its own plugin that can be run once for all layers. -int* BuildGlobalIndex(cudaStream_t stream, const int* global_attention, int batch_size, int sequence_length, void* workspace, size_t softmax_workspace_size) { - int* global_idx = reinterpret_cast(workspace); - int* batch_global_num = global_idx + batch_size * sequence_length; // Number of global tokens in each batch, shape is (batch_size) - int* sequence_index = batch_global_num + batch_size * sequence_length; - int* tmp_storage = sequence_index + batch_size * sequence_length; - - InitSequenceIndexKernel<<>>(sequence_index, sequence_length); - - // Determine temporary device storage requirements - // Find the global attention indices and number of global attention tokens - size_t temp_storage_bytes = 0; - cub::DevicePartition::Flagged(NULL, temp_storage_bytes, sequence_index, - global_attention, global_idx, batch_global_num, sequence_length, stream); - assert(temp_storage_bytes <= softmax_workspace_size - static_cast(3 * batch_size * sequence_length)); - - for (int i = 0; i < batch_size; ++i) { - cub::DevicePartition::Flagged(reinterpret_cast(tmp_storage), temp_storage_bytes, sequence_index, - global_attention + i * sequence_length, global_idx + i * sequence_length, - batch_global_num + i, sequence_length, stream); - } - - return global_idx; +// Size of buffer of pinned memory in CPU. The buffer is used to copy memory between CPU and GPU. +// The buffer includes two parts: [global_count (copy of batch_global_num): int Bx1] [copy of scratch2] +size_t GetPinnedBufferSize(int batch_size) { + return sizeof(int) * batch_size + GetScratch2Size(); } template __launch_bounds__(blockSize) __global__ void LongformerSoftmaxKernel(const int* global_attention, - const int* global_idx, + const int* global_index, const int* batch_global_num, - const T* input, + void* input_pointers, const T* attention_mask, - T* output, float scaler, int dim0, int sequence_length, - int attention_window) { + int window, + int num_heads) { 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; + int tid = threadIdx.x; const int batch_index = blockIdx.x / dim0; const int row_index = blockIdx.x % sequence_length; + const int head_index = (blockIdx.x / sequence_length) % num_heads; + + // Adjust the pointers for the batch + const T* mask_block = attention_mask + sequence_length * batch_index; + const int* global_index_block = global_index + sequence_length * batch_index; 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; + size_t* p_inputs = (size_t*)(input_pointers); + size_t* p_outputs = (size_t*)(input_pointers) + 5; + size_t* input_sizes = (size_t*)(input_pointers) + 10; + size_t* input_strides = (size_t*)(input_pointers) + 15; + + const T* inputs[5]; + T* outputs[5]; + for (int i = 0; i < 5; ++i) { + inputs[i] = (T*)p_inputs[i] + batch_index * num_heads * input_sizes[i]; + outputs[i] = (T*)p_outputs[i] + batch_index * num_heads * input_sizes[i]; } - // local attention token + // 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); + bool is_local_row = (global_attention[batch_index * sequence_length + row_index] == static_cast(0)); if (is_local_row) { - col_start = row_index - attention_window; + col_start = row_index - window; if (col_start < 0) { col_start = 0; } - col_end = row_index + attention_window + 1; + col_end = row_index + 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; + // If mask is set then set everything to zero to match huggingface transformers implementation + if ((float)mask_block[row_index] != 0.f) { + if (is_local_row) { + T* output_block = nullptr; + T* output_global = nullptr; + int local_offset = row_index % window; + int local_start = 0; + int local_end = 3 * window; + if (row_index < window) { + local_start = 0; + local_end = 2 * window; + output_block = outputs[0] + row_index * input_strides[0] + head_index * input_sizes[0]; + } else if (row_index < sequence_length - window) { + output_block = outputs[1] + (row_index - window) * input_strides[1] + head_index * input_sizes[1]; + } else { + local_start = 0; + local_end = 2 * window; + output_block = outputs[2] + local_offset * input_strides[2] + head_index * input_sizes[2]; + } - // 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; - } - } + for (int i = local_start + tid; i < local_end; i += blockSize) { + output_block[i] = 0; + } - if (is_local_row) { - for (int g = tid; g < global_num; g += blockSize) { - int i = global_idx[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; + if ((row_index - 2 * window) >= 0) { + output_global = outputs[3] + (row_index - window) * input_strides[3] + head_index * input_sizes[3]; + } + + if (output_global != nullptr) { + for (int i = tid; i < global_num; i += blockSize) { + output_global[i] = 0; } } - } - } - //__syncthreads(); - 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; + } else { + T* output_block = outputs[4]; + for (int i = tid; i < sequence_length; i += blockSize) + output_block[i] = 0; + } + return; } + float sum_input = 0.; + + // Calculate max input + float max_input = -CUDART_INF_F; + if (is_local_row) { - for (int g = tid; g < global_num; g += blockSize) { - int i = global_idx[g]; - if (i < col_start || i > col_end) { - float x = input_block[i]; - x = expf((x)*scaler + (float)mask_block[i] - max_shared); + const T* input_block = nullptr; + T* output_block = nullptr; + T* output_global = nullptr; + int local_offset = row_index % window; + int local_start = local_offset; + int local_end = local_start + 2 * window + 1; + int zero_start = 0; + int zero_end = 3 * window; + if (row_index < window) { + local_start = 0; + local_end = local_offset + window + 1; + zero_end = 2 * window; + + input_block = inputs[0] + row_index * input_strides[0] + head_index * input_sizes[0]; + output_block = outputs[0] + row_index * input_strides[0] + head_index * input_sizes[0]; + } else if (row_index < sequence_length - window) { + input_block = inputs[1] + (row_index - window) * input_strides[1] + head_index * input_sizes[1]; + output_block = outputs[1] + (row_index - window) * input_strides[1] + head_index * input_sizes[1]; + } else { + local_start = local_offset; + local_end = 2 * window; + zero_end = 2 * window; + + input_block = inputs[2] + local_offset * input_strides[2] + head_index * input_sizes[2]; + output_block = outputs[2] + local_offset * input_strides[2] + head_index * input_sizes[2]; + } + + const T* input_global = nullptr; + int local_global = row_index - window; + if (local_global > global_num) local_global = global_num; + if (local_global > 0) { + input_global = inputs[3] + (row_index - window) * input_strides[3] + head_index * input_sizes[3]; + } + + if (row_index < window) { + output_global = (T*)outputs[0] + row_index * input_strides[0] + head_index * input_sizes[0]; + } else if (row_index < 2 * window) { + output_global = outputs[1] + (row_index - window) * input_strides[1] + head_index * input_sizes[1]; + } else { + output_global = outputs[3] + (row_index - window) * input_strides[3] + head_index * input_sizes[3]; + } + + for (int i = local_start + tid, j = col_start + tid; i < local_end; i += blockSize, j += blockSize) { + float x = input_block[i]; + x = x * scaler + (float)mask_block[j]; + if (max_input < x) + max_input = x; + } + + if (input_global != nullptr) { + for (int i = tid; i < local_global; i += blockSize) { + float x = input_global[global_index_block[i]]; + x = x * scaler + (float)mask_block[global_index_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(); + + for (int i = local_start + tid, j = col_start + tid; i < local_end; i += blockSize, j += blockSize) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[j] - max_shared); + sum_input += x; + } + + if (input_global != nullptr) { + for (int i = tid, j = col_start + tid; i < local_global; i += blockSize, j += blockSize) { + float x = input_global[global_index_block[i]]; + x = expf((x)*scaler + (float)mask_block[j] - max_shared); sum_input += x; } } - } - //__syncthreads(); - 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; + 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; - 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) { + for (int i = tid + zero_start; i < local_start; i += blockSize) { output_block[i] = (T)(0.); } - for (int g = tid; g < global_num; g += blockSize) { - int i = global_idx[g]; + for (int i = tid + local_end; i < zero_end; i += blockSize) { + output_block[i] = (T)(0.); + } + + __syncthreads(); + + for (int i = local_start + tid, j = col_start + tid; i < local_end; i += blockSize, j += blockSize) { + float x = input_block[i]; + x = expf((x)*scaler + (float)mask_block[j] - max_shared); + output_block[i] = (T)(recip_sum * x); + } + + if (input_global != nullptr) { + for (int i = tid; i < local_global; i += blockSize) { + float x = input_global[global_index_block[i]]; + x = expf((x)*scaler + (float)mask_block[global_index_block[i]] - max_shared); + output_global[i] = (T)(recip_sum * x); + } + } + } else { + // Global tokens + const T* input_block = inputs[4] + row_index * input_strides[4] + head_index * input_sizes[4]; + T* output_block = outputs[4] + row_index * input_strides[4] + head_index * input_sizes[4]; + + for (int i = tid; i < sequence_length; i += blockSize) { + 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(); + + for (int i = tid; i < sequence_length; i += blockSize) { + 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; + + for (int i = tid; i < sequence_length; i += blockSize) { 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); - } } bool launchSoftmaxKernel( @@ -264,37 +359,31 @@ bool launchSoftmaxKernel( 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 not masked and -10000 masked. - const void* global_q, // Q for global tokens with shape (B, N, G, H) + 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. It has two parts: Number of global tokens per batch with shape (B, 1), and a buffer to copy data to scratch2 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 + 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 - if (batch_size > MAX_LONGFORMER_BATCH_SIZE) { - ORT_THROW("LongformerAttention CUDA operator does not support batch size > 128."); - } + assert(max_num_global <= window); + + const int* global_count = reinterpret_cast(pinned_buffer); bool is_fp16 = (element_size == 2); - void* scratch1 = reinterpret_cast(workspace) + 3 * sizeof(int) * batch_size * sequence_length; - void* scratch2 = reinterpret_cast(scratch1) + GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length); + void* scratch1 = reinterpret_cast(workspace); + char* scratch2 = (char*)scratch1 + GetScratch1Size(element_size, batch_size, num_heads, sequence_length, window); - // Build index for global tokens - int* global_idx = BuildGlobalIndex(stream, global_attention, batch_size, sequence_length, workspace, softmax_workspace_size); - int* batch_global_num = global_idx + batch_size * sequence_length; - - int num_global[MAX_LONGFORMER_BATCH_SIZE] = {-1}; - if (!CUDA_CALL(cudaMemcpyAsync(&num_global[0], batch_global_num, batch_size * sizeof(int), cudaMemcpyDeviceToHost, stream))) { - return false; - } - - // setup shared parameters for two strided batched matrix multiplies + // Setup shared parameters for two strided batched matrix multiplies cudaDataType_t Atype; cudaDataType_t Btype; cudaDataType_t Ctype; @@ -332,6 +421,8 @@ bool launchSoftmaxKernel( // 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 + int elements_per_batch = num_heads * sequence_length * head_size; + int stride_per_head = sequence_length * head_size; // stride for Q, K, V and output // Local attention part // S x S is calculated using sliding block WxW (W is one sided window size) like the following: @@ -339,77 +430,94 @@ bool launchSoftmaxKernel( // [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 first and last rows have 2 blocks per row, 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. - // TODO: Save space by not storing the whole matrix. Instead only allocate space for these blocks. - - 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; + // To save space, we do not store the whole matrix. Instead, we only allocate space for these blocks. + // + // For global attention part, we have two assumptions: + // (1) Global tokens are at the beginging of sequence + // (2) Number of global tokens <= attention window + // + // The results are stored in scratch1 buffer: + // Number of elements for local attention are (3*S/W-2)*W*W*N*B, or (3S-2W)*W*N*B + // Number of elements for local attends to global are (S-W)*W*N*B + // Number of elements for global attends to everything are S*W*N*B + // Total elements (FP16 or FP32) are (5S-3W)*W*N*B - // 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)); + const int w = window; + const int middle_count = (sequence_length - 2 * w) / w; + int last_block = (sequence_length / w) - 1; + + // Determine the non-zero block dimensions and pointers + + // Buffer size per head for a single batch + size_t buffer_sizes[5] = { + static_cast(w * w * 2), // first row of blocks has 2 WxW blocks + static_cast(w * w * middle_count * 3), // middle rows of blocks have 3 WxW blocks per row + static_cast(w * w * 2), // last row of blocks has 2 WxW blocks + static_cast(w * (sequence_length - w)), // local attends to global: global tokens are assumed to be smaller than window size + static_cast(w * sequence_length)}; // global attends to everything. + + size_t buffer_strides[5] = { + static_cast(w * 2), + static_cast(w * 3), + static_cast(w * 2), + static_cast(w), // global tokens are assumed to be smaller than window size + static_cast(sequence_length)}; + + void* input_pointers[5]; + void* output_pointers[5]; + + char* current_pointer = (char*)scratch1; + for (int i = 0; i < 5; ++i) { + input_pointers[i] = (void*)current_pointer; + output_pointers[i] = (void*)current_pointer; // output pointer is same as input + current_pointer += buffer_sizes[i] * num_heads * batch_size * element_size; + } + assert(current_pointer == scratch2); + + // Copy to a continues buffer first so that we only need call cudaMemcpyAsync once + + constexpr size_t totalBytes = 10 * (sizeof(size_t) + sizeof(void*)); + char* temp_buffer = reinterpret_cast(pinned_buffer) + sizeof(int) * batch_size; + memcpy(temp_buffer, &input_pointers[0], 5 * sizeof(void*)); + memcpy(temp_buffer + 5 * sizeof(void*), &output_pointers[0], 5 * sizeof(void*)); + memcpy(temp_buffer + 10 * sizeof(void*), &buffer_sizes[0], 5 * sizeof(size_t)); + memcpy(temp_buffer + 10 * sizeof(void*) + 5 * sizeof(size_t), &buffer_strides[0], 5 * sizeof(size_t)); + CHECK_CUDA(cudaMemcpyAsync(scratch2, temp_buffer, totalBytes, cudaMemcpyHostToDevice, stream)); + + // Local attention part + { + if (middle_count > 0) { + for (int i = 0; i < batch_size; ++i) { + for (int j = 0; j < num_heads; ++j) { + void* q_head = (char*)q + (i * elements_per_batch + j * sequence_length * head_size + w * head_size) * element_size; + void* k_head = (char*)k + (i * elements_per_batch + j * sequence_length * head_size) * element_size; + void* qk_head = (char*)input_pointers[1] + (i * num_heads + j) * buffer_sizes[1] * element_size; + 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 + 3 * w, // ldc + 3 * w * w, // strideC + middle_count, // batch count + resultType, + algo)); + } } } @@ -423,106 +531,107 @@ bool launchSoftmaxKernel( k, // A Atype, // A type head_size, // lda - strideA, // strideA + stride_per_head, // strideA q, // B Btype, // B type head_size, // ldb - strideB, // strideB + stride_per_head, // strideB beta_0, // beta - scratch1, // C + input_pointers[0], // C Ctype, // C type - sequence_length, // ldc - strideC, // strideC + 2 * w, // ldc + buffer_sizes[0], // 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, + 2 * w, // m + w, // n + head_size, // k + alpha, // alpha + k_head, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + q_head, // B + Btype, // B type + head_size, // ldb + stride_per_head, // strideB + beta_0, // beta + input_pointers[2], // C + Ctype, // C type + 2 * w, // ldc + buffer_sizes[2], // strideC + batch_size * num_heads, // batch count resultType, algo)); } // Global attention part for (int i = 0; i < batch_size; ++i) { - if (num_global[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; + if (global_count[i] > 0) { + void* q_batch = (char*)q + (i * elements_per_batch + w * head_size) * element_size; + void* k_batch = (char*)k + (i * elements_per_batch) * element_size; + void* qk_batch = (char*)input_pointers[3] + (i * buffer_sizes[3]) * num_heads * element_size; + // Local tokens attending global tokens CHECK(cublasGemmStridedBatchedEx(cublas, CUBLAS_OP_T, CUBLAS_OP_N, - num_global[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, + global_count[i], // m + sequence_length - w, // n + head_size, // k + alpha, // alpha + k_batch, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + q_batch, // B + Btype, // B type + head_size, // ldb + stride_per_head, // strideB + beta_0, // beta + qk_batch, // C + Ctype, // C type + w, // ldc + buffer_sizes[3], // strideC + num_heads, // batch count resultType, algo)); - void* global_q_batch = (char*)global_q + (i * num_heads * max_num_global * head_size) * element_size; - void* global_k_batch = (char*)global_k + (i * x_offset) * element_size; - int strideB_global = max_num_global * head_size; + 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 + 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; // 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, - num_global[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, + sequence_length, // m + global_count[i], // n + head_size, // k + alpha, // alpha + global_k_batch, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + global_q_batch, // B + Btype, // B type + head_size, // ldb + stride_per_head, // strideB. For compact format: max_num_global * head_size. + beta_0, // beta + qk_batch, // C + Ctype, // C type + sequence_length, // ldc + buffer_sizes[4], // strideC + num_heads, // batch count resultType, algo)); } @@ -530,205 +639,177 @@ bool launchSoftmaxKernel( 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) { LongformerSoftmaxKernel<__half, blockSize><<>>( global_attention, - global_idx, + global_index, batch_global_num, - static_cast(scratch1), + scratch2, static_cast(attention_mask), - static_cast<__half*>(softmax_out), scaler, dim0, dim1, attention_window); + scaler, dim0, dim1, window, num_heads); } else { LongformerSoftmaxKernel<<>>( global_attention, - global_idx, + global_index, batch_global_num, - static_cast(scratch1), + scratch2, static_cast(attention_mask), - static_cast(softmax_out), scaler, dim0, dim1, attention_window); + scaler, dim0, dim1, window, num_heads); } // 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. + // Calculation uses sliding blocks 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)); + { + if (middle_count > 0) { + for (int i = 0; i < batch_size; ++i) { + for (int j = 0; j < num_heads; ++j) { + void* v_head = (char*)v + (i * elements_per_batch + j * head_size * sequence_length) * element_size; + void* prob_head = (char*)output_pointers[1] + (i * num_heads * buffer_sizes[1] + j * buffer_sizes[1]) * element_size; + void* out_head = (char*)output + (i * elements_per_batch + j * head_size * sequence_length + w * head_size) * element_size; + CHECK(cublasGemmStridedBatchedEx(cublas, + CUBLAS_OP_N, + CUBLAS_OP_N, + head_size, // m + w, // n + 3 * w, // k + alpha, // alpha + v_head, // A + Atype, // A type + head_size, // lda + w * head_size, // strideA + prob_head, // B + Btype, // B type + (int)buffer_strides[1], // ldb + 3 * w * w, // strideB + beta_0, // beta + out_head, // C + Ctype, // C type + head_size, // ldc + w * head_size, // strideC + middle_count, // batch 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, + head_size, // m + w, // n + 2 * w, // k + alpha, // alpha + v, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + output_pointers[0], // B + Btype, // B type + (int)buffer_strides[0], // ldb + buffer_sizes[0], // strideB + beta_0, // beta + output, // C + Ctype, // C type + head_size, // ldc + stride_per_head, // strideC + batch_size * num_heads, // batch count 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, + head_size, // m + w, // n + 2 * w, // k + alpha, // alpha + v_head, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + output_pointers[2], // B + Btype, // B type + (int)buffer_strides[2], // ldb + buffer_sizes[2], // strideB + beta_0, // beta + out_head, // C + Ctype, // C type + head_size, // ldc + stride_per_head, // strideC + batch_size * num_heads, // batch count resultType, algo)); } - for (int i = 0; i < batch_size; ++i) { - if (num_global[i] > 0) { - int glob_longdim_mm = (last_block - 1) * w; + if (global_count[i] > 0) { + int glob_longdim_mm = sequence_length - 2 * 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; + void* v_head = (char*)v + (i * elements_per_batch) * element_size; + void* prob_head = (char*)output_pointers[3] + (i * buffer_sizes[3] * num_heads + w * buffer_strides[3]) * element_size; + void* out_head = (char*)output + (i * elements_per_batch + 2 * w * head_size) * element_size; CHECK(cublasGemmStridedBatchedEx(cublas, CUBLAS_OP_N, CUBLAS_OP_N, - head_size, - glob_longdim_mm, - num_global[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, + head_size, // m + glob_longdim_mm, // n + global_count[i], // k + alpha, // alpha + v_head, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + prob_head, // B + Btype, // B type + (int)buffer_strides[3], // ldb + buffer_sizes[3], // strideB + beta_1, // beta + out_head, // C + Ctype, // C type + head_size, // ldc + stride_per_head, // strideC + num_heads, // batch count 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; + v_head = (char*)global_v + (i * elements_per_batch) * element_size; + prob_head = (char*)output_pointers[4] + (i * buffer_sizes[4] * num_heads) * element_size; + out_head = (char*)output + (i * elements_per_batch) * element_size; CHECK(cublasGemmStridedBatchedEx(cublas, CUBLAS_OP_N, CUBLAS_OP_N, - head_size, - num_global[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, + head_size, // m + global_count[i], // n + sequence_length, // k: re-write entries completely + alpha, // alpha + v_head, // A + Atype, // A type + head_size, // lda + stride_per_head, // strideA + prob_head, // B + Btype, // B type + (int)buffer_strides[4], // ldb + buffer_sizes[4], // strideB + beta_0, // beta: overwrite + out_head, // C: assumes global tokens are at the beginning of sequence + Ctype, // C type + head_size, // ldc + stride_per_head, // strideC + num_heads, // batch count resultType, algo)); } @@ -739,23 +820,27 @@ bool launchSoftmaxKernel( template bool LongformerQkvToContext( - const cudaDeviceProp& prop, cublasHandle_t& cublas, cudaStream_t stream, + cublasHandle_t& cublas, cudaStream_t stream, const int batch_size, const int sequence_length, const int num_heads, const int head_size, const int window, const size_t element_size, const T* input, const T* attention_mask, - const T* global_input, const int* global_attention, const int max_num_global, - T* workspace, + 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); + size_t softmax_workspace_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window); 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 + const int elements = batch_size * num_heads * sequence_length * head_size; + // Input should be BxSx3xNxH => qkv: 3xBxNxSxH if (!LaunchTransQkv(stream, sequence_length, batch_size, head_size, num_heads, input, qkv)) { return false; } // Input 'global_input' should be BxSx3xNxH => global_qkv: 3xBxNxSxH - T* global_qkv = qkv + 3 * batch_size * sequence_length * num_heads * head_size * element_size; + T* global_qkv = qkv + 3 * elements; // When there is no global token, no need to process global Q, K and V if (max_num_global > 0 && nullptr != global_input) { @@ -765,7 +850,6 @@ bool LongformerQkvToContext( } // Now qkv has Q, K, V: each has size BxNxSxH - const int elements = batch_size * num_heads * sequence_length * head_size; const T* q = qkv; const T* k = q + elements; const T* v = k + elements; @@ -774,14 +858,9 @@ bool LongformerQkvToContext( const T* global_k = global_q + elements; const T* global_v = global_k + elements; - cublasSetStream(cublas, stream); - CublasMathModeSetter helper(prop, cublas, CUBLAS_TENSOR_OP_MATH); - // Q*K' are scaled by 1/sqrt(H) const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); - cudaDeviceSynchronize(); - T* temp_output = qkv; // Q will be overwritten if (!launchSoftmaxKernel( stream, @@ -792,17 +871,20 @@ bool LongformerQkvToContext( 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 G x H + 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) windows size + window, // Half (one-sided) window size max_num_global, // Maximum number of global tokens (G) element_size)) { return false; @@ -813,12 +895,17 @@ bool LongformerQkvToContext( } bool LaunchLongformerAttentionKernel( - const cudaDeviceProp& prop, + const cudaDeviceProp& device_prop, + cublasHandle_t& cublas, cudaStream_t stream, const void* input, const void* attention_mask, const void* global_input, const int* global_attention, + const int* global_index, + const int* batch_global_num, + void* pinned_buffer, + void* workspace, void* output, int batch_size, int sequence_length, @@ -826,27 +913,32 @@ bool LaunchLongformerAttentionKernel( int head_size, int window, int max_num_global, - void* workspace, - cublasHandle_t& cublas, const size_t element_size) { + CublasMathModeSetter helper(device_prop, cublas, CUBLAS_TENSOR_OP_MATH); if (element_size == 2) { - return LongformerQkvToContext(prop, cublas, stream, + return LongformerQkvToContext(cublas, stream, batch_size, sequence_length, num_heads, head_size, window, element_size, reinterpret_cast(input), reinterpret_cast(attention_mask), reinterpret_cast(global_input), global_attention, + global_index, + batch_global_num, max_num_global, + pinned_buffer, reinterpret_cast(workspace), reinterpret_cast(output)); } else { - return LongformerQkvToContext(prop, cublas, stream, + return LongformerQkvToContext(cublas, stream, batch_size, sequence_length, num_heads, head_size, window, element_size, reinterpret_cast(input), reinterpret_cast(attention_mask), reinterpret_cast(global_input), global_attention, + global_index, + batch_global_num, max_num_global, + pinned_buffer, reinterpret_cast(workspace), reinterpret_cast(output)); } diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h index c08461e800..1cc0d80313 100644 --- a/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h @@ -8,21 +8,30 @@ namespace onnxruntime { namespace contrib { namespace cuda { + size_t GetPinnedBufferSize( + int batch_size); + size_t GetLongformerAttentionWorkspaceSize( size_t element_size, - int batchsize, + int batch_size, int num_heads, int head_size, int sequence_length, - int max_num_global); + int max_num_global, + int window); -bool LaunchLongformerAttentionKernel( + bool LaunchLongformerAttentionKernel( const cudaDeviceProp& device_prop, // Device Properties + cublasHandle_t& cublas, // Cublas handle cudaStream_t stream, // CUDA stream const void* input, // Input tensor const void* attention_mask, // Attention mask with shape (B, S) const void* global_input, // Global attention input, or nullptr when max_num_global == 0. const int* global_attention, // Global attention flags with shape (B, S) + const int* global_index, // Global index + const int* batch_global_num, // Number of global tokens per batch. It is in device memory. + void* pinned_buffer, // Buffer in pinned memory of CPU with two parts: a copy of batch_global_num, and buffer for copy to scratch2. + void* workspace, // Temporary buffer void* output, // Output tensor int batch_size, // Batch size (B) int sequence_length, // Sequence length (S) @@ -30,8 +39,6 @@ bool LaunchLongformerAttentionKernel( 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) - void* workspace, // Temporary buffer - cublasHandle_t& cublas, // Cublas handle const size_t element_size // Element size of input tensor ); diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu b/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu new file mode 100644 index 0000000000..395c19f39b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.cu @@ -0,0 +1,76 @@ +/* +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. +*/ + +#include +#include "core/providers/cuda/cuda_common.h" +#include "longformer_global_impl.h" + +using namespace onnxruntime::cuda; +using namespace cub; + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +size_t GetGlobalScratchSize(int batch_size, int sequence_length) { + // Global Index scratch layout: + // [sequence_index: int BxS][tmp_storage: int 1024x1] + return sizeof(int) * (batch_size * sequence_length + 1024); +} + +__global__ void InitSequenceIndexKernel(int* sequence_index, int sequence_length) { + int batch_index = blockIdx.x; + for (int i = threadIdx.x; i < sequence_length; i += blockDim.x) { + sequence_index[batch_index * sequence_length + i] = i; + } +} + +void BuildGlobalIndex( + cudaStream_t stream, + const int* global_attention, + int batch_size, + int sequence_length, + int* global_index, + int* batch_global_num, + void* scratch, + size_t scratch_size) { + int* sequence_index = (int*)scratch; + int* tmp_storage = sequence_index + batch_size * sequence_length; + + InitSequenceIndexKernel<<>>(sequence_index, sequence_length); + + // Determine temporary device storage size. + // For int* inputs/outputs, it need 767 bytes. When data type changes, its size will be different. + size_t temp_storage_bytes = 0; + cub::DevicePartition::Flagged(NULL, temp_storage_bytes, sequence_index, + global_attention, global_index, batch_global_num, sequence_length, stream); + if (temp_storage_bytes + sizeof(int) * batch_size * sequence_length > scratch_size) { + ORT_THROW("LongformerAttention scratch space is not large enough. Temp storage bytes are", temp_storage_bytes); + } + + // Find the global attention indices and number of global attention tokens + for (int i = 0; i < batch_size; ++i) { + cub::DevicePartition::Flagged(reinterpret_cast(tmp_storage), temp_storage_bytes, sequence_index, + global_attention + i * sequence_length, global_index + i * sequence_length, + batch_global_num + i, sequence_length, stream); + } + + return; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.h b/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.h new file mode 100644 index 0000000000..4e3bd13cd4 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/longformer_global_impl.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Size of global Index scratch in bytes. +size_t GetGlobalScratchSize(int batch_size, int sequence_length); + +// Find the global attention indices and number of global attention tokens +void BuildGlobalIndex( + cudaStream_t stream, + const int* global_attention, + int batch_size, + int sequence_length, + int* global_index, + int* batch_global_num, + void* scratch, + size_t scratch_size); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/transformers/longformer/__init__.py b/onnxruntime/python/tools/transformers/longformer/__init__.py new file mode 100644 index 0000000000..ad5632855c --- /dev/null +++ b/onnxruntime/python/tools/transformers/longformer/__init__.py @@ -0,0 +1,4 @@ +import os +import sys + +sys.path.append(os.path.dirname(__file__)) diff --git a/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py b/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py index fe0278a28c..80f9d44a70 100644 --- a/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py +++ b/onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py @@ -15,6 +15,8 @@ import os import sys import torch import onnxruntime +import numpy as np +import pprint # Mapping from model name to pretrained model name MODELS = { @@ -25,19 +27,15 @@ MODELS = { is_debug = False sys.path.append(os.path.join(os.path.dirname(__file__), '..')) - -# Run onnx model with ORT import benchmark_helper -def get_dummy_inputs(sequence_length, num_global_tokens, device): - # 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) + +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 - # TODO: support more inputs like token_type_ids, position_ids return input_ids, attention_mask, global_attention_mask @@ -52,11 +50,8 @@ def diff_outputs(ort_outputs, torch_outputs): return max_diff -def test_torch(device, model, model_name, batch_sizes, sequence_lengths, global_lengths, test_times, num_threads): - # Comment the following so that PyTorch use default setting as well. - #if num_threads <= 0: - # import psutil - # num_threads = psutil.cpu_count(logical=False) +def test_torch(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) @@ -65,8 +60,8 @@ def test_torch(device, model, model_name, batch_sizes, sequence_lengths, global_ 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 print(f"batch_size={batch_size} sequence_length={sequence_length} global_length={global_length}...") - input_ids, attention_mask, global_attention_mask = get_dummy_inputs(sequence_length, global_length, - device) + input_ids, attention_mask, global_attention_mask = get_dummy_inputs(batch_size, sequence_length, + global_length, device) # Run PyTorch _ = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask) @@ -105,7 +100,8 @@ def test_onnxruntime(device, test_times, num_threads, optimizer=False, - precision='fp32'): + precision='fp32', + verbose=True): results = [] for batch_size in batch_sizes: for sequence_length in sequence_lengths: # This is total length of . @@ -113,8 +109,8 @@ def test_onnxruntime(device, print( f"Testing batch_size={batch_size} sequence_length={sequence_length} global_length={global_length} optimizer={optimizer}, precision={precision}..." ) - input_ids, attention_mask, global_attention_mask = get_dummy_inputs(sequence_length, global_length, - device) + input_ids, attention_mask, global_attention_mask = get_dummy_inputs(batch_size, sequence_length, + global_length, device) # Run OnnxRuntime ort_inputs = { @@ -123,10 +119,13 @@ def test_onnxruntime(device, "global_attention_mask": global_attention_mask.cpu().numpy() } + if verbose: + pprint.pprint(ort_inputs) + # run one query for warm up ort_outputs = ort_session.run(None, ort_inputs) - if is_debug: + if verbose: # Run PyTorch then compare the results with OnnxRuntime. torch_outputs = model(input_ids, attention_mask=attention_mask, @@ -157,7 +156,7 @@ def test_onnxruntime(device, 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, @@ -169,7 +168,14 @@ def test_onnxruntime(device, output_buffer_max_sizes=[max_last_state_size, max_pooler_size], batch_size=batch_size, device=device) - print(result) + """ + 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) return results @@ -204,7 +210,7 @@ def test_all(args): 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.global_lengths, args.test_times, num_threads, args.verbose) if "onnxruntime" in args.engines: session = benchmark_helper.create_onnxruntime_session(onnx_model_path, @@ -212,7 +218,8 @@ def test_all(args): 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.global_lengths, args.test_times, num_threads, optimized, precision, + args.verbose) return results @@ -258,6 +265,8 @@ def parse_arguments(): parser.add_argument("-n", "--num_threads", required=False, nargs="+", type=int, default=[0], help="Threads to use") + parser.add_argument("--verbose", required=False, action="store_true", help="Print more information") + args = parser.parse_args() return args @@ -282,9 +291,9 @@ def output_summary(results, csv_filename, args): for threads in args.num_threads: row = {} for result in results: - if result["model_name"] == model and result["inputs"] == input_count and result[ - "engine"] == engine_name and result["io_binding"] == io_binding and result[ - "threads"] == threads: + if result["model_name"] == model and result["inputs"] == input_count and \ + result["engine"] == engine_name and result["io_binding"] == io_binding and \ + result["threads"] == threads: headers = {k: v for k, v in result.items() if k in header_names} if not row: row.update(headers) 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 bcc1ff31a7..d0cdea4302 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,9 @@ # 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. -# Conversion tested in Ubuntu 18.04 in WSL (Windows Subsystem for Linux), python 3.6, onnxruntime 1.5.2, PyTorch 1.6.0+cpu, transformers 3.0.2 +# 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 onnxruntime-gpu 1.6.0 (or nightly build). +# For inference of the onnx model, you will need latest onnxruntime-gpu 1.7.0 or above. import torch import numpy as np @@ -109,6 +109,43 @@ example_outputs = model(input_ids, attention_mask=attention_mask, global_attenti # A new function to replace LongformerSelfAttention.forward +#For transformer 4.3 +def my_longformer_self_attention_forward_4_3(self, + hidden_states, + attention_mask=None, + is_index_masked=None, + 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 + + #For transformers 4.0 def my_longformer_self_attention_forward_4(self, hidden_states, @@ -181,6 +218,23 @@ def my_longformer_attention_forward_3(self, hidden_states, attention_mask, outpu # 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 @@ -189,6 +243,7 @@ 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)