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

Fix longformer parity and perf regression (#6760)

Browse source 
* add fast kernel back, update benchmark and conversion scripts
This commit is contained in:
Tianlei Wu 2021-02-19 21:47:36 -08:00 committed by GitHub
parent 47519623cd
commit 3bda7f4d36
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 1460 additions and 401 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h
View file

@ -25,5 +25,10 @@ class LongformerAttentionBase {
int window_; // Attention windows length (W). It is half (one-sided) of total window size.
};
namespace longformer {
// Environment variable to give a hint about choosing kernels for less memory or latency.
constexpr const char* kUseCompactMemory = "ORT_LONGFORMER_COMPACT_MEMORY";
} // namespace longformer
} // namespace contrib
} // namespace onnxruntime

20
onnxruntime/contrib_ops/cuda/bert/longformer_attention.cc
View file

@ -5,6 +5,7 @@
#include "core/framework/tensorprotoutils.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/shared_inc/fpgeneric.h"
#include "core/platform/env_var_utils.h"
#include "longformer_global_impl.h"
#include "longformer_attention_impl.h"
@ -49,7 +50,9 @@ class AutoDestoryCudaEvent {
};
template <typename T>
LongformerAttention<T>::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {}
LongformerAttention<T>::LongformerAttention(const OpKernelInfo& info) : CudaKernel(info), LongformerAttentionBase(info) {
use_compact_memory_ = ParseEnvironmentVariableWithDefault<bool>(longformer::kUseCompactMemory, false);
}
template <typename T>
Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
@ -80,6 +83,7 @@ Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
constexpr size_t element_size = sizeof(T);
// TODO: only calculate once per model.
// Build Global Index
auto global_index_buffer = GetScratchBuffer<int>(batch_size * sequence_length);
auto batch_global_num_buffer = GetScratchBuffer<int>(batch_size);
@ -148,10 +152,11 @@ Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
}
}
// Cuda kernel implementation has a limitation of number of global tokens.
if (max_num_global > window_) {
ORT_THROW("LongformerAttention CUDA operator does not support number of global tokens > attention window.");
}
// Force to use fast kernel in two situations:
// (1) global tokens > windows size. In that case, compact memory kernel cannot be used.
// (2) sequence_length == 2 * attention_window. Use fast kernel to walk around parity issue of compact memory kernel.
// In other case, we will choose according to user's environment variable setting (default is fast kernel).
bool use_fast_kernel = (max_num_global > window_ || sequence_length == 2 * window_ || !use_compact_memory_);
// Fully connection for global projection.
// Note that Q only need handle global query tokens if we split GEMM to global Q/K/V separately.
@ -172,7 +177,7 @@ Status LongformerAttention<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, window_);
size_t workSpaceSize = GetLongformerAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, max_num_global, window_, use_fast_kernel);
auto workspace_buffer = GetScratchBuffer<void>(workSpaceSize);
if (!LaunchLongformerAttentionKernel(
device_prop,
@ -193,7 +198,8 @@ Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
head_size,
window_,
max_num_global,
element_size)) {
element_size,
use_fast_kernel)) {
// Get last error to reset it to cudaSuccess.
CUDA_CALL(cudaGetLastError());
return Status(common::ONNXRUNTIME, common::FAIL);

3
onnxruntime/contrib_ops/cuda/bert/longformer_attention.h
View file

@ -18,6 +18,9 @@ class LongformerAttention final : public CudaKernel, public LongformerAttentionB
public:
LongformerAttention(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* context) const override;
private:
bool use_compact_memory_;
};
} // namespace cuda

135
onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.cu
View file

