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

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 <typename T>
 LongformerAttention<T>::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {}
 
@@ -56,8 +75,41 @@ Status LongformerAttention<T>::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<int>(batch_size * sequence_length);
+  auto batch_global_num_buffer = GetScratchBuffer<int>(batch_size);
+
+  size_t global_scratch_bytes = GetGlobalScratchSize(batch_size, sequence_length);
+  auto global_scratch_buffer = GetScratchBuffer<void>(global_scratch_bytes);
+
+  BuildGlobalIndex(
+      stream,
+      global_attention->template Data<int>(),
+      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<void>(pinned_buffer_bytes);
+  int* batch_global_num_pinned = reinterpret_cast<int*>(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<T>::ComputeInternal(OpKernelContext* context) const {
       reinterpret_cast<const CudaT*>(input->template Data<T>()), k,
       &one, reinterpret_cast<CudaT*>(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<T>::ComputeInternal(OpKernelContext* context) const {
         &one, reinterpret_cast<CudaT*>(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<void>(workSpaceSize);
   if (!LaunchLongformerAttentionKernel(
           device_prop,
-          Stream(),
+          cublas,
+          stream,
           reinterpret_cast<const CudaT*>(gemm_buffer.get()),
           reinterpret_cast<const CudaT*>(mask->template Data<T>()),
           reinterpret_cast<const CudaT*>(global_gemm_buffer.get()),
           global_attention->template Data<int>(),
+          global_index_buffer.get(),
+          batch_global_num_buffer.get(),
+          pinned_buffer.get(),
+          workspace_buffer.get(),
           output->template MutableData<T>(),
           batch_size,
           sequence_length,
@@ -123,14 +193,14 @@ Status LongformerAttention<T>::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();
 }
 

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
 );
 

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 <cub/cub.cuh>
+#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<<<batch_size, 128, 0, stream>>>(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<void*>(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

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

onnxruntime/python/tools/transformers/longformer/__init__.py

@@ -0,0 +1,4 @@
+import os
+import sys
+
+sys.path.append(os.path.dirname(__file__))

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 <query, document>.
             for global_length in global_lengths:  # This is length of <query>. 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 <query, document>.
@@ -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)

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)