@ -28,7 +28,7 @@ limitations under the License.
#include "core/providers/cuda/cuda_common.h"
#include "longformer_attention_impl.h"
#include "attention_impl.h"
#include "attention_softmax.h"
#include "longformer_attention_softmax.h"
using namespace onnxruntime::cuda;
using namespace cub;
@ -53,11 +53,8 @@ namespace cuda {
// [SoftmaxSpace: see below] [Q:BxNxSxH] [K:BxNxSxH] [V:BxNxSxH] [Global_Q:BxNxSxH] [Global_K:BxNxSxH] [Global_V:BxNxSxH]
// where Global_Q, Global_K and Global_V are optional. They are not allocated when there is no global token.
//
// It is feasible to use compact format for Global_Q with shape BxNxGxH to save space. We do not use compact format for now.
//
// SoftmaxSpace layout:
// [scratch1: (5S-3W)*W*N*B][scratch2: size_t 20]
//
// Scratch1 has 5 buffers for local and global attention calculation.
// Scratch2 has 5 input pointers, 5 output pointers, 5 buffer sizes and 5 strides related to scratch1.
@ -74,10 +71,17 @@ size_t GetLongformerSoftmaxWorkspaceSize(
int batch_size,
int num_heads,
int sequence_length,
int window) {
size_t scratch1_size = GetScratch1Size(element_size, batch_size, num_heads, sequence_length, window);
size_t scratch2_size = 10 * (sizeof(void*) + sizeof(size_t));
return scratch1_size + scratch2_size;
int window,
bool use_fast_kernel) {
if (!use_fast_kernel) {
size_t scratch1_size = GetScratch1Size(element_size, batch_size, num_heads, sequence_length, window);
size_t scratch2_size = 10 * (sizeof(void*) + sizeof(size_t));
return scratch1_size + scratch2_size;
} else {
// Non-compact layout when environment variable ORT_LONGFORMER_COMPACT_MEMORY=0 is set.
// [scratch1: BxNxSxS] [scratch2: BxNxSxS]
return 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length);
}
}
size_t GetLongformerAttentionWorkspaceSize(
@ -87,8 +91,9 @@ size_t GetLongformerAttentionWorkspaceSize(
int head_size,
int sequence_length,
int max_num_global,
int window) {
size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window);
int window,
bool use_fast_kernel) {
size_t softmax_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window, use_fast_kernel);
size_t qkv_size = 3 * batch_size * sequence_length * num_heads * head_size * element_size;
size_t global_qkv_size = max_num_global > 0 ? qkv_size : 0;
return softmax_size + qkv_size + global_qkv_size;
@ -100,6 +105,7 @@ size_t GetPinnedBufferSize(int batch_size) {
return sizeof(int) * batch_size + GetScratch2Size();
}
// Softmax kernel for compact format
template <typename T, int blockSize>
__launch_bounds__(blockSize)
__global__ void LongformerSoftmaxKernel(const int* global_attention,
@ -354,7 +360,6 @@ bool launchSoftmaxKernel(
cudaStream_t stream,
cublasHandle_t cublas,
void* workspace,
size_t softmax_workspace_size,
const void* q, // transposed Q with shape (B, N, S, H)
const void* k, // transposed K with shape (B, N, S, H)
const void* v, // transposed V with shape (B, N, S, H)
@ -373,10 +378,7 @@ bool launchSoftmaxKernel(
int num_heads, // number of heads
int head_size, // hidden size per head
int window, // one sided window size
int max_num_global, // maximum number of global tokens (G) in all batches
size_t element_size) { // size of element: 2 for half, and 4 for float
assert(max_num_global <= window);
const int* global_count = reinterpret_cast<const int*>(pinned_buffer);
bool is_fp16 = (element_size == 2);
@ -605,7 +607,10 @@ bool launchSoftmaxKernel(
resultType,
algo));
void* global_q_batch = (char*)global_q + (i * elements_per_batch) * element_size; // For compact format: replace elements_per_batch by num_heads * max_num_global * head_size
// It is feasible to use compact format for Global_Q with shape BxNxGxH to save space.
// In that case, elements_per_batch is num_heads * max_num_global * head_size, and stride_per_head is max_num_global * head_size.
void* global_q_batch = (char*)global_q + (i * elements_per_batch) * element_size;
void* global_k_batch = (char*)global_k + (i * elements_per_batch) * element_size;
qk_batch = (char*)input_pointers[4] + (i * buffer_sizes[4] * num_heads) * element_size;
@ -625,7 +630,7 @@ bool launchSoftmaxKernel(
global_q_batch, // B
Btype, // B type
head_size, // ldb
stride_per_head, // strideB. For compact format: max_num_global * head_size.
stride_per_head, // strideB.
beta_0, // beta
qk_batch, // C
Ctype, // C type
@ -827,8 +832,9 @@ bool LongformerQkvToContext(
const T* global_input, const int* global_attention,
const int* global_index, const int* batch_global_num, const int max_num_global,
void* pinned_buffer, T* workspace,
T* output) {
size_t softmax_workspace_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window);
T* output,
size_t softmax_workspace_size,
bool use_fast_kernel) {
T* qkv = reinterpret_cast<T*>((char*)workspace + softmax_workspace_size);
// Number of elements in Q, K, V, Global_Q, Global_K or Global_V are same: BxNxSxH
@ -862,33 +868,62 @@ bool LongformerQkvToContext(
const float rsqrt_head_size = 1.f / sqrt(static_cast<float>(head_size));
T* temp_output = qkv; // Q will be overwritten
if (!launchSoftmaxKernel(
stream,
cublas,
workspace,
softmax_workspace_size,
q, // Transposed Q with shape B x N x S x H
k, // Transposed K with shape B x N x S x H
v, // Transposed V with shape B x N x S x H
attention_mask, // Attention mask flags with shape B x S
global_q, // Transposed global Q with shape B x N x S x H.
global_k, // Transposed global K with shape B x N x S x H
global_v, // Transposed global V with shape B x N x S x H
global_attention, // Global attention flags with shape B x S
global_index, // Global index with shape B x S
batch_global_num, // Number of global token per batch with shape B x 1
pinned_buffer, // Pinned Memory Buffer
temp_output, // Output with shape B x N x S x H
rsqrt_head_size, // Scaler
batch_size, // Batch size
sequence_length, // Sequence length
num_heads, // Number of attention heads
head_size, // Hidden size per head
window, // Half (one-sided) window size
max_num_global, // Maximum number of global tokens (G)
element_size)) {
return false;
if (use_fast_kernel) {
if (!launchSoftmaxFastKernel(
stream,
cublas,
workspace, // softmax space
q, // transposed Q with shape (B, N, S, H)
k, // transposed K with shape (B, N, S, H)
v, // transposed V with shape (B, N, S, H)
attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked.
global_q, // Q for global tokens with shape (B, N, S, H)
global_k, // K for global tokens with shape (B, N, S, H)
global_v, // V for global tokens with shape (B, N, S, H)
global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention.
global_index, // Global index with shape (B, S)
batch_global_num, // Number of global tokens per batch with shape (B, 1)
pinned_buffer, // Pinned memory in CPU. Number of global tokens per batch with shape (B, 1)
temp_output, // output with shape (B, N, S, H)
rsqrt_head_size, // scalar
batch_size, // batch size
sequence_length, // sequence length
num_heads, // number of heads
head_size, // hidden size per head
window, // Half (one-sided) window size
element_size)) {
return false;
}
} else {
assert(max_num_global <= window);
if (!launchSoftmaxKernel(
stream,
cublas,
workspace, // softmax space
q, // Transposed Q with shape B x N x S x H
k, // Transposed K with shape B x N x S x H
v, // Transposed V with shape B x N x S x H
attention_mask, // Attention mask flags with shape B x S. Value -10000.0 means masked, and 0.0 not mased.
global_q, // Transposed global Q with shape B x N x S x H.
global_k, // Transposed global K with shape B x N x S x H
global_v, // Transposed global V with shape B x N x S x H
global_attention, // Global attention flags with shape B x S
global_index, // Global index with shape B x S
batch_global_num, // Number of global token per batch with shape B x 1
pinned_buffer, // Pinned Memory Buffer
temp_output, // Output with shape B x N x S x H
rsqrt_head_size, // Scaler
batch_size, // Batch size
sequence_length, // Sequence length
num_heads, // Number of attention heads
head_size, // Hidden size per head
window, // Half (one-sided) window size
element_size)) {
return false;
}
}
// The temp_output is BxNxSxH, transpose it to final output BxSxNxH
return LaunchTransCtx(stream, sequence_length, batch_size, head_size, num_heads, temp_output, output);
@ -913,8 +948,10 @@ bool LaunchLongformerAttentionKernel(
int head_size,
int window,
int max_num_global,
const size_t element_size) {
const size_t element_size,
bool use_fast_kernel) {
CublasMathModeSetter helper(device_prop, cublas, CUBLAS_TENSOR_OP_MATH);
size_t softmax_workspace_size = GetLongformerSoftmaxWorkspaceSize(element_size, batch_size, num_heads, sequence_length, window, use_fast_kernel);
if (element_size == 2) {
return LongformerQkvToContext(cublas, stream,
batch_size, sequence_length, num_heads, head_size, window, element_size,
@ -927,7 +964,9 @@ bool LaunchLongformerAttentionKernel(
max_num_global,
pinned_buffer,
reinterpret_cast<half*>(workspace),
reinterpret_cast<half*>(output));
reinterpret_cast<half*>(output),
softmax_workspace_size,
use_fast_kernel);
} else {
return LongformerQkvToContext(cublas, stream,
batch_size, sequence_length, num_heads, head_size, window, element_size,
@ -940,7 +979,9 @@ bool LaunchLongformerAttentionKernel(
max_num_global,
pinned_buffer,
reinterpret_cast<float*>(workspace),
reinterpret_cast<float*>(output));
reinterpret_cast<float*>(output),
softmax_workspace_size,
use_fast_kernel);
}
}

8
onnxruntime/contrib_ops/cuda/bert/longformer_attention_impl.h
View file

@ -18,9 +18,10 @@ size_t GetLongformerAttentionWorkspaceSize(
int head_size,
int sequence_length,
int max_num_global,
int window);
int window,
bool use_fast_kernel);
bool LaunchLongformerAttentionKernel(
bool LaunchLongformerAttentionKernel(
const cudaDeviceProp& device_prop, // Device Properties
cublasHandle_t& cublas, // Cublas handle
cudaStream_t stream, // CUDA stream
@ -39,7 +40,8 @@ size_t GetLongformerAttentionWorkspaceSize(
int head_size, // Hidden layer size per head (H)
int window, // One sided attention window (W)
int max_num_global, // Maximum number of global tokens (G)
const size_t element_size // Element size of input tensor
const size_t element_size, // Element size of input tensor,
bool use_fast_kernel // Use compact memory
);
} // namespace cuda

657
onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.cu Normal file
View file

@ -0,0 +1,657 @@
/*
Copyright (c) NVIDIA Corporation and Microsoft Corporation
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
// This is fast cuda kernels for longformer attention softmax.
// It uses two temporary matrix of BxNxSxS, and consumes more memory when sequence length is large.
#include <cub/cub.cuh>
#include <cublas_v2.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <math_constants.h>
#include "core/providers/cuda/cu_inc/common.cuh"
#include "core/providers/cuda/cuda_common.h"
#include "longformer_attention_softmax.h"
#include "attention_impl.h"
using namespace onnxruntime::cuda;
using namespace cub;
#define CHECK(expr) \
if (!CUBLAS_CALL(expr)) { \
return false; \
}
namespace onnxruntime {
namespace contrib {
namespace cuda {
template <typename T, int blockSize>
__launch_bounds__(blockSize)
__global__ void LongformerSoftmaxFastKernel(const int* global_attention,
const int* global_index,
const int* batch_global_num,
const T* input,
const T* attention_mask,
T* output,
float scaler,
int dim0,
int sequence_length,
int attention_window) {
typedef cub::BlockReduce<float, blockSize> BlockReduce;
__shared__ typename BlockReduce::TempStorage block_reduce_temp;
__shared__ float max_shared;
__shared__ float sum_shared;
const T* input_block = input + sequence_length * blockIdx.x;
T* output_block = output + sequence_length * blockIdx.x;
const int batch_index = blockIdx.x / dim0;
const int row_index = blockIdx.x % sequence_length;
const int global_num = batch_global_num[batch_index];
// To be consistent with Huggingface Longformer, the row of maksed word are set as zero.
if ((float)attention_mask[batch_index * sequence_length + row_index] < 0.0f) {
for (int i = threadIdx.x; i < sequence_length; i += blockSize) {
output_block[i] = (T)(0);
}
return;
}
// local attention token
int col_start = 0;
int col_end = sequence_length;
bool is_local_row = (global_attention[batch_index * sequence_length + row_index] == (int)0);
if (is_local_row) {
col_start = row_index - attention_window;
if (col_start < 0) {
col_start = 0;
}
col_end = row_index + attention_window + 1;
if (col_end > sequence_length) {
col_end = sequence_length;
}
}
const T* mask_block = attention_mask + sequence_length * batch_index;
int tid = threadIdx.x;
// calculate max input
float max_input = -CUDART_INF_F;
// #pragma unroll 16
for (int i = tid + col_start; i < col_end; i += blockSize) {
float x = input_block[i];
x = x * scaler + (float)mask_block[i];
if (max_input < x) {
max_input = x;
}
}
if (is_local_row) {
for (int g = tid; g < global_num; g += blockSize) {
int i = global_index[g];
if (i < col_start || i > col_end) {
float x = input_block[i];
x = x * scaler + (float)mask_block[i];
if (max_input < x) {
max_input = x;
}
}
}
}
float max_block = BlockReduce(block_reduce_temp).Reduce(max_input, cub::Max());
if (tid == 0) {
max_shared = max_block;
}
__syncthreads();
float sum_input = 0.f;
// #pragma unroll 16
for (int i = tid + col_start; i < col_end; i += blockSize) {
float x = input_block[i];
x = expf((x)*scaler + (float)mask_block[i] - max_shared);
sum_input += x;
}
if (is_local_row) {
for (int g = tid; g < global_num; g += blockSize) {
int i = global_index[g];
if (i < col_start || i > col_end) {
float x = input_block[i];
x = expf((x)*scaler + (float)mask_block[i] - max_shared);
sum_input += x;
}
}
}
float sum_block = BlockReduce(block_reduce_temp).Reduce(sum_input, cub::Sum());
if (tid == 0) {
sum_shared = sum_block;
}
__syncthreads();
float recip_sum = 1.f / sum_shared;
if (is_local_row) {
// We only need to fill in zeros for blocks that will be used in the matrix multiplication
// following the Softmax.
//
// For now zero-out only [row_index - 2*attention_window, row_index + 2*attention_window],
// we can even be more agressive and reduce the zeroing out window size since
// each row has entries in 3 blocks (3*attention_window size instead of 4*attention_window)
int zero_start = row_index - 2 * attention_window;
if (zero_start < 0) {
zero_start = 0;
}
int zero_end = row_index + 2 * attention_window;
if (zero_end > sequence_length) {
zero_end = sequence_length;
}
for (int i = tid + zero_start; i < zero_end; i += blockSize) {
output_block[i] = (T)(0.);
}
for (int g = tid; g < global_num; g += blockSize) {
int i = global_index[g];
float x = input_block[i];
x = expf((x)*scaler + (float)mask_block[i] - max_shared);
output_block[i] = (T)(recip_sum * x);
}
}
// #pragma unroll 16
for (int i = tid + col_start; i < col_end; i += blockSize) {
float x = input_block[i];
x = expf((x)*scaler + (float)mask_block[i] - max_shared);
output_block[i] = (T)(recip_sum * x);
}
}
// Launch the softmax kernel for non compact memory.
bool launchSoftmaxFastKernel(
cudaStream_t stream,
cublasHandle_t cublas,
void* workspace, // softmax space
const void* q, // transposed Q with shape (B, N, S, H)
const void* k, // transposed K with shape (B, N, S, H)
const void* v, // transposed V with shape (B, N, S, H)
const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked.
const void* global_q, // Q for global tokens with shape (B, N, S, H)
const void* global_k, // K for global tokens with shape (B, N, S, H)
const void* global_v, // V for global tokens with shape (B, N, S, H)
const int* global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention.
const int* global_index, // Global index with shape (B, S)
const int* batch_global_num, // Number of global tokens per batch with shape (B, 1)
void* pinned_buffer, // Pinned memory in CPU. Number of global tokens per batch with shape (B, 1)
void* output, // output with shape (B, N, S, H)
float scaler, // scalar
int batch_size, // batch size
int sequence_length, // sequence length
int num_heads, // number of heads
int head_size, // hidden size per head
int attention_window, // one sided windows size
size_t element_size) { // size of element: 2 for half, and 4 for float
bool is_fp16 = (element_size == 2);
void* scratch1 = reinterpret_cast<char*>(workspace);
void* scratch2 = reinterpret_cast<char*>(scratch1) + GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length, sequence_length);
// setup shared parameters for two strided batched matrix multiplies
cudaDataType_t Atype;
cudaDataType_t Btype;
cudaDataType_t Ctype;
cudaDataType_t resultType;
cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT;
__half one_fp16, zero_fp16;
float one_fp32, zero_fp32;
void *alpha, *beta_0, *beta_1;
if (is_fp16) {
one_fp16 = __float2half(1.f);
zero_fp16 = __float2half(0.f);
alpha = static_cast<void*>(&one_fp16);
beta_0 = static_cast<void*>(&zero_fp16);
beta_1 = static_cast<void*>(&one_fp16);
Atype = CUDA_R_16F;
Btype = CUDA_R_16F;
Ctype = CUDA_R_16F;
resultType = CUDA_R_16F;
algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP;
} else {
one_fp32 = 1.f;
zero_fp32 = 0.f;
alpha = static_cast<void*>(&one_fp32);
beta_0 = static_cast<void*>(&zero_fp32);
beta_1 = static_cast<void*>(&one_fp32);
Atype = CUDA_R_32F;
Btype = CUDA_R_32F;
Ctype = CUDA_R_32F;
resultType = CUDA_R_32F;
}
// Strided batch matrix multiply
// qk = q * k^T
// Shapes: q and k = B x N x S x H, qk = B x N x S x S
// Convert col-major to row-major by swapping q and k in Gemm
// Local attention part
// S x S is calculated using sliding block WxW (W is one sided window size) like the following:
// [W][W]
// [W][W][W]
// [W][W][W]
// [W][W]
// The first and last rows have 2 blocks, and the remaining has 3 blocks per row.
// The calculation are splited into 3 parts. Firstly, fill the middle rows, then the first row and finally the last row.
// The results are stored in scratch1.
int w = attention_window;
int x_offset = num_heads * sequence_length * head_size;
int y_offset = num_heads * sequence_length * sequence_length;
int last_block = (sequence_length / w) - 1;
int strideA = sequence_length * head_size;
int strideB = sequence_length * head_size;
int strideC = sequence_length * sequence_length;
// When S == 2W, there is no middle rows of blocks:
// [W][W]
// [W][W]
// We can use normal matrix multiplication in this case.
if (sequence_length == 2 * w) {
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_T,
CUBLAS_OP_N,
sequence_length,
sequence_length,
head_size,
alpha,
k,
Atype,
head_size,
sequence_length * head_size,
q,
Btype,
head_size,
sequence_length * head_size,
beta_0,
scratch1,
Ctype,
sequence_length,
sequence_length * sequence_length,
batch_size * num_heads,
resultType,
algo));
} else { // sequence_length > 2 * w
for (int i = 0; i < batch_size; ++i) {
for (int j = 0; j < num_heads; ++j) {
void* q_head = (char*)q + (i * x_offset + j * sequence_length * head_size + w * head_size) * element_size;
void* k_head = (char*)k + (i * x_offset + j * sequence_length * head_size) * element_size;
void* qk_head = (char*)scratch1 + (i * y_offset + j * sequence_length * sequence_length + w * sequence_length) * element_size;
int count = (sequence_length - 2 * w) / w;
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_T,
CUBLAS_OP_N,
3 * w, // m
w, // n
head_size, // k
alpha, // alpha
k_head, // A
Atype, // A type
head_size, // lda
w * head_size, // strideA
q_head, // B
Btype, // B type
head_size, // ldb
w * head_size, // strideB
beta_0, // beta
qk_head, // C
Ctype, // C type
sequence_length, // ldc
sequence_length * w + w, // strideC
count, // batch count
resultType,
algo));
}
}
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_T,
CUBLAS_OP_N,
2 * w, // m
w, // n
head_size, // k
alpha, // alpha
k, // A
Atype, // A type
head_size, // lda
strideA, // strideA
q, // B
Btype, // B type
head_size, // ldb
strideB, // strideB
beta_0, // beta
scratch1, // C
Ctype, // C type
sequence_length, // ldc
strideC, // strideC
batch_size * num_heads, // batch count
resultType,
algo));
void* q_head = (char*)q + (last_block * w * head_size) * element_size;
void* k_head = (char*)k + ((last_block - 1) * w * head_size) * element_size;
void* qk_head = (char*)scratch1 + (last_block * w * sequence_length + (last_block - 1) * w) * element_size;
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_T,
CUBLAS_OP_N,
2 * w,
w,
head_size,
alpha,
k_head,
Atype,
head_size,
strideA,
q_head,
Btype,
head_size,
strideB,
beta_0,
qk_head,
Ctype,
sequence_length,
strideC,
batch_size * num_heads,
resultType,
algo));
}
const int* batch_global_count = reinterpret_cast<const int*>(pinned_buffer);
// Global attention part
for (int i = 0; i < batch_size; ++i) {
if (batch_global_count[i] > 0) {
void* q_batch = (char*)q + (i * x_offset) * element_size;
void* k_batch = (char*)k + (i * x_offset) * element_size;
void* qk_batch = (char*)scratch1 + (i * y_offset) * element_size;
// Local tokens attending global tokens
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_T,
CUBLAS_OP_N,
batch_global_count[i],
sequence_length,
head_size,
alpha,
k_batch,
Atype,
head_size,
strideA,
q_batch,
Btype,
head_size,
strideB,
beta_0,
qk_batch,
Ctype,
sequence_length,
strideC,
num_heads,
resultType,
algo));
void* global_q_batch = (char*)global_q + (i * num_heads * sequence_length * head_size) * element_size;
void* global_k_batch = (char*)global_k + (i * x_offset) * element_size;
int strideB_global = sequence_length * head_size;
// Global tokens attending everything
// This GEMMs need to be last to make sure all global token entries are re-written.
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_T,
CUBLAS_OP_N,
sequence_length,
batch_global_count[i],
head_size,
alpha,
global_k_batch,
Atype,
head_size,
strideA,
global_q_batch,
Btype,
head_size,
strideB_global,
beta_0,
qk_batch,
Ctype,
sequence_length,
strideC,
num_heads,
resultType,
algo));
}
}
int dim0 = sequence_length * num_heads;
int dim1 = sequence_length;
void* softmax_out = scratch2;
const int blockSize = 64;
const int gridSize = batch_size * num_heads * sequence_length;
if (is_fp16) {
LongformerSoftmaxFastKernel<__half, blockSize><<<gridSize, blockSize, 0, stream>>>(
global_attention,
global_index,
batch_global_num,
static_cast<const __half*>(scratch1),
static_cast<const __half*>(attention_mask),
static_cast<__half*>(softmax_out), scaler, dim0, dim1, attention_window);
} else {
LongformerSoftmaxFastKernel<float, blockSize><<<gridSize, blockSize, 0, stream>>>(
global_attention,
global_index,
batch_global_num,
static_cast<const float*>(scratch1),
static_cast<const float*>(attention_mask),
static_cast<float*>(softmax_out), scaler, dim0, dim1, attention_window);
}
// Run the matrix multiply: output = softmax_out * v
// softmax_out: B x N x S x S
// v: B x N x S x H
// attn_out: B x N x S x H
// Calculation uses full Gemm (S == 2W) or sliding blocks (S > 2W) in a way similar to local attention part.
if (sequence_length == 2 * w) {
// convert col-major to row-major by swapping softmax_out and v
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
head_size,
sequence_length,
sequence_length,
alpha,
v,
Atype,
head_size,
sequence_length * head_size,
softmax_out,
Btype,
sequence_length,
sequence_length * sequence_length,
beta_0,
output,
Ctype,
head_size,
sequence_length * head_size,
batch_size * num_heads,
resultType,
algo));
} else { // sequence_length > 2 * w
for (int i = 0; i < batch_size; ++i) {
for (int j = 0; j < num_heads; ++j) {
void* v_head = (char*)v + (i * x_offset + j * head_size * sequence_length) * element_size;
void* prob_head = (char*)softmax_out + (i * y_offset + j * sequence_length * sequence_length + w * sequence_length) * element_size;
void* out_head = (char*)output + (i * x_offset + j * head_size * sequence_length + w * head_size) * element_size;
int count = (sequence_length - 2 * w) / w;
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
head_size,
w,
3 * w,
alpha,
v_head,
Atype,
head_size,
w * head_size,
prob_head,
Btype,
sequence_length,
sequence_length * w + w,
beta_0,
out_head,
Ctype,
head_size,
w * head_size,
count,
resultType,
algo));
}
}
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
head_size,
w,
2 * w,
alpha,
v,
Atype,
head_size,
sequence_length * head_size,
softmax_out,
Btype,
sequence_length,
sequence_length * sequence_length,
beta_0,
output,
Ctype,
head_size,
sequence_length * head_size,
batch_size * num_heads,
resultType,
algo));
void* v_head = (char*)v + (last_block - 1) * w * head_size * element_size;
void* prob_head = (char*)softmax_out + (sequence_length * last_block * w + (last_block - 1) * w) * element_size;
void* out_head = (char*)output + last_block * w * head_size * element_size;
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
head_size,
w,
2 * w,
alpha,
v_head,
Atype,
head_size,
sequence_length * head_size,
prob_head,
Btype,
sequence_length,
sequence_length * sequence_length,
beta_0,
out_head,
Ctype,
head_size,
sequence_length * head_size,
batch_size * num_heads,
resultType,
algo));
}
for (int i = 0; i < batch_size; ++i) {
if (batch_global_count[i] > 0) {
int glob_longdim_mm = (last_block - 1) * w;
void* v_head = (char*)v + (i * x_offset) * element_size;
void* prob_head = (char*)softmax_out + (i * y_offset + 2 * w * sequence_length) * element_size;
void* out_head = (char*)output + (i * x_offset + 2 * w * head_size) * element_size;
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
head_size,
glob_longdim_mm,
batch_global_count[i],
alpha,
v_head,
Atype,
head_size,
sequence_length * head_size,
prob_head,
Btype,
sequence_length,
sequence_length * sequence_length,
beta_1,
out_head,
Ctype,
head_size,
sequence_length * head_size,
num_heads,
resultType,
algo));
// Global tokens
v_head = (char*)global_v + (i * x_offset) * element_size;
prob_head = (char*)softmax_out + (i * y_offset) * element_size;
out_head = (char*)output + (i * x_offset) * element_size;
CHECK(cublasGemmStridedBatchedEx(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
head_size,
batch_global_count[i],
sequence_length, // Re-write entries completely
alpha,
v_head,
Atype,
head_size,
sequence_length * head_size,
prob_head,
Btype,
sequence_length,
sequence_length * sequence_length,
beta_0, // Use beta=0 to overwrite
out_head, // Here assumes global tokens are at the beginning of sequence.
Ctype,
head_size,
sequence_length * head_size,
num_heads,
resultType,
algo));
}
}
return true;
}
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

51
onnxruntime/contrib_ops/cuda/bert/longformer_attention_softmax.h Normal file
View file

@ -0,0 +1,51 @@
/*
Copyright (c) NVIDIA Corporation and Microsoft Corporation
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
// This is fast cuda kernels for longformer attention softmax.
// It uses two temporary matrix of BxNxSxS, and consumes more memory when sequence length is large.
namespace onnxruntime {
namespace contrib {
namespace cuda {
// Launch the softmax kernel for non compact memory.
bool launchSoftmaxFastKernel(
cudaStream_t stream,
cublasHandle_t cublas,
void* workspace, // softmax space
const void* q, // transposed Q with shape (B, N, S, H)
const void* k, // transposed K with shape (B, N, S, H)
const void* v, // transposed V with shape (B, N, S, H)
const void* attention_mask, // attention mask with shape (B, S), with value 0.0 not masked, and -10000.0 masked.
const void* global_q, // Q for global tokens with shape (B, N, S, H)
const void* global_k, // K for global tokens with shape (B, N, S, H)
const void* global_v, // V for global tokens with shape (B, N, S, H)
const int* global_attention, // global attention with shape (B, S), with value 0 for local attention and 1 for global attention.
const int* global_index, // Global index with shape (B, S)
const int* batch_global_num, // Number of global tokens per batch with shape (B, 1)
void* pinned_buffer, // Pinned memory in CPU. Number of global tokens per batch with shape (B, 1)
void* output, // output with shape (B, N, S, H)
float scaler, // scalar
int batch_size, // batch size
int sequence_length, // sequence length
int num_heads, // number of heads
int head_size, // hidden size per head
int attention_window, // one sided windows size
size_t element_size);
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

82
onnxruntime/python/tools/transformers/benchmark_helper.py
View file

@ -235,3 +235,85 @@ def allocateOutputBuffers(output_buffers, output_buffer_max_sizes, device):
for i in output_buffer_max_sizes:
output_buffers.append(torch.empty(i, dtype=torch.float32, device=device))
def set_random_seed(seed=123):
"""Set random seed manully to get deterministic results"""
import random
random.seed(seed)
numpy.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
#torch.backends.cudnn.enabled = False
#torch.backends.cudnn.benchmark = False
#torch.backends.cudnn.deterministic = True
def measure_memory(is_gpu, func):
import os
import psutil
from time import sleep
class MemoryMonitor:
def __init__(self, keep_measuring=True):
self.keep_measuring = keep_measuring
def measure_cpu_usage(self):
max_usage = 0
while True:
max_usage = max(max_usage, psutil.Process(os.getpid()).memory_info().rss / 1024**2)
sleep(0.005) # 5ms
if not self.keep_measuring:
break
return max_usage
def measure_gpu_usage(self):
from py3nvml.py3nvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, \
nvmlDeviceGetMemoryInfo, nvmlDeviceGetName, nvmlShutdown, NVMLError
max_gpu_usage = []
gpu_name = []
try:
nvmlInit()
deviceCount = nvmlDeviceGetCount()
max_gpu_usage = [0 for i in range(deviceCount)]
gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(deviceCount)]
while True:
for i in range(deviceCount):
info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2)
sleep(0.005) # 5ms
if not self.keep_measuring:
break
nvmlShutdown()
return [{
"device_id": i,
"name": gpu_name[i],
"max_used_MB": max_gpu_usage[i]
} for i in range(deviceCount)]
except NVMLError as error:
if not self.silent:
self.logger.error("Error fetching GPU information using nvml: %s", error)
return None
monitor = MemoryMonitor(False)
if is_gpu:
print(f"GPU memory usage before testing: {monitor.measure_gpu_usage()}")
else:
print(f"Peak CPU memory usage before testing: {monitor.measure_cpu_usage():.2f} MB")
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor() as executor:
monitor = MemoryMonitor()
mem_thread = executor.submit(monitor.measure_gpu_usage if is_gpu else monitor.measure_cpu_usage)
try:
fn_thread = executor.submit(func)
result = fn_thread.result()
finally:
monitor.keep_measuring = False
max_usage = mem_thread.result()
if is_gpu:
print(f"Peak GPU memory usage: {max_usage}")
else:
print(f"Peak CPU memory usage: {max_usage:.2f} MB")
return max_usage

6
onnxruntime/python/tools/transformers/longformer/__init__.py
View file

@ -1,3 +1,9 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import os
import sys

317
onnxruntime/python/tools/transformers/longformer/benchmark_longformer.py
View file

@ -1,11 +1,29 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
#
# This script run benchmark of latency or peak memory usage of Longformer model inference.
#
# Please run convert_longformer_to_onnx.py to get onnx model before running this script.
# Tested with python 3.7, onnxruntime-gpu 1.6.0 (or nightly), PyTorch 1.7.0, transformers 4.0, CUDA 10.2, CUDNN 8.0
# Example step by step command lines for benchmarking longformer base model (without/with optimizer) in Linux:
# Tested with python 3.6, onnxruntime-gpu 1.7.0, PyTorch 1.7.1, transformers 4.3.2, CUDA 10.2.
#
# Example commands for exporting longformer base model in Linux or WSL:
# cd ../torch_extensions
# python setup.py install
# python convert_longformer_to_onnx.py -m longformer-base-4096
# python benchmark_longformer.py -m longformer-base-4096
# python convert_longformer_to_onnx.py -m longformer-base-4096 -o
# python benchmark_longformer.py -m longformer-base-4096
# cd ../longformer
# python convert_longformer_to_onnx.py --model longformer-base-4096 --precision fp32 --optimize_onnx
#
# Benchmark the latency (Exported onnx model is in the current directory):
# python benchmark_longformer.py --models longformer-base-4096 --batch_sizes 1 --sequence_lengths 512 1024 2048 4096 --global_lengths 8 --onnx_dir . --validate_onnx -t 100
#
# Benchmark GPU peak memory:
# export ORT_LONGFORMER_COMPACT_MEMORY=0
# python benchmark_longformer.py --models longformer-base-4096 --batch_sizes 1 --sequence_lengths 4096 --global_lengths 8 --onnx_dir . --memory -t 10
# export ORT_LONGFORMER_COMPACT_MEMORY=1
# python benchmark_longformer.py --models longformer-base-4096 --batch_sizes 1 --sequence_lengths 4096 --global_lengths 8 --onnx_dir . --memory -t 10
# By default, compact memory kernel is not enabled since it is slower. You need set an environment variable ORT_LONGFORMER_COMPACT_MEMORY=1 to enable it, which uses less memory in this test.
import timeit
from datetime import datetime
@ -17,57 +35,30 @@ import torch
import onnxruntime
import numpy as np
import pprint
import math
# Mapping from model name to pretrained model name
MODELS = {
"longformer-base-4096": "allenai/longformer-base-4096",
"longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging
}
is_debug = False
from longformer_helper import LongformerHelper, PRETRAINED_LONGFORMER_MODELS
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
import benchmark_helper
def get_dummy_inputs(batch_size, sequence_length, num_global_tokens, device):
input_ids = torch.randint(low=0, high=100, size=(batch_size, sequence_length), dtype=torch.long, device=device)
attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)
global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device)
global_token_index = list(range(num_global_tokens))
global_attention_mask[:, global_token_index] = 1
return input_ids, attention_mask, global_attention_mask
def diff_outputs(ort_outputs, torch_outputs):
max_diff = []
# Compare the outputs to find max difference
for i in range(2):
print(f"output {i} shape: ORT={ort_outputs[i].shape}, Torch={torch_outputs[i].shape}")
diff = (torch.from_numpy(ort_outputs[i]) - torch_outputs[i].to('cpu')).abs().max()
max_diff.append(diff)
print(f"max diff for output: {max_diff}")
return max_diff
def test_torch(device, model, model_name, batch_sizes, sequence_lengths, global_lengths, test_times, num_threads,
verbose):
def test_torch_latency(device, model, model_name, batch_sizes, sequence_lengths, global_lengths, test_times,
num_threads, verbose):
if num_threads > 0:
torch.set_num_threads(num_threads)
results = []
for batch_size in batch_sizes:
for sequence_length in sequence_lengths: # This is total length of <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
for sequence_length in sequence_lengths:
for global_length in global_lengths:
print(f"batch_size={batch_size} sequence_length={sequence_length} global_length={global_length}...")
input_ids, attention_mask, global_attention_mask = get_dummy_inputs(batch_size, sequence_length,
global_length, device)
inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, global_length,
device)
input_list = inputs.to_list()
# Run PyTorch
_ = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)
runtimes = timeit.repeat(lambda: model(input_ids, attention_mask, global_attention_mask),
repeat=test_times,
number=1)
_ = model(*input_list)
runtimes = timeit.repeat(lambda: model(*input_list), repeat=test_times, number=1)
result = {
"engine": "torch", #TODO: test torchscript
"version": torch.__version__,
@ -90,7 +81,24 @@ def test_torch(device, model, model_name, batch_sizes, sequence_lengths, global_
return results
def test_onnxruntime(device,
def test_parity(device, model, ort_session, batch_size, sequence_length, global_length, verbose=True):
print(
f"Comparing Torch and ORT outputs for batch_size={batch_size} sequence_length={sequence_length} global_length={global_length}..."
)
dummy_inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, global_length,
device)
ort_inputs = dummy_inputs.get_ort_inputs()
ort_outputs = ort_session.run(None, ort_inputs)
input_list = dummy_inputs.to_list()
torch_outputs = model(*input_list)
max_diff = np.amax(torch_outputs[0].cpu().numpy() - ort_outputs[0])
print(f"last_state max diff = {max_diff}")
if verbose and (math.isnan(max_diff) or max_diff > 0.001):
print("torch last_state:", torch_outputs[0])
print("ort last_state:", ort_outputs[0])
def test_ort_latency(device,
model,
model_name,
ort_session,
@ -101,43 +109,30 @@ def test_onnxruntime(device,
num_threads,
optimizer=False,
precision='fp32',
validate_onnx=True,
disable_io_binding=False,
verbose=True):
results = []
for batch_size in batch_sizes:
for sequence_length in sequence_lengths: # This is total length of <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
for sequence_length in sequence_lengths:
for global_length in global_lengths:
assert global_length <= model.config.attention_window[
0], "Limitation of current implementation: number of global token <= attention_window"
print(
f"Testing batch_size={batch_size} sequence_length={sequence_length} global_length={global_length} optimizer={optimizer}, precision={precision}..."
f"Testing batch_size={batch_size} sequence_length={sequence_length} global_length={global_length} optimizer={optimizer}, precision={precision} io_binding={not disable_io_binding}..."
)
input_ids, attention_mask, global_attention_mask = get_dummy_inputs(batch_size, sequence_length,
global_length, device)
dummy_inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length,
global_length, device)
# Run OnnxRuntime
ort_inputs = {
"input_ids": input_ids.cpu().numpy(),
"attention_mask": attention_mask.cpu().numpy(),
"global_attention_mask": global_attention_mask.cpu().numpy()
}
ort_inputs = dummy_inputs.get_ort_inputs()
if verbose:
pprint.pprint(ort_inputs)
print(ort_inputs)
# run one query for warm up
ort_outputs = ort_session.run(None, ort_inputs)
if verbose:
# Run PyTorch then compare the results with OnnxRuntime.
torch_outputs = model(input_ids,
attention_mask=attention_mask,
global_attention_mask=global_attention_mask)
max_diff = diff_outputs(ort_outputs, torch_outputs)
print("max diff for outputs", max_diff)
if max(max_diff) > 0.001:
print("ort_inputs", ort_inputs)
print("ort_outputs", ort_outputs)
device = input_ids.device
result_template = {
"model_name": model_name,
"inputs": 3,
@ -154,32 +149,59 @@ def test_onnxruntime(device,
"datetime": str(datetime.now()),
}
max_last_state_size = max(batch_sizes) * max(sequence_lengths) * model.config.hidden_size
max_pooler_size = max(batch_sizes) * max(sequence_lengths)
"""
result = benchmark_helper.inference_ort_with_io_binding(
ort_session,
ort_inputs,
result_template=result_template,
repeat_times=test_times,
ort_output_names=["last_state", "pooler"],
ort_outputs=ort_outputs,
output_buffers=[],
output_buffer_max_sizes=[max_last_state_size, max_pooler_size],
batch_size=batch_size,
device=device)
"""
result = benchmark_helper.inference_ort(ort_session,
ort_inputs,
result_template=result_template,
repeat_times=test_times,
batch_size=batch_size)
if not disable_io_binding:
max_last_state_size = max(batch_sizes) * max(sequence_lengths) * model.config.hidden_size
max_pooler_size = max(batch_sizes) * max(sequence_lengths)
result = benchmark_helper.inference_ort_with_io_binding(
ort_session,
ort_inputs,
result_template=result_template,
repeat_times=test_times,
ort_output_names=["last_state", "pooler"],
ort_outputs=ort_outputs,
output_buffers=[],
output_buffer_max_sizes=[max_last_state_size, max_pooler_size],
batch_size=batch_size,
device=device,
data_type=np.longlong, #input data type
)
else:
result = benchmark_helper.inference_ort(ort_session,
ort_inputs,
result_template=result_template,
repeat_times=test_times,
batch_size=batch_size)
pprint.pprint(result)
results.append(result)
if validate_onnx:
test_parity(device, model, ort_session, batch_size, sequence_length, global_length, verbose)
return results
def test_ort_memory(device, onnx_model_path, batch_size, sequence_length, global_length, test_times, num_threads):
print(
f"Testing memory for model={onnx_model_path}, batch_size={batch_size}, sequence_length={sequence_length}, global_length={global_length}, test_times={test_times}, num_threads={num_threads}"
)
def inference():
session = benchmark_helper.create_onnxruntime_session(onnx_model_path,
use_gpu=True,
enable_all_optimization=True,
num_threads=num_threads)
dummy_inputs: LongforerInputs = LongformerHelper.get_dummy_inputs(batch_size, sequence_length, global_length,
device)
ort_inputs = dummy_inputs.get_ort_inputs()
for _ in range(test_times):
ort_outputs = session.run(None, ort_inputs)
benchmark_helper.measure_memory(is_gpu=True, func=inference)
print("Memory test is done")
def test_all(args):
# Currently, the longformer attention operator could only run in GPU (no CPU implementation yet).
device = torch.device('cuda:0')
@ -188,17 +210,18 @@ def test_all(args):
for model_name in args.models:
# Here we run an example input
from transformers import LongformerModel
torch_model_name_or_dir = MODELS[model_name]
torch_model_name_or_dir = PRETRAINED_LONGFORMER_MODELS[model_name]
model = LongformerModel.from_pretrained(torch_model_name_or_dir) # pretrained model name or directory
model.to(device)
# Search onnx model in the following order: optimized fp16 model, optimized fp32 model, raw model
# TODO: call convert_longformer_to_onnx to export onnx instead.
import os.path
optimized = False
precision = 'fp32'
onnx_model_path = model_name + ".onnx"
optimized_fp32_model = model_name + "_fp32.onnx"
optimized_fp16_model = model_name + "_fp16.onnx"
import os.path
onnx_model_path = os.path.join(args.onnx_dir, model_name + ".onnx")
optimized_fp32_model = os.path.join(args.onnx_dir, model_name + "_fp32.onnx")
optimized_fp16_model = os.path.join(args.onnx_dir, model_name + "_fp16.onnx")
if os.path.isfile(optimized_fp16_model):
onnx_model_path = optimized_fp16_model
optimized = True
@ -206,24 +229,33 @@ def test_all(args):
elif os.path.isfile(optimized_fp32_model):
onnx_model_path = optimized_fp32_model
optimized = True
print("ONNX model path:", onnx_model_path)
for num_threads in args.num_threads:
if "torch" in args.engines:
results += test_torch(device, model, model_name, args.batch_sizes, args.sequence_lengths,
args.global_lengths, args.test_times, num_threads, args.verbose)
results += test_torch_latency(device, model, model_name, args.batch_sizes, args.sequence_lengths,
args.global_lengths, args.test_times, num_threads, args.verbose)
if "onnxruntime" in args.engines:
session = benchmark_helper.create_onnxruntime_session(onnx_model_path,
use_gpu=True,
enable_all_optimization=True,
num_threads=num_threads)
results += test_onnxruntime(device, model, model_name, session, args.batch_sizes, args.sequence_lengths,
args.global_lengths, args.test_times, num_threads, optimized, precision,
args.verbose)
if args.memory:
test_ort_memory(device, onnx_model_path, args.batch_sizes[0], args.sequence_lengths[0],
args.global_lengths[0], args.test_times, num_threads)
else: # test latency
session = benchmark_helper.create_onnxruntime_session(onnx_model_path,
use_gpu=True,
enable_all_optimization=True,
num_threads=num_threads)
if session is None:
raise RuntimeError(f"Failed to create ORT sesssion from ONNX file {onnx_model_path}")
results += test_ort_latency(device, model, model_name, session, args.batch_sizes,
args.sequence_lengths, args.global_lengths, args.test_times,
num_threads, optimized, precision, args.validate_onnx,
args.disable_io_binding, args.verbose)
return results
def parse_arguments():
def parse_arguments(argv=None):
parser = argparse.ArgumentParser()
parser.add_argument("-m",
@ -231,9 +263,9 @@ def parse_arguments():
required=False,
nargs="+",
type=str,
default=["longformer-random-tiny"] if is_debug else ["longformer-base-4096"],
choices=list(MODELS.keys()),
help="Pre-trained models in the list: " + ", ".join(MODELS.keys()))
default=["longformer-base-4096"],
help="Checkpoint directory or pre-trained model names in the list: " +
", ".join(PRETRAINED_LONGFORMER_MODELS.keys()))
parser.add_argument("-e",
"--engines",
@ -253,21 +285,52 @@ def parse_arguments():
parser.add_argument("-b", "--batch_sizes", nargs="+", type=int, default=[1])
# If multiple of window size is used during exporting onnx model, there is no padding in ONNX model so you will need padding by yourself before running onnx model.
# In that case, you can only test sequence length that is multiple of window size (4 or 512 for these two models).
parser.add_argument("-s",
"--sequence_lengths",
# If --export_padding is not used in exporting onnx model, there is no padding in ONNX model so you will need padding inputs by yourself before running onnx model.
# In that case, you can only test sequence length that is multiple of attention window size.
parser.add_argument(
"-s",
"--sequence_lengths",
nargs="+",
type=int,
default=[512, 1024, 2048, 4096],
help=
"Sequence lengths. It could have multiple values in latency test. If --export_padding is not used in exporting onnx model, sequence length shall be multiple of window size."
)
parser.add_argument("--onnx_dir",
required=False,
type=str,
default=os.path.join('.', 'onnx_models'),
help="Directory to search onnx models.")
parser.add_argument("-g",
"--global_lengths",
nargs="+",
type=int,
default=[4] if is_debug else [512, 1024, 2048, 4096])
default=[0],
help="Number of global tokens. It could have multiple values in latency test.")
parser.add_argument("-g", "--global_lengths", nargs="+", type=int, default=[1] if is_debug else [8])
parser.add_argument("-n",
"--num_threads",
required=False,
nargs="+",
type=int,
default=[0],
help="Threads to use. It could have multiple values in latency test.")
parser.add_argument("-n", "--num_threads", required=False, nargs="+", type=int, default=[0], help="Threads to use")
parser.add_argument("-v",
"--validate_onnx",
required=False,
action="store_true",
help="Validate that ONNX model generates same output as PyTorch model.")
parser.add_argument("--verbose", required=False, action="store_true", help="Print more information")
parser.add_argument("--disable_io_binding", required=False, action="store_true", help="Do not use IO Binding.")
args = parser.parse_args()
parser.add_argument("--memory", required=False, action="store_true", help="Test memory usage instead of latency.")
parser.add_argument("--verbose", required=False, action="store_true", help="Print more information.")
args = parser.parse_args(argv)
return args
@ -327,16 +390,27 @@ def output_details(results, csv_filename):
print(f"Detail results are saved to csv file: {csv_filename}")
def main():
args = parse_arguments()
def main(args):
assert len(args.models) == 1, "run only one model at a time"
if args.memory:
if len(args.batch_sizes) > 1:
raise RuntimeError("For memory test, only one batch_size (-b) is allowed.")
if len(args.sequence_lengths) > 1:
raise RuntimeError("For memory test, only one sequence_length (-s) is allowed.")
if len(args.global_lengths) > 1:
raise RuntimeError("For memory test, only one global_length (-g) is allowed.")
if len(args.num_threads) > 1:
raise RuntimeError("For memory test, only one value of --num_threads is allowed.")
if not torch.cuda.is_available():
raise RuntimeError("Please install PyTorch with Cuda, and use a machine with GPU for testing gpu performance.")
torch.set_grad_enabled(False)
# set random seed manully to get deterministic results
#benchmark_helper.set_random_seed(123)
all_results = test_all(args)
time_stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
@ -348,4 +422,7 @@ def main():
if __name__ == "__main__":
main()
args = parse_arguments()
#args = parse_arguments("-e onnxruntime -t 1 -b 1 -s 4 -g 2 --onnx_dir . -t 1 -m longformer-random-tiny".split(' '))
main(args)

415
onnxruntime/python/tools/transformers/longformer/convert_longformer_to_onnx.py
View file

@ -1,9 +1,19 @@
# Before running this script, please run "python setup.py install" in ../torch_extensions to build longformer_attention.cpp
# under a python environment with PyTorch installed. Then you can update the path of longformer_attention.cpython-*.so
# and run this script in same environment.
# Tested in Ubuntu 18.04, python 3.6, PyTorch 1.7.1, transformers 4.3.0.
# GPU is not needed for this script. You can run it in CPU.
# For inference of the onnx model, you will need latest onnxruntime-gpu 1.7.0 or above.
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# This script converts Longformer model from huggingface transformers 4.0 or later to ONNX.
# Unlike normal ONNX model exporting, it will directly translate LongformerSelfAttention to the LongformerAttention operator in ONNX Runtime.
#
# Before running this script, please run "python setup.py install" in ../torch_extensions under Linux with PyTorch installed.
# Then you can update the path of longformer_attention.cpython-*.so and run this script in same environment.
#
# It is tested in Ubuntu 18.04, python 3.6, PyTorch 1.7.1, transformers 4.3.0 or 4.3.2.
# GPU is not needed for this script. You can run it in CPU. For --optimize_onnx, you can use either onnxruntime or onnxruntime-gpu package.
#
# For inference of the onnx model, you will need onnxruntime-gpu 1.7.0 or above.
import torch
import numpy as np
@ -13,6 +23,8 @@ from torch.onnx import register_custom_op_symbolic
from torch.onnx.symbolic_helper import parse_args
from packaging import version
from longformer_helper import LongformerHelper, PRETRAINED_LONGFORMER_MODELS
@parse_args('v', 'v', 'v', 'v', 'v', 'v', 'v', 'i', 'i')
def my_longformer_attention(g, input, weight, bias, mask, global_weight, global_bias, global_mask, num_heads, window):
@ -31,18 +43,10 @@ def my_longformer_attention(g, input, weight, bias, mask, global_weight, global_
# namespace is onnxruntime which is registered in longformer_attention.cpp
register_custom_op_symbolic('onnxruntime::LongformerAttention', my_longformer_attention, 9)
# TODO: update the path according to output of "python setup.py install" when your python version is not 3.6
# TODO: search the directory to find correct output filename of "python setup.py install" when python version is not 3.6
torch.ops.load_library(
r'../torch_extensions/build/lib.linux-x86_64-3.6/longformer_attention.cpython-36m-x86_64-linux-gnu.so')
# mapping from model name to pretrained model name
MODELS = {
"longformer-base-4096": "allenai/longformer-base-4096",
"longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging
}
is_debug = False
def parse_arguments():
parser = argparse.ArgumentParser()
@ -51,21 +55,24 @@ def parse_arguments():
"--model",
required=False,
type=str,
default="longformer-random-tiny" if is_debug else "longformer-base-4096",
choices=list(MODELS.keys()),
help="Pre-trained models in the list: " + ", ".join(MODELS.keys()))
default="longformer-base-4096",
help="Checkpoint directory or pre-trained model names in the list: " +
", ".join(PRETRAINED_LONGFORMER_MODELS.keys()))
# Sequence length shall choose properly.
# If multiple of windows size is used, there is no padding in ONNX model so you will need padding by yourself before running onnx model.
parser.add_argument("-s", "--sequence_length", type=int, default=4 if is_debug else 512)
parser.add_argument("-g", "--global_length", type=int, default=1 if is_debug else 8)
parser.add_argument(
'--export_padding',
required=False,
action='store_true',
help=
'Export padding logic to ONNX graph. If not enabled, user need pad input so that sequence length is multiple of window size.'
)
parser.set_defaults(export_padding=False)
parser.add_argument('-o',
'--optimize_onnx',
required=False,
action='store_true',
help='Use optimizer.py to optimize onnx model')
help='Use optimizer.py to optimize onnx model.')
parser.set_defaults(optimize_onnx=False)
parser.add_argument("-p",
@ -80,36 +87,77 @@ def parse_arguments():
return args
def get_dummy_inputs(sequence_length, num_global_tokens, device):
# Create a dummy input for ONNX export.
def get_dummy_inputs(config, export_padding, device):
# When sequence length is multiple of windows size, there is no padding logic in ONNX graph
sequence_length = config.attention_window[0] + 1 if export_padding else config.attention_window[0]
# Create dummy inputs
input_ids = torch.arange(sequence_length).unsqueeze(0).to(device)
attention_mask = torch.ones(input_ids.shape, dtype=torch.long,
device=input_ids.device) # TODO: use random word ID. #TODO: simulate masked word
global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=input_ids.device)
if num_global_tokens > 0:
global_token_index = list(range(num_global_tokens))
global_attention_mask[:, global_token_index] = 1
# TODO: support more inputs like token_type_ids, position_ids
attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)
attention_mask[:, sequence_length - 1] = 0 # last token is masked
global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device)
global_attention_mask[:, 0] = 1 # first token is global token
return input_ids, attention_mask, global_attention_mask
args = parse_arguments()
model_name = args.model
onnx_model_path = model_name + ".onnx"
from transformers import LongformerModel
model = LongformerModel.from_pretrained(MODELS[model_name]) # pretrained model name or directory
input_ids, attention_mask, global_attention_mask = get_dummy_inputs(sequence_length=args.sequence_length,
num_global_tokens=args.global_length,
device=torch.device('cpu'))
example_outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)
# A new function to replace LongformerSelfAttention.forward
#For transformer 4.3
# For transformers 4.0.0
def my_longformer_self_attention_forward_4(self,
hidden_states,
attention_mask=None,
is_index_masked=None,
is_index_global_attn=None,
is_global_attn=None):
global_mask = is_index_global_attn.int()
# The following check is based on the dummy inputs (only the first token is global).
assert len(global_mask.shape) == 2 and global_mask.shape[0] == 1 and global_mask.count_nonzero().item(
) == 1 and global_mask.tolist()[0][0] == 1
input_mask = is_index_masked.float()
input_mask = input_mask.masked_fill(is_index_masked, -10000.0)
# Yet another way to generate input_mask = torch.masked_fill(attention_mask, is_index_global_attn, 0.0)
# TODO: add postprocess of ONNX model to calculate based on graph input: input_mask = (attention_mask - 1) * 10000.0
# TODO: add postprocess of ONNX model to use graph input directly: glboal_mask = global_attention_mask
# The following check is based on the dummy inputs (only the last token is masked).
assert len(input_mask.shape) == 2 and input_mask.shape[0] == 1 and input_mask.count_nonzero().item(
) == 1 and input_mask.tolist()[0][-1] == -10000.0
weight = torch.stack(
(self.query.weight.transpose(0, 1), \
self.key.weight.transpose(0, 1), \
self.value.weight.transpose(0, 1)), dim=1)
weight = weight.reshape(self.embed_dim, 3 * self.embed_dim)
bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0)
bias = bias.reshape(3 * self.embed_dim)
global_weight = torch.stack((self.query_global.weight.transpose(0, 1), \
self.key_global.weight.transpose(0, 1), \
self.value_global.weight.transpose(0, 1)),
dim=1)
global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim)
global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0)
global_bias = global_bias.reshape(3 * self.embed_dim)
attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, input_mask, global_weight,
global_bias, global_mask, self.num_heads,
self.one_sided_attn_window_size)
assert attn_output.size() == hidden_states.size(), "Unexpected size"
outputs = (attn_output, )
return outputs
# For transformers 4.3.0
def my_longformer_self_attention_forward_4_3(self,
hidden_states,
attention_mask=None,
@ -117,190 +165,135 @@ def my_longformer_self_attention_forward_4_3(self,
is_index_global_attn=None,
is_global_attn=None,
output_attentions=False):
# TODO: move mask calculation to LongFormerModel class to avoid calculating it again and again in each layer.
global_mask = is_index_global_attn.int()
torch.masked_fill(attention_mask, is_index_global_attn, 0.0)
weight = torch.stack(
(self.query.weight.transpose(0, 1), self.key.weight.transpose(0, 1), self.value.weight.transpose(0, 1)), dim=1)
weight = weight.reshape(self.embed_dim, 3 * self.embed_dim)
bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0)
bias = bias.reshape(3 * self.embed_dim)
global_weight = torch.stack((self.query_global.weight.transpose(0, 1), self.key_global.weight.transpose(
0, 1), self.value_global.weight.transpose(0, 1)),
dim=1)
global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim)
global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0)
global_bias = global_bias.reshape(3 * self.embed_dim)
attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, attention_mask, global_weight,
global_bias, global_mask, self.num_heads,
self.one_sided_attn_window_size)
assert attn_output.size() == hidden_states.size(), "Unexpected size"
outputs = (attn_output, )
return outputs
assert output_attentions == False
return my_longformer_self_attention_forward_4(self, hidden_states, attention_mask, is_index_masked,
is_index_global_attn, is_global_attn)
#For transformers 4.0
def my_longformer_self_attention_forward_4(self,
hidden_states,
attention_mask=None,
is_index_masked=None,
is_index_global_attn=None,
is_global_attn=None):
# TODO: move mask calculation to LongFormerModel class to avoid calculating it again and again in each layer.
global_mask = is_index_global_attn.int()
torch.masked_fill(attention_mask, is_index_global_attn, 0.0)
weight = torch.stack(
(self.query.weight.transpose(0, 1), self.key.weight.transpose(0, 1), self.value.weight.transpose(0, 1)), dim=1)
weight = weight.reshape(self.embed_dim, 3 * self.embed_dim)
bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0)
bias = bias.reshape(3 * self.embed_dim)
global_weight = torch.stack((self.query_global.weight.transpose(0, 1), self.key_global.weight.transpose(
0, 1), self.value_global.weight.transpose(0, 1)),
dim=1)
global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim)
global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0)
global_bias = global_bias.reshape(3 * self.embed_dim)
attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, attention_mask, global_weight,
global_bias, global_mask, self.num_heads,
self.one_sided_attn_window_size)
assert attn_output.size() == hidden_states.size(), "Unexpected size"
outputs = (attn_output, )
return outputs
# For transformers 4.3.2
def my_longformer_self_attention_forward_4_3_2(self,
hidden_states,
attention_mask=None,
layer_head_mask=None,
is_index_masked=None,
is_index_global_attn=None,
is_global_attn=None,
output_attentions=False):
assert output_attentions == False
assert layer_head_mask is None
return my_longformer_self_attention_forward_4(self, hidden_states, attention_mask, is_index_masked,
is_index_global_attn, is_global_attn)
# For transformers 3.0
def my_longformer_attention_forward_3(self, hidden_states, attention_mask, output_attentions=False):
def export_longformer(model, onnx_model_path, export_padding):
input_ids, attention_mask, global_attention_mask = get_dummy_inputs(model.config,
export_padding,
device=torch.device('cpu'))
assert output_attentions is False
example_outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)
# TODO: attention_mask can directly be passed from inputs of model to avoid these processing.
attention_mask = attention_mask.squeeze(dim=2).squeeze(dim=1)
global_mask = (attention_mask > 0).int()
torch.masked_fill(attention_mask, attention_mask > 0, 0.0)
if version.parse(transformers.__version__) < version.parse("4.0.0"):
raise RuntimeError("This tool requires transformers 4.0.0 or later.")
weight = torch.stack(
(self.query.weight.transpose(0, 1), self.key.weight.transpose(0, 1), self.value.weight.transpose(0, 1)), dim=1)
weight = weight.reshape(self.embed_dim, 3 * self.embed_dim)
# Here we replace LongformerSelfAttention.forward using our implmentation for exporting ONNX model
from transformers import LongformerSelfAttention
import inspect
key = ' '.join(inspect.getfullargspec(LongformerSelfAttention.forward).args)
args_to_func = {
'self hidden_states attention_mask layer_head_mask is_index_masked is_index_global_attn is_global_attn output_attentions':
my_longformer_self_attention_forward_4_3_2,
'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn output_attentions':
my_longformer_self_attention_forward_4_3,
'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn':
my_longformer_self_attention_forward_4,
}
bias = torch.stack((self.query.bias, self.key.bias, self.value.bias), dim=0)
bias = bias.reshape(3 * self.embed_dim)
if key not in args_to_func:
print("Current arguments", inspect.getfullargspec(LongformerSelfAttention.forward).args)
raise RuntimeError(
"LongformerSelfAttention.forward arguments are different. Please install supported version (like transformers 4.3.0)."
)
global_weight = torch.stack((self.query_global.weight.transpose(0, 1), self.key_global.weight.transpose(
0, 1), self.value_global.weight.transpose(0, 1)),
dim=1)
global_weight = global_weight.reshape(self.embed_dim, 3 * self.embed_dim)
# Store for restoring later
original_forward = LongformerSelfAttention.forward
global_bias = torch.stack((self.query_global.bias, self.key_global.bias, self.value_global.bias), dim=0)
global_bias = global_bias.reshape(3 * self.embed_dim)
LongformerSelfAttention.forward = args_to_func[key]
attn_output = torch.ops.onnxruntime.LongformerAttention(hidden_states, weight, bias, attention_mask, global_weight,
global_bias, global_mask, self.num_heads,
self.one_sided_attn_window_size)
example_inputs = (input_ids, attention_mask, global_attention_mask)
assert attn_output.size() == hidden_states.size(), "Unexpected size"
Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
outputs = (attn_output, )
return outputs
# Here we replace LongformerSelfAttention.forward using our implmentation for exporting ONNX model
from transformers.modeling_longformer import LongformerSelfAttention
key = ' '.join(inspect.getfullargspec(LongformerSelfAttention.forward).args)
args_to_func = {
'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn output_attentions':
my_longformer_self_attention_forward_4_3,
'self hidden_states attention_mask is_index_masked is_index_global_attn is_global_attn':
my_longformer_self_attention_forward_4,
'self hidden_states attention_mask output_attentions': my_longformer_self_attention_forward_3,
}
if key not in args_to_func:
raise RuntimeError(
"LongformerSelfAttention.forward arguments are different. Please install supported version (like 4.3.0) of transformers package."
)
LongformerSelfAttention.forward = args_to_func[key]
"""
if version.parse(transformers.__version__) < version.parse("4.0.0"):
from transformers.modeling_longformer import LongformerSelfAttention
#original_forward = LongformerSelfAttention.forward
LongformerSelfAttention.forward = my_longformer_attention_forward_3
else:
from transformers.models.longformer.modeling_longformer import LongformerSelfAttention
#original_forward = LongformerSelfAttention.forward
LongformerSelfAttention.forward = my_longformer_self_attention_forward_4
"""
# TODO: support more inputs like (input_ids, attention_mask, global_attention_mask, token_type_ids, position_ids)
example_inputs = (input_ids, attention_mask, global_attention_mask)
torch.onnx.export(model,
example_inputs,
onnx_model_path,
opset_version=11,
example_outputs=example_outputs,
input_names=["input_ids", "attention_mask", "global_attention_mask"],
output_names=["last_state", "pooler"],
dynamic_axes={
'input_ids': {
0: 'batch_size',
1: 'sequence_length'
torch.onnx.export(model,
example_inputs,
onnx_model_path,
opset_version=11,
example_outputs=example_outputs,
input_names=["input_ids", "attention_mask", "global_attention_mask"],
output_names=["last_state", "pooler"],
dynamic_axes={
'input_ids': {
0: 'batch_size',
1: 'sequence_length'
},
'attention_mask': {
0: 'batch_size',
1: 'sequence_length'
},
'global_attention_mask': {
0: 'batch_size',
1: 'sequence_length'
},
'last_state': {
0: 'batch_size',
1: 'sequence_length'
},
'pooler': {
0: 'batch_size',
1: 'sequence_length'
}
},
'attention_mask': {
0: 'batch_size',
1: 'sequence_length'
},
'global_attention_mask': {
0: 'batch_size',
1: 'sequence_length'
},
'last_state': {
0: 'batch_size',
1: 'sequence_length'
},
'pooler': {
0: 'batch_size',
1: 'sequence_length'
}
},
custom_opsets={"com.microsoft": 1})
print(f"ONNX model exported to {onnx_model_path}")
custom_opsets={"com.microsoft": 1})
print(f"ONNX model exported to {onnx_model_path}")
if args.sequence_length % model.config.attention_window[0] == 0:
print(
f"*Attention*: You need input padding for inference: input sequece length shall be multiple of {model.config.attention_window[0]}. It is because the example input for export ONNX model does not need padding so padding logic is not in onnx model."
)
# Restore original implementaiton:
LongformerSelfAttention.forward = original_forward
# Restore Huggingface implementaiton like the following:
# LongformerSelfAttention.forward = original_forward
if args.precision != 'fp32' or args.optimize_onnx:
def optimize_longformer(onnx_model_path, fp32_model_path, fp16_model_path=None):
from onnx import load_model
from onnxruntime.transformers.onnx_model_bert import BertOnnxModel, BertOptimizationOptions
model = load_model(onnx_model_path, format=None, load_external_data=True)
optimization_options = BertOptimizationOptions('bert')
optimizer = BertOnnxModel(model, num_heads=16, hidden_size=768)
optimizer = BertOnnxModel(model, num_heads=16,
hidden_size=768) # paramters does not matter since attention fusion is not needed.
optimizer.optimize(optimization_options)
optimized_model_path = model_name + "_fp32.onnx"
optimizer.save_model_to_file(optimized_model_path)
print(f"optimized fp32 model saved to {optimized_model_path}")
if args.precision == 'fp16':
use_external_data_format = False
if fp32_model_path:
optimizer.save_model_to_file(fp32_model_path, use_external_data_format)
print(f"optimized fp32 model saved to {fp32_model_path}")
if fp16_model_path:
optimizer.convert_model_float32_to_float16(cast_input_output=True)
optimized_model_path = model_name + "_fp16.onnx"
optimizer.save_model_to_file(optimized_model_path)
print(f"optimized fp16 model saved to {optimized_model_path}")
optimizer.save_model_to_file(fp16_model_path, use_external_data_format)
print(f"optimized fp16 model saved to {fp16_model_path}")
def main(args):
model_name = args.model
onnx_model_path = model_name + ".onnx"
from transformers import LongformerModel
model = LongformerModel.from_pretrained(PRETRAINED_LONGFORMER_MODELS[model_name])
export_longformer(model, onnx_model_path, args.export_padding)
if args.optimize_onnx or args.precision != 'fp32':
fp32_model_path = model_name + "_fp32.onnx"
fp16_model_path = model_name + "_fp16.onnx" if args.precision == 'fp16' else None
optimize_longformer(onnx_model_path, fp32_model_path, fp16_model_path)
if __name__ == "__main__":
args = parse_arguments()
main(args)

76
onnxruntime/python/tools/transformers/longformer/longformer_helper.py Normal file
View file

@ -0,0 +1,76 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# This script helps creating dummy inputs for Longformer model.
import os
import logging
import torch
import onnx
import random
import numpy
import time
import re
from pathlib import Path
from typing import List, Dict, Tuple, Union
logger = logging.getLogger(__name__)
PRETRAINED_LONGFORMER_MODELS = {
"longformer-base-4096": "allenai/longformer-base-4096",
"longformer-large-4096": "allenai/longformer-large-4096",
"longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging
}
class LongformerInputs:
def __init__(self, input_ids, attention_mask, global_attention_mask):
self.input_ids: torch.LongTensor = input_ids
self.attention_mask: Union[torch.FloatTensor, torch.HalfTensor] = attention_mask
self.global_attention_mask: Union[torch.FloatTensor, torch.HalfTensor] = global_attention_mask
def to_list(self) -> List:
return [v for v in [self.input_ids, self.attention_mask, self.global_attention_mask] if v is not None]
def to_tuple(self) -> Tuple:
return tuple(v for v in self.to_list())
def get_ort_inputs(self) -> Dict:
return {
"input_ids": numpy.ascontiguousarray(self.input_ids.cpu().numpy()),
"attention_mask": numpy.ascontiguousarray(self.attention_mask.cpu().numpy()),
"global_attention_mask": numpy.ascontiguousarray(self.global_attention_mask.cpu().numpy()),
}
class LongformerHelper:
""" A helper class for Longformer model conversion, inference and verification.
"""
@staticmethod
def get_dummy_inputs(batch_size: int,
sequence_length: int,
num_global_tokens: int,
device: torch.device,
vocab_size: int = 100) -> LongformerInputs:
""" Create random inputs for Longformer model.
Returns torch tensors of input_ids, attention_mask and global_attention_mask tensors.
"""
input_ids = torch.randint(low=0,
high=vocab_size - 1,
size=(batch_size, sequence_length),
dtype=torch.long,
device=device)
attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)
global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device)
global_token_index = list(range(num_global_tokens))
global_attention_mask[:, global_token_index] = 1
return LongformerInputs(input_ids, attention_mask, global_attention_mask)
@staticmethod
def get_output_shapes(batch_size: int, sequence_length: int, hidden_size: int) -> Dict[str, List[int]]:
""" Returns a dictionary with output name as key, and shape as value.
"""
return {"last_state": [batch_size, sequence_length, hidden_size], "pooler": [batch_size, sequence_length]}

5
onnxruntime/python/tools/transformers/onnx_model.py
View file

@ -405,7 +405,10 @@ class OnnxModel:
self.model.opset_import[0].version = original_opset_version
def convert_model_float32_to_float16(self, cast_input_output=True):
""" Convert a graph to FLOAT16
"""Convert a graph to FLOAT16. By default, we will keep data types of inputs and outputs.
For decoder model with past_key_values, it is recommended to set cast_input_output=False for better performance.
Args:
cast_input_output (bool, optional): keep data type of inputs and outputs, and add Cast nodes to convert float32 inputs to float16, and float16 to float32 for outputs. Defaults to True.
"""
from packaging.version import Version
import onnxconverter_common as oc

79
onnxruntime/test/contrib_ops/longformer_attention_op_test.cc
View file

@ -5,6 +5,8 @@
#include "test/common/tensor_op_test_utils.h"
#include "test/common/cuda_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
#include "test/util/include/scoped_env_vars.h"
#include "contrib_ops/cpu/bert/longformer_attention_base.h"
namespace onnxruntime {
namespace test {
@ -148,16 +150,33 @@ static void GetTinyLongformerData(
static void RunTinyLongformerBatch1(
std::vector<float>& mask_data,
std::vector<int>& global_data,
std::vector<float>& input_data,
std::vector<float>& output_data,
bool use_float16,
bool window_cover_whole_sequence = false) {
bool use_float16) {
int batch_size = 1;
int one_sided_attention_window_size = 2;
int hidden_size = 8;
int number_of_heads = 2;
std::vector<float> weight_data;
std::vector<float> bias_data;
std::vector<float> global_weight_data;
std::vector<float> global_bias_data;
GetTinyLongformerData(weight_data, bias_data, global_weight_data, global_bias_data);
int sequence_length = static_cast<int>(mask_data.size()) / batch_size;
RunAttentionTest(input_data, weight_data, bias_data, mask_data, global_weight_data, global_bias_data, global_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, one_sided_attention_window_size, use_float16);
}
static void RunTinyLongformerBatch1(
std::vector<float>& mask_data,
std::vector<int>& global_data,
std::vector<float>& output_data,
bool use_float16,
bool window_cover_whole_sequence = false) {
// Total windows size 4 will cover the whole sequence length 4
int sequence_length = window_cover_whole_sequence ? 4 : 8;
std::vector<float> input_data;
if (window_cover_whole_sequence) {
input_data = {
@ -176,15 +195,7 @@ static void RunTinyLongformerBatch1(
-1.0536f, -0.0425f, -1.1194f, -0.6423f, 2.1825f, 0.2547f, 0.6015f, -0.1809f,
0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f};
}
std::vector<float> weight_data;
std::vector<float> bias_data;
std::vector<float> global_weight_data;
std::vector<float> global_bias_data;
GetTinyLongformerData(weight_data, bias_data, global_weight_data, global_bias_data);
RunAttentionTest(input_data, weight_data, bias_data, mask_data, global_weight_data, global_bias_data, global_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, one_sided_attention_window_size, use_float16);
return RunTinyLongformerBatch1(mask_data, global_data, input_data, output_data, use_float16);
}
TEST(LongformerAttentionTest, LongformerAttention_NoGlobal) {
@ -238,6 +249,50 @@ TEST(LongformerAttentionTest, LongformerAttention_GlobalStart) {
RunTinyLongformerBatch1(mask_data, global_data, output_data, false);
}
/*
* This following case is generated from the first self-attention of a tiny longformer model like the following:
import torch
from transformers import AutoModel
model = AutoModel.from_pretrained("patrickvonplaten/longformer-random-tiny")
input_ids = torch.LongTensor([[97, 24, 71, 5, 8]]))
attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device)
global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=input_ids.device)
global_attention_mask[:, [0, 1,]] = 1
outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)
*/
TEST(LongformerAttentionTest, LongformerAttention_UseCompactMemory) {
std::vector<float> mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f, -10000.0f, -10000.0f};
std::vector<int> global_data = {1, 1, 0, 0, 0, 0, 0, 0};
std::vector<float> input_data = {
-1.6886f, -0.0712f, 0.7017f, -1.5513f, 0.1059f, 0.8119f, 1.1139f, 0.5776f,
-1.0874f, 0.2870f, -0.3767f, -1.8203f, 0.0116f, 1.4285f, 0.6643f, 0.8931f,
-1.0425f, 1.0148f, -0.6387f, -1.6496f, -0.3027f, 0.3839f, 1.0230f, 1.2119f,
0.3576f, -1.5848f, 0.1713f, -0.9284f, 1.7172f, -0.8574f, 0.8390f, 0.2855f,
-1.7443f, 0.4251f, 1.6560f, -0.5054f, -0.8734f, 0.4769f, 0.8434f, -0.2781f,
0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f,
0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f,
0.5219f, 0.1777f, 0.7090f, -2.1933f, 0.5258f, -0.0639f, -0.8511f, 1.1738f};
std::vector<float> output_data = {
-0.0330f, 0.0517f, -0.0292f, 0.0194f, -0.0154f, 0.0012f, -0.0041f, -0.0518f,
-0.0331f, 0.0517f, -0.0291f, 0.0194f, -0.0154f, 0.0011f, -0.0041f, -0.0517f,
0.0421f, 0.0627f, 0.0445f, 0.0328f, 0.0196f, -0.0712f, -0.0233f, -0.0875f,
0.0419f, 0.0626f, 0.0447f, 0.0328f, 0.0196f, -0.0712f, -0.0233f, -0.0875f,
0.0421f, 0.0627f, 0.0445f, 0.0328f, 0.0196f, -0.0712f, -0.0233f, -0.0875f,
0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f,
0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f,
0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f, 0.0000f};
ScopedEnvironmentVariables scoped_env_vars{
EnvVarMap{
{onnxruntime::contrib::longformer::kUseCompactMemory, "1"},
}};
RunTinyLongformerBatch1(mask_data, global_data, input_data, output_data, false);
}
TEST(LongformerAttentionTest, LongformerAttention_Float16) {
std::vector<float> mask_data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, -10000.0f};

2
tools/ci_build/amd_hipify.py
View file

@ -26,6 +26,8 @@ contrib_ops_excluded_files = [
'bert/layer_norm.cuh',
'bert/longformer_attention.cc',
'bert/longformer_attention.h',
'bert/longformer_attention_softmax.cu',
'bert/longformer_attention_softmax.h',
'bert/longformer_attention_impl.cu',
'bert/longformer_attention_impl.h',
'bert/longformer_global_impl.cu',