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

Add past state support in Attention Op for GPT-2 (#4107)

Browse source 
Update Attention op to allow past state input and output.
Add fusion script and tests
This commit is contained in:
Tianlei Wu 2020-06-11 14:19:55 -07:00 committed by GitHub
parent e6ccb1ac28
commit 2605faef88
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
31 changed files with 1733 additions and 404 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

89
onnxruntime/contrib_ops/cpu/bert/attention.cc
View file

@ -27,18 +27,21 @@ AttentionBase::AttentionBase(const OpKernelInfo& info) {
int64_t num_heads = 0;
ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0);
num_heads_ = static_cast<int>(num_heads);
is_unidirectional_ = info.GetAttrOrDefault<int64_t>("unidirectional", 0) == 1;
}
Status AttentionBase::CheckInputs(const Tensor* input,
const Tensor* weights,
const Tensor* bias,
const Tensor* mask_index) const {
// Input and output shapes:
const Tensor* mask_index,
const Tensor* past) const {
// Input shapes:
// input : (batch_size, sequence_length, hidden_size)
// weights : (hidden_size, 3 * hidden_size)
// bias : (3 * hidden_size)
// mask_index : (batch_size) if presented
// past : (2, batch_size, num_heads, past_sequence_length, head_size)
const auto dims = input->Shape().GetDims();
if (dims.size() != 3) {
@ -91,9 +94,59 @@ Status AttentionBase::CheckInputs(const Tensor* input,
}
}
if (past != nullptr) { // past is optional
if (!is_unidirectional_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 (past) is only allowed for unidirectional");
}
const auto past_dims = past->Shape().GetDims();
if (past_dims.size() != 5) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 is expected to have 5 dimension, got ",
past_dims.size());
}
if (static_cast<int>(past_dims[0]) != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 0 shall have length of 2");
}
if (static_cast<int>(past_dims[1]) != batch_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 1 shall have same length as dimension 0 of input 0");
}
if (static_cast<int>(past_dims[2]) != num_heads_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 2 shall have length of num_heads", num_heads_);
}
if (static_cast<int>(past_dims[4]) != hidden_size / num_heads_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 2 shall have length of ", hidden_size / num_heads_);
}
}
return Status::OK();
}
Tensor* AttentionBase::GetPresent(OpKernelContext* context,
const Tensor* past,
int batch_size,
int head_size,
int sequence_length,
int& past_sequence_length) const {
// Input and output shapes:
// past : (2, batch_size, num_heads, past_sequence_length, head_size)
// present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)
std::vector<int64_t> present_dims{2, batch_size, num_heads_, sequence_length, head_size};
if (nullptr != past) {
const auto past_dims = past->Shape().GetDims();
past_sequence_length = static_cast<int>(past_dims[3]);
present_dims[3] += past_dims[3];
}
TensorShape present_shape(present_dims);
Tensor* present = context->Output(1, present_shape);
if (nullptr != past && nullptr == present) {
ORT_THROW("Expect to have present state output when past state input is given");
}
return present;
}
template <typename T>
Attention<T>::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionBase(info) {
}
@ -104,7 +157,9 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
const Tensor* weights = context->Input<Tensor>(1);
const Tensor* bias = context->Input<Tensor>(2);
const Tensor* mask_index = context->Input<Tensor>(3);
ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index));
const Tensor* past = context->Input<Tensor>(4);
ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index, past));
const auto dims = input->Shape().GetDims();
const int batch_size = static_cast<int>(dims[0]);
@ -115,6 +170,12 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
TensorShape output_shape(dims);
Tensor* output = context->Output(0, output_shape);
int past_sequence_length = 0;
Tensor* present = GetPresent(context, past, batch_size, head_size, sequence_length, past_sequence_length);
// Total sequence length including that of past state: S* = S' + S
const int all_sequence_length = past_sequence_length + sequence_length;
constexpr size_t element_size = sizeof(T);
AllocatorPtr allocator;
@ -182,18 +243,18 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
}
// STEP.2: compute the attention score. It does 2 things:
// I. attention_probs(B, N, S, S) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S, H -> B, N, H, S) +
// 1 x mask_data(B, N, S, S)
// II.attention_probs(B, N, S, S) = Softmax(attention_probs)
size_t attention_probs_bytes = SafeInt<size_t>(batch_size) * num_heads_ * sequence_length * sequence_length * element_size;
// I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) +
// 1 x mask_data(B, N, S, S*)
// II.attention_probs(B, N, S, S*) = Softmax(attention_probs)
size_t attention_probs_bytes = SafeInt<size_t>(batch_size) * num_heads_ * sequence_length * all_sequence_length * element_size;
auto attention_probs = allocator->Alloc(attention_probs_bytes);
BufferUniquePtr scratch_buffer(attention_probs, BufferDeleter(allocator));
size_t mask_data_bytes = 0;
if (mask_index != nullptr) {
mask_data_bytes = SafeInt<size_t>(batch_size) * sequence_length * sequence_length * element_size;
mask_data_bytes = SafeInt<size_t>(batch_size) * sequence_length * all_sequence_length * element_size;
} else if (is_unidirectional_) {
mask_data_bytes = SafeInt<size_t>(sequence_length) * sequence_length * element_size;
mask_data_bytes = SafeInt<size_t>(sequence_length) * all_sequence_length * element_size;
}
void* mask_data = nullptr;
@ -204,17 +265,21 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
BufferUniquePtr mask_data_buffer(mask_data, BufferDeleter(allocator));
const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data<int32_t>() : nullptr;
const T* past_data = past != nullptr ? past->template Data<T>() : nullptr;
T* present_data = present != nullptr ? present->template MutableData<T>() : nullptr;
ComputeAttentionProbs<T>(static_cast<T*>(attention_probs), Q, K, mask_index_data, static_cast<T*>(mask_data),
batch_size, sequence_length, head_size, num_heads_, is_unidirectional_, tp);
batch_size, sequence_length, past_sequence_length, head_size, num_heads_, is_unidirectional_,
past_data, present_data, tp);
// STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S) x V(B, N, S, H)
// STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S*) x V(B, N, S*, H)
auto out_tmp_data =
allocator->Alloc(SafeInt<size_t>(batch_size) * num_heads_ * sequence_length * head_size * element_size);
BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator));
ComputeVxAttentionScore(output->template MutableData<T>(), static_cast<T*>(out_tmp_data), static_cast<T*>(attention_probs), V,
batch_size, sequence_length, head_size, num_heads_, hidden_size, tp);
batch_size, sequence_length, past_sequence_length, head_size, num_heads_, hidden_size,
past_data, present_data, tp);
return Status::OK();
}

14
onnxruntime/contrib_ops/cpu/bert/attention.h
View file

@ -15,10 +15,18 @@ class AttentionBase {
Status CheckInputs(const Tensor* input,
const Tensor* weights,
const Tensor* bias,
const Tensor* mask_index) const;
const Tensor* mask_index,
const Tensor* past) const;
int num_heads_; // number of attention heads
bool is_unidirectional_; // whether every token can only attend to previous tokens.
Tensor* GetPresent(OpKernelContext* context,
const Tensor* past,
int batch_size,
int head_size,
int sequence_length,
int& past_sequence_length) const;
int num_heads_; // number of attention heads
bool is_unidirectional_; // whether every token can only attend to previous tokens.
};
template <typename T>

183
onnxruntime/contrib_ops/cpu/bert/attention_helper.h
View file

@ -59,111 +59,176 @@ inline void ComputeAttentionSoftmaxInplace(float* score, int N, int D, ThreadPoo
MlasComputeSoftmax(score, score, N, D, false, tp);
}
// Helper function to compute the attention probs. It does 2 things:
// I. attention_probs(B, N, S, S) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S, H -> B, N, H, S) +
// 1 x mask_data(B, N, S, S)
// II.attention_probs(B, N, S, S) = Softmax(attention_probs)
template <typename T>
void ComputeAttentionProbs(T* attention_probs, // output buffer for the attention probs. Its size is B*N*S*S
const T* Q, // Q data. Its size is B*N*S*H
const T* K, // k data. Its size is B*N*S*H
void PrepareMask(const int32_t* mask_index,
T* mask_data,
bool is_unidirectional,
int batch_size,
int sequence_length,
int past_sequence_length) {
const int all_sequence_length = past_sequence_length + sequence_length;
T* p_mask = mask_data;
if (is_unidirectional) {
// unidirectional mask has shape SxS*
for (int s_i = 0; s_i < sequence_length - 1; s_i++) {
for (int m_i = past_sequence_length + s_i + 1; m_i < all_sequence_length; m_i++) {
p_mask[s_i * all_sequence_length + m_i] = static_cast<T>(-10000.0);
}
}
return;
}
ORT_ENFORCE(mask_index, "mask index should not be null.");
for (int b_i = 0; b_i < batch_size; b_i++) {
// TODO: mask_index can be used in softmax to save some calculation.
// Convert mask_index to mask (-10000 means out of range, which will be 0 after softmax): B => BxS*
int valid_length = mask_index[b_i];
for (int m_i = valid_length; m_i < all_sequence_length; m_i++) {
p_mask[m_i] = static_cast<T>(-10000.0);
}
// Broadcast mask from BxS* to BxSxS*
for (int s_i = 1; s_i < sequence_length; s_i++) {
memcpy(p_mask + s_i * all_sequence_length, p_mask, all_sequence_length * sizeof(T));
}
p_mask += sequence_length * sequence_length;
}
}
// Concatenate a past state chunk S'xH with input state chunk SxH into present state chunk S*xH
// Returns a pointer to the start of present state chunk.
template<typename T>
T* ConcatStateChunk(const T* past, const T* chunk, T* present, size_t past_chunk_length, size_t present_chunk_length, std::ptrdiff_t i) {
T* start = present + i * present_chunk_length;
T* p = start;
if (nullptr != past) {
const T* src_past = past + i * past_chunk_length;
memcpy(p, src_past, past_chunk_length * sizeof(T));
p += past_chunk_length;
}
memcpy(p, chunk, (present_chunk_length - past_chunk_length) * sizeof(T));
return start;
}
// Helper function to compute the attention probs. It does 2 things:
// I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) +
// 1 x mask_data(B, N, S, S*)
// II.attention_probs(B, N, S, S*) = Softmax(attention_probs)
template <typename T>
void ComputeAttentionProbs(T* attention_probs, // output buffer for the attention probs. Its size is BxNxSxS
const T* Q, // Q data. Its size is BxNxSxH
const T* K, // k data. Its size is BxNxSxH
const int32_t* mask_index, // mask index. nullptr if no mask or its size is B
T* mask_data, // buffer for mask data. Its size is: S*S if is_unidirectiona; B*S*S if mask_index; null otherwise
T* mask_data, // buffer for mask data. Its size is: SxS* if is_unidirectional; BxSxS* if mask_index; null otherwise
int batch_size, // batch size of self-attention
int sequence_length, // sequence length of self-attention
int past_sequence_length, // sequence length of past state
int head_size, // head size of self-attention
int num_heads, // number of heads of self-attention
bool is_unidirectional, // indicate if it is unidrectional.
const T* past, // past state
T* present, // present state
ThreadPool* tp) {
const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S
const size_t past_chunk_length = static_cast<size_t>(past_sequence_length * head_size); // S' x H
const size_t input_chunk_length = static_cast<size_t>(sequence_length * head_size); // S x H
const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H
{
if (mask_data != nullptr) {
if (is_unidirectional) {
for (int s_i = 0; s_i < sequence_length - 1; s_i++) {
for (int m_i = s_i + 1; m_i < sequence_length; m_i++) {
mask_data[s_i * sequence_length + m_i] = static_cast<T>(-10000.0);
}
}
} else {
ORT_ENFORCE(mask_index, "mask index should not be null.");
T* p_mask = mask_data;
for (int b_i = 0; b_i < batch_size; b_i++) {
// TODO: mask_index can be used in softmax to save some calculation.
// Convert mask_index to mask (-10000 means out of range, which will be 0 after softmax): B => BxS
int valid_length = mask_index[b_i];
for (int m_i = valid_length; m_i < sequence_length; m_i++) {
p_mask[m_i] = static_cast<T>(-10000.0);
}
// Broadcast mask from BxS to BxSxS
for (int s_i = 1; s_i < sequence_length; s_i++) {
memcpy(p_mask + s_i * sequence_length, p_mask, sequence_length * sizeof(T));
}
p_mask += sequence_length * sequence_length;
}
}
PrepareMask(mask_index, mask_data, is_unidirectional, batch_size, sequence_length, past_sequence_length);
} else { // no any mask
memset(attention_probs, 0, batch_size * num_heads * sequence_length * sequence_length * sizeof(T));
memset(attention_probs, 0, batch_size * num_heads * sequence_length * all_sequence_length * sizeof(T));
}
const int loop_len = batch_size * num_heads;
const float alpha = 1.0f / sqrt(static_cast<float>(head_size));
// The cost of Gemm
const double cost =
static_cast<double>(head_size) * static_cast<double>(sequence_length) * static_cast<double>(sequence_length);
const double cost = static_cast<double>(head_size * sequence_length * all_sequence_length);
ThreadPool::TryParallelFor(tp, loop_len, cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) {
for (std::ptrdiff_t i = begin; i != end; ++i) {
const std::ptrdiff_t batch_index = i / num_heads;
// broadcast mask data: SxS or (Bx)SxS -> (BxNx)SxS
// broadcast mask data: SxS* or (Bx)SxS* -> (BxNx)SxS*
if (mask_data != nullptr) {
const T* broadcast_data_src = is_unidirectional ? reinterpret_cast<T*>(mask_data) : reinterpret_cast<T*>(mask_data) + batch_index * sequence_length * sequence_length;
T* broadcast_data_dest = reinterpret_cast<T*>(attention_probs) + sequence_length * sequence_length * i;
memcpy(broadcast_data_dest, broadcast_data_src, sequence_length * sequence_length * sizeof(T));
const T* broadcast_data_src = is_unidirectional ? reinterpret_cast<T*>(mask_data) : reinterpret_cast<T*>(mask_data) + batch_index * sequence_length * all_sequence_length;
T* broadcast_data_dest = reinterpret_cast<T*>(attention_probs) + sequence_length * all_sequence_length * i;
memcpy(broadcast_data_dest, broadcast_data_src, sequence_length * all_sequence_length * sizeof(T));
}
const T* k = K + input_chunk_length * i;
if (nullptr != present) {
// concatenate past_K and K : (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH
k = ConcatStateChunk(past, k, present, past_chunk_length, present_chunk_length, i);
}
// gemm
// original transposed iteration
// A: Q (BxNxSxH) (B.N.)S x H S x H
// B: K' (BxNxSxH) (B.N.)H x S H x S
// C: attention_probs (BxNxSxS) (B.N.)S x S S x S
math::Gemm<T, ThreadPool>(CblasNoTrans, CblasTrans, sequence_length, sequence_length, head_size, alpha,
Q + sequence_length * head_size * i, K + sequence_length * head_size * i, 1.0,
reinterpret_cast<T*>(attention_probs) + sequence_length * sequence_length * i, nullptr);
// original transposed each iteration
// A: Q (B x N x) S x H (B x N x) S x H S x H
// B: K' (B x N x) S* x H (B x N x) H x S* H x S*
// C: attention_probs (B x N x) S x S* (B x N x) S x S* S x S*
math::Gemm<T, ThreadPool>(CblasNoTrans, CblasTrans, sequence_length, all_sequence_length, head_size, alpha,
Q + input_chunk_length * i, k, 1.0,
reinterpret_cast<T*>(attention_probs) + sequence_length * all_sequence_length * i, nullptr);
}
});
}
// attention_probs(B, N, S, S) = Softmax(attention_probs)
// attention_probs(B, N, S, S*) = Softmax(attention_probs)
{
const int N = batch_size * num_heads * sequence_length;
const int D = sequence_length;
const int D = all_sequence_length;
ComputeAttentionSoftmaxInplace(attention_probs, N, D, tp);
}
}
template <typename T>
void ComputeVxAttentionScore(T* output, // buffer for the result with size B*S*N*H
T* tmp_buffer, // buffer for temp use with size is B*N*S*H
const T* attention_probs, // Attention probs with size B*N*S*S
const T* V, // V valuee with size B*N*S*H
void ComputeVxAttentionScore(T* output, // buffer for the result with size BxSxNxH
T* tmp_buffer, // buffer for temp use with size is BxNxSxH
const T* attention_probs, // Attention probs with size BxNxSxS*
const T* V, // V value with size BxNxSxH
int batch_size, // batch size
int sequence_length, // sequence length
int past_sequence_length, // sequence length in past state
int head_size, // head size
int num_heads, // number of heads
int hidden_size, // hidden size
const T* past, // past state
T* present, // present state
ThreadPool* tp) {
const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S
const size_t past_chunk_length = static_cast<size_t>(past_sequence_length * head_size); // S' x H
const size_t input_chunk_length = static_cast<size_t>(sequence_length * head_size); // S x H
const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H
// Move the pointer of past and present to start of v values.
if (nullptr != past) {
past += batch_size * num_heads * past_sequence_length * head_size;
}
if (nullptr != present) {
present += batch_size * num_heads * all_sequence_length * head_size;
}
const double cost =
static_cast<double>(sequence_length) * static_cast<double>(head_size) * static_cast<double>(sequence_length);
ThreadPool::TryParallelFor(tp, batch_size * num_heads, cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) {
const int sequence_length_mul_head_size = sequence_length * head_size;
for (std::ptrdiff_t i = begin; i != end; ++i) {
T* current_tmp_data = tmp_buffer + sequence_length_mul_head_size * i;
math::MatMul<T>(sequence_length, head_size, sequence_length,
attention_probs + sequence_length * sequence_length * i,
V + sequence_length_mul_head_size * i, current_tmp_data, nullptr);
const T* v = V + input_chunk_length * i;
if (nullptr != present) {
// concatenate past_V and V: (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH
v = ConcatStateChunk(past, v, present, past_chunk_length, present_chunk_length, i);
}
T* current_tmp_data = reinterpret_cast<T*>(tmp_buffer) + input_chunk_length * i;
math::MatMul<T>(sequence_length, head_size, all_sequence_length,
attention_probs + sequence_length * all_sequence_length * i,
v, current_tmp_data, nullptr);
// transpose: out(B, S, N, H) = transpose out_tmp(B, N, S, H)
const int batch_index = static_cast<int>(i / num_heads);

10
onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc
View file

@ -59,7 +59,7 @@ Status QAttention<T, QInput, QWeight>::Compute(OpKernelContext* context) const {
const Tensor* i_zp_tensor = context->Input<Tensor>(6);
const Tensor* w_zp_tensor = context->Input<Tensor>(7);
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index));
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr));
ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor),
"input scale must be a scalar or 1D tensor of size 1");
@ -193,8 +193,12 @@ Status QAttention<T, QInput, QWeight>::Compute(OpKernelContext* context) const {
const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data<int32_t>() : nullptr;
int past_sequence_length = 0;
const T* past_data = nullptr;
T* present_data = nullptr;
ComputeAttentionProbs<T>(static_cast<T*>(attention_probs), Q, K, mask_index_data, static_cast<T*>(mask_data),
batch_size, sequence_length, head_size, num_heads_, is_unidirectional_, tp);
batch_size, sequence_length, past_sequence_length, head_size, num_heads_, is_unidirectional_,
past_data, present_data, tp);
// STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S) x V(B, N, S, H)
auto out_tmp_data =
@ -202,7 +206,7 @@ Status QAttention<T, QInput, QWeight>::Compute(OpKernelContext* context) const {
BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator));
ComputeVxAttentionScore(output->template MutableData<T>(), static_cast<T*>(out_tmp_data), static_cast<T*>(attention_probs), V,
batch_size, sequence_length, head_size, num_heads_, hidden_size, tp);
batch_size, sequence_length, past_sequence_length, head_size, num_heads_, hidden_size, past_data, present_data, tp);
return Status::OK();
}

25
onnxruntime/contrib_ops/cuda/bert/attention.cc
View file

@ -34,18 +34,16 @@ Attention<T>::Attention(const OpKernelInfo& info) : CudaKernel(info), AttentionB
template <typename T>
Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
// Input and output shapes:
// Input 0 - input : (batch_size, sequence_length, hidden_size)
// Input 1 - weights : (hidden_size, 3 * hidden_size)
// Input 2 - bias : (3 * hidden_size)
// Input 3 - mask_index : (batch_size) if presented
// Output : (batch_size, sequence_length, hidden_size)
const Tensor* input = context->Input<Tensor>(0);
const Tensor* weights = context->Input<Tensor>(1);
const Tensor* bias = context->Input<Tensor>(2);
const Tensor* mask_index = context->Input<Tensor>(3);
ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index));
const Tensor* past = context->Input<Tensor>(4);
ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index, past));
// Input and output shapes:
// Input 0 - input : (batch_size, sequence_length, hidden_size)
// Output 0 - output : (batch_size, sequence_length, hidden_size)
const auto dims = input->Shape().GetDims();
int batch_size = static_cast<int>(dims[0]);
int sequence_length = static_cast<int>(dims[1]);
@ -55,8 +53,11 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
TensorShape output_shape(dims);
Tensor* output = context->Output(0, output_shape);
int past_sequence_length = 0;
Tensor* present = GetPresent(context, past, batch_size, head_size, sequence_length, past_sequence_length);
cublasHandle_t cublas = CublasHandle();
const size_t element_size = sizeof(T);
constexpr size_t element_size = sizeof(T);
// Use GEMM for fully connection.
int m = batch_size * sequence_length;
@ -84,7 +85,7 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
reinterpret_cast<const CudaT*>(input->template Data<T>()), k,
&one, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop));
size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length);
size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, past_sequence_length);
auto temp_buffer = GetScratchBuffer<void>(workSpaceSize);
if (!LaunchAttentionKernel(
reinterpret_cast<const CudaT*>(gemm_buffer.get()),
@ -97,7 +98,11 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
temp_buffer.get(),
cublas,
element_size,
is_unidirectional_)) {
is_unidirectional_,
past_sequence_length,
nullptr == past ? nullptr : past->template Data<T>(),
nullptr == present ? nullptr : present->template MutableData<T>()
)) {
// Get last error to reset it to cudaSuccess.
CUDA_CALL(cudaGetLastError());
return Status(common::ONNXRUNTIME, common::FAIL);

188
onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
View file

@ -40,8 +40,8 @@ static size_t AlignTo(size_t a, size_t b) {
return CeilDiv(a, b) * b;
}
size_t ScratchSize(size_t element_size, int batch_size, int num_heads, int sequence_length) {
const size_t len = batch_size * num_heads * sequence_length * sequence_length;
size_t ScratchSize(size_t element_size, int batch_size, int num_heads, int sequence_length, int past_sequence_length) {
const size_t len = batch_size * num_heads * sequence_length * (sequence_length + past_sequence_length);
const size_t bytes = len * element_size;
const size_t alignment = 256;
@ -49,13 +49,19 @@ size_t ScratchSize(size_t element_size, int batch_size, int num_heads, int seque
return bytesAligned;
}
size_t GetAttentionWorkspaceSize(size_t element_size, int batch_size, int num_heads, int head_size, int sequence_length) {
size_t GetAttentionWorkspaceSize(
size_t element_size,
int batch_size,
int num_heads,
int head_size,
int sequence_length,
int past_sequence_length) {
size_t qkv_size = 3 * batch_size * sequence_length * num_heads * head_size * element_size;
return qkv_size + 2 * ScratchSize(element_size, batch_size, num_heads, sequence_length);
return qkv_size + 2 * ScratchSize(element_size, batch_size, num_heads, sequence_length, past_sequence_length);
}
template <typename T, unsigned TPB>
__device__ inline void Softmax(const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) {
__device__ inline void Softmax(const int past_sequence_length, const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
@ -64,13 +70,14 @@ __device__ inline void Softmax(const int sequence_length, const int valid_length
float thread_data_max(-CUDART_INF_F);
const int num_valid = is_unidirectional ? (blockIdx.x % sequence_length) + 1 : valid_length;
const int num_valid = is_unidirectional ? past_sequence_length + (blockIdx.x % sequence_length) + 1 : valid_length;
// e^x is represented as infinity if x is large enough, like 100.f.
// Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough.
// a math transform as below is leveraged to get a stable softmax:
// e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max))
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length;
const int all_sequence_length = past_sequence_length + sequence_length;
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
for (int i = threadIdx.x; i < num_valid; i += TPB) {
const int index = offset + i;
if (thread_data_max < float(input[index])) {
@ -99,7 +106,7 @@ __device__ inline void Softmax(const int sequence_length, const int valid_length
}
__syncthreads();
for (int i = threadIdx.x; i < sequence_length; i += TPB) {
for (int i = threadIdx.x; i < all_sequence_length; i += TPB) {
const int index = offset + i;
const float val = (i < num_valid) ? expf(float(input[index]) - max_block) * sum_reverse_block : 0.f;
output[index] = T(val);
@ -107,17 +114,19 @@ __device__ inline void Softmax(const int sequence_length, const int valid_length
}
template <typename T, unsigned TPB>
__device__ inline void SoftmaxSmall(const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) {
__device__ inline void SoftmaxSmall(const int past_sequence_length, const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
__shared__ float sum_reverse_block;
__shared__ float max_block;
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length;
// Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S;
const int all_sequence_length = past_sequence_length + sequence_length;
const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length;
const int index = offset + threadIdx.x;
const int num_valid = is_unidirectional ? (blockIdx.x % sequence_length) + 1 : valid_length;
const int num_valid = is_unidirectional ? past_sequence_length + (blockIdx.x % sequence_length) + 1 : valid_length;
// e^x is represented as infinity if x is large enough, like 100.f.
// Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough.
@ -146,43 +155,44 @@ __device__ inline void SoftmaxSmall(const int sequence_length, const int valid_l
// Store max value
if (threadIdx.x == 0) {
sum_reverse_block = (1.f) / sum;
sum_reverse_block = (num_valid == 0) ? 0.f : (1.f) / sum;
}
__syncthreads();
if (threadIdx.x < sequence_length) {
// threadIdx.x might be larger than all_sequence_length due to alignment to 32x.
if (threadIdx.x < all_sequence_length) {
// this will be 0 for threadIdx.x >= num_valid
output[index] = T(thread_data_exp * sum_reverse_block);
}
}
template <typename T, unsigned TPB>
__global__ void SoftmaxKernelSmall(const int sequence_length, const T* input, T* output, bool is_unidirectional) {
SoftmaxSmall<T, TPB>(sequence_length, sequence_length, input, output, is_unidirectional);
__global__ void SoftmaxKernelSmall(const int past_sequence_length, const int sequence_length, const T* input, T* output, bool is_unidirectional) {
SoftmaxSmall<T, TPB>(past_sequence_length, sequence_length, sequence_length, input, output, is_unidirectional);
}
template <typename T, unsigned TPB>
__global__ void SoftmaxKernel(const int sequence_length, const T* input, T* output, bool is_unidirectional) {
Softmax<T, TPB>(sequence_length, sequence_length, input, output, is_unidirectional);
__global__ void SoftmaxKernel(const int past_sequence_length, const int sequence_length, const T* input, T* output, bool is_unidirectional) {
Softmax<T, TPB>(past_sequence_length, sequence_length, sequence_length, input, output, is_unidirectional);
}
template <typename T>
bool ComputeSoftmax(
cudaStream_t stream, const int sequence_length, const int batch_size, const int num_heads,
cudaStream_t stream, const int past_sequence_length, const int sequence_length, const int batch_size, const int num_heads,
const T* input, T* output, bool is_unidirectional) {
const dim3 grid(sequence_length * num_heads, batch_size, 1);
if (sequence_length <= 32) {
const int blockSize = 32;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(sequence_length, input, output, is_unidirectional);
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(past_sequence_length, sequence_length, input, output, is_unidirectional);
} else if (sequence_length <= 128) {
const int blockSize = 128;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(sequence_length, input, output, is_unidirectional);
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(past_sequence_length, sequence_length, input, output, is_unidirectional);
} else if (sequence_length == 384) {
const int blockSize = 384;
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(sequence_length, input, output, is_unidirectional);
SoftmaxKernelSmall<T, blockSize><<<grid, blockSize, 0, stream>>>(past_sequence_length, sequence_length, input, output, is_unidirectional);
} else {
const int blockSize = 256;
SoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(sequence_length, input, output, is_unidirectional);
SoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(past_sequence_length, sequence_length, input, output, is_unidirectional);
}
return CUDA_CALL(cudaPeekAtLastError());
@ -197,7 +207,7 @@ __global__ void MaskedSoftmaxKernelSmall(const int sequence_length, const int* m
}
__syncthreads();
SoftmaxSmall<T, TPB>(sequence_length, num_valid, input, output, false);
SoftmaxSmall<T, TPB>(0, sequence_length, num_valid, input, output, false);
}
template <typename T, unsigned TPB>
@ -209,7 +219,7 @@ __global__ void MaskedSoftmaxKernel(const int sequence_length, const int* mask_i
}
__syncthreads();
Softmax<T, TPB>(sequence_length, num_valid, input, output, false);
Softmax<T, TPB>(0, sequence_length, num_valid, input, output, false);
}
template <typename T>
@ -370,6 +380,89 @@ bool LaunchTransQkv(cudaStream_t stream,
return CUDA_CALL(cudaPeekAtLastError());
}
template <typename T>
__global__ void ConcatPastToPresent(const int sequence_length,
const T* past,
const T* k_v,
T* present) {
const int h = threadIdx.x;
const int n = threadIdx.y;
const int s = blockIdx.x;
const int b = blockIdx.y;
const int is_v = blockIdx.z; // 0 for k, 1 for v
const int all_sequence_length = gridDim.x;
const int batch_size = gridDim.y;
const int num_heads = blockDim.y;
const int H = blockDim.x;
// past: 2 x BxNxS'xH (past_k and past_v)
// k_v: 2 x BxNxSxH (k and v)
// present: 2 x BxNxS*xH (present_k and present_v)
const int past_sequence_length = all_sequence_length - sequence_length;
const int present_SH = all_sequence_length * H;
const int present_NSH = num_heads * present_SH;
int out_offset = b * present_NSH + n * present_SH + s * H + h + is_v * (present_NSH * batch_size);
if (s < past_sequence_length) {
const int past_SH = past_sequence_length * H;
const int past_NSH = num_heads * past_SH;
const int in_offset = b * past_NSH + n * past_SH + s * H + h + is_v * (past_NSH * batch_size);
present[out_offset] = past[in_offset];
} else if (s < all_sequence_length) {
const int SH = sequence_length * H;
const int NSH = num_heads * SH;
const int in_offset = b * NSH + n * SH + (s - past_sequence_length) * H + h + is_v * (NSH * batch_size);
present[out_offset] = k_v[in_offset];
}
}
bool LaunchConcatPastToPresent(cudaStream_t stream,
const int past_sequence_length,
const int sequence_length,
const int batch_size,
const int head_size,
const int num_heads,
const float* past,
const float* k_v,
float* present) {
const int all_sequence_length = past_sequence_length + sequence_length;
const dim3 grid(all_sequence_length, batch_size, 2);
if (0 == (head_size & 1)) {
const dim3 block(head_size / 2, num_heads, 1);
ConcatPastToPresent<float2><<<grid, block, 0, stream>>>(sequence_length, reinterpret_cast<const float2*>(past), reinterpret_cast<const float2*>(k_v), reinterpret_cast<float2*>(present));
} else
{
const dim3 block(head_size, num_heads, 1);
ConcatPastToPresent<float><<<grid, block, 0, stream>>>(sequence_length, past, k_v, present);
}
return CUDA_CALL(cudaPeekAtLastError());
}
bool LaunchConcatPastToPresent(cudaStream_t stream,
const int past_sequence_length,
const int sequence_length,
const int batch_size,
const int head_size,
const int num_heads,
const half* past,
const half* k_v,
half* present) {
const int all_sequence_length = past_sequence_length + sequence_length;
const dim3 grid(all_sequence_length, batch_size, 2);
if (0 == (head_size % 4)) {
const dim3 block(head_size / 4, num_heads, 1);
ConcatPastToPresent<float2><<<grid, block, 0, stream>>>(sequence_length, reinterpret_cast<const float2*>(past), reinterpret_cast<const float2*>(k_v), reinterpret_cast<float2*>(present));
} else if (0 == (head_size & 1)) {
const dim3 block(head_size / 2, num_heads, 1);
ConcatPastToPresent<half2><<<grid, block, 0, stream>>>(sequence_length, reinterpret_cast<const half2*>(past), reinterpret_cast<const half2*>(k_v), reinterpret_cast<half2*>(present));
} else { // this should be an "odd" case. probably not worth catching it in the half2 kernel.
const dim3 block(head_size, num_heads, 1);
ConcatPastToPresent<half><<<grid, block, 0, stream>>>(sequence_length, past, k_v, present);
}
return CUDA_CALL(cudaPeekAtLastError());
}
cublasStatus_t inline CublasGemmStridedBatched(
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
int m, int n, int k, const float alpha,
@ -394,8 +487,8 @@ bool QkvToContext(
const int batch_size, const int sequence_length, const int num_heads, const int head_size, const size_t element_size,
const T* input, T* output, T* workspace,
const int* mask_index,
bool is_unidirectional) {
const size_t bytes = ScratchSize(element_size, batch_size, num_heads, sequence_length);
bool is_unidirectional, int past_sequence_length, const T* past, T* present) {
const size_t bytes = ScratchSize(element_size, batch_size, num_heads, sequence_length, past_sequence_length);
T* scratch1 = workspace;
T* scratch2 = scratch1 + (bytes / element_size);
T* scratch3 = scratch2 + (bytes / element_size);
@ -409,7 +502,6 @@ bool QkvToContext(
const int batches = batch_size * num_heads;
const int size_per_batch = sequence_length * head_size;
const int total_size = batches * size_per_batch;
const int temp_matrix_size = sequence_length * sequence_length;
const T* q = scratch3;
const T* k = q + total_size;
@ -418,29 +510,46 @@ bool QkvToContext(
cublasSetStream(cublas, stream);
CublasMathModeSetter helper(cublas, CUBLAS_TENSOR_OP_MATH);
// compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS
// Concat past (2xBxNxS'xH) to present (2xBxNxS*xH):
// past_k (BxNxS'xH) + k (BxNxSxH) => present_k (BxNxS*xH)
// past_v (BxNxS'xH) + v (BxNxSxH) => present_v (BxNxS*xH)
const int present_size_per_batch = (past_sequence_length + sequence_length) * head_size;
if (nullptr != present) {
if (!LaunchConcatPastToPresent(stream, past_sequence_length, sequence_length, batch_size, head_size, num_heads, past, k, present)) {
return false;
}
// update pointers to present_k and present_v.
k = present;
v = present + batches * present_size_per_batch;
}
// compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS*
// Q: BxNxSxH, K (present_k): BxNxS*xH, Q*K': BxNxSxS*
const float rsqrt_head_size = 1.f / sqrt(static_cast<float>(head_size));
const int all_sequence_length = past_sequence_length + sequence_length;
const int temp_matrix_size = sequence_length * all_sequence_length;
if (!CUBLAS_CALL(CublasGemmStridedBatched(
cublas, CUBLAS_OP_T, CUBLAS_OP_N, sequence_length, sequence_length, head_size, rsqrt_head_size, k, head_size, size_per_batch,
q, head_size, size_per_batch, 0.f, scratch1, sequence_length, temp_matrix_size, batches))) {
cublas, CUBLAS_OP_T, CUBLAS_OP_N, all_sequence_length, sequence_length, head_size, rsqrt_head_size, k, head_size, present_size_per_batch,
q, head_size, size_per_batch, 0.f, scratch1, all_sequence_length, temp_matrix_size, batches))) {
return false;
}
// apply softmax and store result P to scratch2: BxNxSxS
// apply softmax and store result P to scratch2: BxNxSxS*
if (nullptr != mask_index) {
if (!ComputeMaskedSoftmax<T>(stream, sequence_length, batch_size, num_heads, mask_index, scratch1, scratch2)) {
return false;
}
} else {
if (!ComputeSoftmax<T>(stream, sequence_length, batch_size, num_heads, scratch1, scratch2, is_unidirectional)) {
if (!ComputeSoftmax<T>(stream, past_sequence_length, sequence_length, batch_size, num_heads, scratch1, scratch2, is_unidirectional)) {
return false;
}
}
// compute P*V (as V*P), and store in scratch3: BxNxSxH
if (!CUBLAS_CALL(CublasGemmStridedBatched(
cublas, CUBLAS_OP_N, CUBLAS_OP_N, head_size, sequence_length, sequence_length, 1.f, v, head_size, size_per_batch,
scratch2, sequence_length, temp_matrix_size, 0.f, scratch3, head_size, size_per_batch, batches))) {
cublas, CUBLAS_OP_N, CUBLAS_OP_N, head_size, sequence_length, all_sequence_length, 1.f, v, head_size, present_size_per_batch,
scratch2, all_sequence_length, temp_matrix_size, 0.f, scratch3, head_size, size_per_batch, batches))) {
return false;
}
@ -459,7 +568,10 @@ bool LaunchAttentionKernel(
void* workspace,
cublasHandle_t& cublas,
const size_t element_size,
bool is_unidirectional) {
bool is_unidirectional,
int past_sequence_length,
const void* past,
void* present) {
// use default stream
const cudaStream_t stream = nullptr;
@ -467,12 +579,14 @@ bool LaunchAttentionKernel(
return QkvToContext(cublas, stream,
batch_size, sequence_length, num_heads, head_size, element_size,
reinterpret_cast<const half*>(input), reinterpret_cast<half*>(output), reinterpret_cast<half*>(workspace),
mask_index, is_unidirectional);
mask_index, is_unidirectional,
past_sequence_length, reinterpret_cast<const half*>(past), reinterpret_cast<half*>(present));
} else {
return QkvToContext(cublas, stream,
batch_size, sequence_length, num_heads, head_size, element_size,
reinterpret_cast<const float*>(input), reinterpret_cast<float*>(output), reinterpret_cast<float*>(workspace),
mask_index, is_unidirectional);
mask_index, is_unidirectional,
past_sequence_length, reinterpret_cast<const float*>(past), reinterpret_cast<float*>(present));
}
}

13
onnxruntime/contrib_ops/cuda/bert/attention_impl.h
View file

@ -7,7 +7,13 @@
namespace onnxruntime {
namespace contrib {
namespace cuda {
size_t GetAttentionWorkspaceSize(size_t element_size, int batchsize, int num_heads, int head_size, int sequence_length);
size_t GetAttentionWorkspaceSize(
size_t element_size,
int batchsize,
int num_heads,
int head_size,
int sequence_length,
int past_sequence_length);
bool LaunchAttentionKernel(
const void* input, // Input tensor
@ -20,7 +26,10 @@ bool LaunchAttentionKernel(
void* workspace, // Temporary buffer
cublasHandle_t& cublas, // Cublas handle
const size_t element_size, // Element size of input tensor
bool is_unidirectional // Whether there is unidirecitonal mask.
bool is_unidirectional, // Whether there is unidirecitonal mask.
int past_sequence_length, // Sequence length in past state
const void* past, // Past state input
void* present // Present state output
);
} // namespace cuda

13
onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc
View file

@ -59,7 +59,7 @@ Status QAttention<T, int8_t>::CheckInputs(const Tensor* input,
// Input 7 - weight_zero_point : scalar
// Output : (batch_size, sequence_length, hidden_size)
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index));
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr));
ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor),
"input scale must be a scalar or 1D tensor of size 1");
@ -161,7 +161,10 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
m,
n);
size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length);
const int past_sequence_length = 0;
const T* past_data = nullptr;
T* present_data = nullptr;
size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, past_sequence_length);
auto temp_buffer = GetScratchBuffer<void>(workSpaceSize);
if (!LaunchAttentionKernel(
reinterpret_cast<const CudaT*>(gemm_buffer.get()),
@ -174,7 +177,11 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
temp_buffer.get(),
cublas,
element_size,
is_unidirectional_)) {
is_unidirectional_,
past_sequence_length,
past_data,
present_data
)) {
// Get last error to reset it to cudaSuccess.
CUDA_CALL(cudaGetLastError());
return Status(common::ONNXRUNTIME, common::FAIL);

43
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -309,10 +309,49 @@ mask_index shall not be provided.)DOC";
.Input(1, "weight", "2D input tensor with shape (hidden_size, 3 * hidden_size)", "T")
.Input(2, "bias", "1D input tensor with shape (3 * hidden_size)", "T")
.Input(3, "mask_index", "Attention mask index with shape (batch_size).", "M", OpSchema::Optional)
.Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T")
.Input(4, "past", "past state for key and value with shape (2, batch_size, num_heads, past_sequence_length, head_size).", "T", OpSchema::Optional)
.Output(0, "output", "3D output tensor with shape (batch_size, append_length, hidden_size)", "T")
.Output(1, "present", "present state for key and value with shape (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)", "T", OpSchema::Optional)
.TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.")
.TypeConstraint("M", {"tensor(int32)"}, "Constrain mask index to integer types")
.TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput);
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (ctx.getNumOutputs() > 1) {
propagateElemTypeFromInputToOutput(ctx, 0, 1);
}
if (hasInputShape(ctx, 0)) {
propagateShapeFromInputToOutput(ctx, 0, 0);
if (ctx.getNumOutputs() > 1) {
auto& input_shape = getInputShape(ctx, 0);
auto& input_dims = input_shape.dim();
if (input_dims.size() != 3) {
fail_shape_inference("Inputs 0 shall be 3 dimensions");
}
if (hasInputShape(ctx, 4)) {
auto& past_shape = getInputShape(ctx, 4);
auto& past_dims = past_shape.dim();
if (past_dims.size() != 5) {
fail_shape_inference("Inputs 4 shall be 5 dimensions");
}
if (past_dims[3].has_dim_value() && input_dims[1].has_dim_value()) {
auto all_sequence_length = past_shape.dim(3).dim_value() + input_shape.dim(1).dim_value();
ONNX_NAMESPACE::TensorShapeProto present_shape;
for (auto& dim : past_dims) {
*present_shape.add_dim() = dim;
}
present_shape.mutable_dim(3)->set_dim_value(all_sequence_length);
updateOutputShape(ctx, 1, present_shape);
}
}
}
}
});
ONNX_CONTRIB_OPERATOR_SCHEMA(QAttention)
.SetDomain(kMSDomain)

22
onnxruntime/python/tools/transformers/BertOnnxModel.py
View file

@ -183,7 +183,7 @@ class BertOnnxModel(OnnxModel):
for node in self.nodes():
# Before:
# input_ids --> Shape --> Gather(indices=0) --> Unsqueeze ------+
# | |
# | |
# | v
# +----> Shape --> Gather(indices=1) --> Unsqueeze---> Concat --> ConstantOfShape -->Cast --> EmbedLayerNormaliation/ReduceSum
# After:
@ -292,8 +292,18 @@ class BertOnnxModel(OnnxModel):
attention = op_count['Attention']
gelu = op_count['Gelu'] + op_count['BiasGelu'] + op_count['FastGelu']
layer_norm = op_count['LayerNormalization'] + op_count['SkipLayerNormalization']
is_optimized = (embed > 0) and (attention > 0) and (attention == gelu) and (layer_norm >= 2 * attention)
logger.info(
f"EmbedLayer={embed}, Attention={attention}, Gelu={gelu}, LayerNormalization={layer_norm}, Successful={is_optimized}"
)
return is_optimized
is_perfect = (embed > 0) and (attention > 0) and (attention == gelu) and (layer_norm >= 2 * attention)
if layer_norm == 0:
logger.debug("Layer Normalization not fused")
if gelu == 0:
logger.debug("Gelu/FastGelu not fused")
if embed == 0:
logger.debug("Embed Layer not fused")
if attention == 0:
logger.debug("Attention not fused")
return is_perfect

142
onnxruntime/python/tools/transformers/Gpt2OnnxModel.py
View file

@ -10,6 +10,8 @@ import numpy as np
from collections import deque
from onnx import ModelProto, TensorProto, numpy_helper
from BertOnnxModel import BertOnnxModel
from fusion_gpt_attention_no_past import FusionGptAttentionNoPast
from fusion_gpt_attention import FusionGptAttention
logger = logging.getLogger(__name__)
@ -19,135 +21,12 @@ class Gpt2OnnxModel(BertOnnxModel):
super().__init__(model, num_heads, hidden_size)
def fuse_attention(self):
"""
Fuse Attention subgraph into one Attention node.
"""
logger.debug(f"start attention fusion...")
input_name_to_nodes = self.input_name_to_nodes()
output_name_to_node = self.output_name_to_node()
attention_count = 0
for normalize_node in self.get_nodes_by_op_type("LayerNormalization"):
return_indice = []
qkv_nodes = self.match_parent_path(
normalize_node,
['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'],
[0, None, 0, 0, 0, 0, 0],
output_name_to_node=output_name_to_node,
return_indice=return_indice
) # yapf: disable
if qkv_nodes is None:
continue
(add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes
another_input = add_qkv.input[1 - return_indice[0]]
v_nodes = self.match_parent_path(
matmul_qkv,
['Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'],
[1, 0, 0, 0, 0, 0]) # yapf: disable
if v_nodes is None:
logger.debug("fuse_attention: failed to match v path")
continue
(transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes
layernorm_before_attention = self.get_parent(reshape_before_gemm, 0, output_name_to_node)
if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization':
logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}")
continue
if not another_input in layernorm_before_attention.input:
logger.debug("Add and LayerNormalization shall have one same input")
continue
qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0])
if qk_nodes is not None:
(softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes
mask_nodes = self.match_parent_path(
sub_qk,
['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
[1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
continue
div_mask = mask_nodes[-1]
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
continue
else:
# New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0.
qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'Div', 'MatMul'], [0, 0, 1, 0])
if qk_nodes is None:
logger.debug("fuse_attention: failed to match qk path")
continue
(softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes
mask_nodes = self.match_parent_path(
where_qk,
['Cast', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
[ 0, 0, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
continue
div_mask = mask_nodes[-1]
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
continue
q_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0])
if q_nodes is None:
logger.debug("fuse_attention: failed to match q path")
continue
(transpose_q, reshape_q, split_q) = q_nodes
if split_v != split_q:
logger.debug("fuse_attention: skip since split_v != split_q")
continue
k_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [1, 0, 0])
if k_nodes is None:
logger.debug("fuse_attention: failed to match k path")
continue
(transpose_k, reshape_k, split_k) = k_nodes
if split_v != split_k:
logger.debug("fuse_attention: skip since split_v != split_k")
continue
self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0],
attention_count == 0)
# we rely on prune_graph() to clean old subgraph nodes:
# qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv]
attention_count += 1
self.prune_graph()
logger.info(f"Fused Attention count:{attention_count}")
def create_attention_node(self, gemm, gemm_qkv, input, output, add_graph_input):
attention_node_name = self.create_node_name('Attention')
attention_node = onnx.helper.make_node('Attention',
inputs=[input, gemm.input[1], gemm.input[2]],
outputs=[attention_node_name + "_output"],
name=attention_node_name)
attention_node.domain = "com.microsoft"
attention_node.attribute.extend(
[onnx.helper.make_attribute("num_heads", self.num_heads),
onnx.helper.make_attribute("unidirectional", 1)])
matmul_node = onnx.helper.make_node('MatMul',
inputs=[attention_node_name + "_output", gemm_qkv.input[1]],
outputs=[attention_node_name + "_matmul_output"],
name=attention_node_name + "_matmul")
add_node = onnx.helper.make_node('Add',
inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]],
outputs=[output],
name=attention_node_name + "_add")
self.add_node(attention_node)
self.add_node(matmul_node)
self.add_node(add_node)
if len(self.model.graph.input) == 1 or len(self.model.graph.output) == 1:
fusion = FusionGptAttentionNoPast(self, self.num_heads)
fusion.apply()
else:
fusion = FusionGptAttention(self, self.num_heads)
fusion.apply()
def postprocess(self):
"""
@ -186,7 +65,12 @@ class Gpt2OnnxModel(BertOnnxModel):
outputs=[add_node_name + "_output"],
name=add_node_name)
self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output")
# Link root node output with MatMul
self.replace_input_of_all_nodes(root_node.output[0], matmul_node_name + "_input")
root_node.output[0] = matmul_node_name + "_input"
self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output")
self.add_node(matmul_node)

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

@ -222,12 +222,13 @@ def get_onnx_file_path(onnx_dir: str, model_name: str, input_count: int, optimiz
def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwrite):
if overwrite or not os.path.exists(ort_model_path):
from optimizer import optimize_model, get_fusion_statistics
# Use onnxruntime to optimize model, which will be saved to *_ort.onnx
opt_model = optimize_by_onnxruntime(onnx_model_path,
use_gpu=use_gpu,
optimized_model_path=ort_model_path,
opt_level=99)
model_fusion_statistics[ort_model_path] = opt_model.get_fused_operator_statistics()
model_fusion_statistics[ort_model_path] = get_fusion_statistics(ort_model_path)
else:
logger.info(f"Skip optimization since model existed: {ort_model_path}")
@ -235,7 +236,7 @@ def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwri
def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_attention_heads, hidden_size, use_gpu,
fp16, overwrite):
if overwrite or not os.path.exists(optimized_model_path):
from optimizer import optimize_model, optimize_by_onnxruntime
from optimizer import optimize_model
from BertOnnxModel import BertOptimizationOptions
optimization_options = BertOptimizationOptions(model_type)
if fp16:

109
onnxruntime/python/tools/transformers/benchmark_gpt2.py
View file

@ -18,10 +18,10 @@ from transformers import GPT2Model, GPT2LMHeadModel, GPT2Tokenizer, AutoConfig
logger = logging.getLogger('')
# Map alias to a tuple of Model Class and pretrained model name
# Map alias to a tuple of Model Class, Tokenizer, pretrained model name, use LMHead or not, use attention mask or not
MODEL_CLASSES = {
"gpt2": (GPT2Model, GPT2Tokenizer, "gpt2"),
"distilgpt2": (GPT2LMHeadModel, GPT2Tokenizer, "distilgpt2"),
"gpt2": (GPT2Model, GPT2Tokenizer, "gpt2", False, False),
"distilgpt2": (GPT2LMHeadModel, GPT2Tokenizer, "distilgpt2", True, True),
}
@ -44,12 +44,12 @@ def setup_environment():
dump_environment()
def pytorch_inference(model, input_ids, past=None, total_runs=100):
def pytorch_inference(model, input_ids, past=None, attention_mask=None, total_runs=100):
latency = []
with torch.no_grad():
for _ in range(total_runs):
start = time.time()
outputs = model(input_ids=input_ids, past=past)
outputs = model(input_ids=input_ids, past=past, attention_mask=attention_mask)
latency.append(time.time() - start)
average_latency = sum(latency) * 1000 / len(latency)
@ -57,10 +57,12 @@ def pytorch_inference(model, input_ids, past=None, total_runs=100):
return outputs, average_latency
def onnxruntime_inference(ort_session, input_ids, past=None, total_runs=100):
def onnxruntime_inference(ort_session, input_ids, past=None, attention_mask=None, total_runs=100):
ort_inputs = {'input_ids': numpy.ascontiguousarray(input_ids.cpu().numpy())}
# TODO: pass input tensor stored in GPU
if attention_mask is not None:
ort_inputs['attention_mask'] = numpy.ascontiguousarray(attention_mask.cpu().numpy())
if past is not None:
for i, past_i in enumerate(past):
ort_inputs[f'past_{i}'] = numpy.ascontiguousarray(past[i].cpu().numpy())
@ -82,6 +84,7 @@ def onnxruntime_inference_with_binded_io(ort_session,
last_state,
last_state_shape,
past=None,
attention_mask=None,
present=None,
present_shape=None,
total_runs=100):
@ -90,13 +93,17 @@ def onnxruntime_inference_with_binded_io(ort_session,
# Bind inputs
io_binding.bind_input('input_ids', input_ids.device.type, 0, numpy.longlong, list(input_ids.size()),
input_ids.data_ptr())
if attention_mask is not None:
io_binding.bind_input('attention_mask', attention_mask.device.type, 0, numpy.float32, list(attention_mask.size()),
attention_mask.data_ptr())
if past is not None:
for i, past_i in enumerate(past):
io_binding.bind_input(f'past_{i}', past[i].device.type, 0, numpy.float32, list(past[i].size()),
past[i].data_ptr())
# Bind outputs
io_binding.bind_output("last_state", last_state.device.type, 0, numpy.float32, last_state_shape,
io_binding.bind_output(ort_session.get_outputs()[0].name, last_state.device.type, 0, numpy.float32, last_state_shape,
last_state.data_ptr())
if present is not None:
for i, present_i in enumerate(present):
@ -125,6 +132,7 @@ def inference(model,
ort_session,
input_ids,
past=None,
attention_mask=None,
last_state=None,
present=None,
last_state_shape=None,
@ -132,12 +140,12 @@ def inference(model,
total_runs=100,
verify_outputs=True,
with_io_binding=False):
outputs, torch_latency = pytorch_inference(model, input_ids, past, total_runs)
ort_outputs, ort_latency = onnxruntime_inference(ort_session, input_ids, past, total_runs)
outputs, torch_latency = pytorch_inference(model, input_ids, past, attention_mask, total_runs)
ort_outputs, ort_latency = onnxruntime_inference(ort_session, input_ids, past, attention_mask, total_runs)
latencies = [torch_latency, ort_latency]
if with_io_binding:
ort_io_outputs, ort_io_latency = onnxruntime_inference_with_binded_io(ort_session, input_ids, last_state,
last_state_shape, past, present,
last_state_shape, past, attention_mask, present,
present_shape, total_runs)
latencies.append(ort_io_latency)
if verify_outputs:
@ -203,7 +211,8 @@ def parse_arguments():
help='Use optimizer.py to optimize onnx model')
parser.set_defaults(optimize_onnx=False)
parser.add_argument('--with_io_binding',
parser.add_argument('-i',
'--with_io_binding',
required=False,
action='store_true',
help='Run ONNX Runtime with binded inputs and outputs. ')
@ -243,7 +252,9 @@ def setup_logger(verbose=True):
logger.setLevel(logging_level)
def export_onnx(model, config, tokenizer, device, output_dir):
def export_onnx(model, config, tokenizer, device, output_dir, use_LMHead=False, use_attention_mask=False):
""" Export GPT-2 model with past state to ONNX model
"""
model.to(device)
inputs = tokenizer.encode_plus("Here is an example input for GPT2 model",
@ -251,56 +262,59 @@ def export_onnx(model, config, tokenizer, device, output_dir):
return_tensors='pt')
input_ids = inputs['input_ids'].to(device)
logger.debug(f"input_ids={input_ids}")
# Use example input to generate an example of past state.
outputs = model(input_ids=input_ids, past=None)
assert len(outputs) == 2
logger.debug(f"output 0 shape={outputs[0].shape}")
logger.debug(f"outputs[1][0] shape={outputs[1][0].shape}")
num_layer = model.config.n_layer
present_names = [f'present_{i}' for i in range(num_layer)]
output_names = ["last_state"] + present_names
input_names = ['input_ids']
# input_ids has only one word for model with past state.
# Shape of input tensors:
# input_ids: (batch_size, 1)
# past_{i}: (2, batch_size, num_heads, seq_len, hidden_size/num_heads)
# Shape of output tensors:
# last_state: (batch_size, seq_len + 1, hidden_size)
# present_{i}: (2, batch_size, num_heads, seq_len + 1, hidden_size/num_heads)
dynamic_axes = {'input_ids': {0: 'batch_size'}, 'last_state': {0: 'batch_size', 1: 'seq_len_plus_1'}}
for name in present_names:
dynamic_axes[name] = {1: 'batch_size', 3: 'seq_len_plus_1'}
past_names = [f'past_{i}' for i in range(num_layer)]
input_names = ['input_ids'] + past_names
dummy_past = [torch.zeros(list(outputs[1][0].shape), dtype=torch.float32, device=device) for _ in range(num_layer)]
present_names = [f'present_{i}' for i in range(num_layer)]
# GPT2Model output last_state has shape (batch_size, all_seq_len, hidden_size)
# GPT2LMHeadModel output prediction_scores has shape (batch_size, all_seq_len, vocab_size)
# where all_seq_len = past_seq_len + seq_len
output_names = ["prediction_scores" if use_LMHead else "last_state"] + present_names
# Shape of input tensors:
# input_ids: (batch_size, seq_len)
# past_{i}: (2, batch_size, num_heads, past_seq_len, hidden_size/num_heads)
# attention_mask: (batch_size, seq_len)
# Shape of output tensors:
# last_state: (batch_size, all_seq_len, hidden_size)
# present_{i}: (2, batch_size, num_heads, all_seq_len, hidden_size/num_heads)
dynamic_axes = {'input_ids': {0: 'batch_size', 1 : 'seq_len'}, output_names[0]: {0: 'batch_size', 1: 'all_seq_len'}}
for name in past_names:
dynamic_axes[name] = {1: 'batch_size', 3: 'seq_len'}
logger.debug(f"vocab_size:{model.config.vocab_size}")
dynamic_axes[name] = {1: 'batch_size', 3: 'past_seq_len'}
for name in present_names:
dynamic_axes[name] = {1: 'batch_size', 3: 'all_seq_len'}
if use_attention_mask:
dynamic_axes['attention_mask'] = {0: 'batch_size', 1: 'seq_len'}
dummy_input_ids = torch.randint(low=0,
high=model.config.vocab_size - 1,
size=(1, 1),
dtype=torch.int64,
device=device)
logger.debug(f"dummy_input_ids={dummy_input_ids}")
export_inputs = (dummy_input_ids, tuple(dummy_past))
# Use the example past state to create dummy past state inputs.
dummy_past = [torch.zeros(list(outputs[1][0].shape), dtype=torch.float32, device=device) for _ in range(num_layer)]
export_model_path = os.path.join(output_dir, 'gpt2_past.onnx')
dummy_mask = torch.ones([1, 1], dtype=torch.float32, device=device) if use_attention_mask else None
model_name = "gpt2{}_past{}.onnx".format("_lm" if use_LMHead else "", "_mask" if use_attention_mask else "")
export_model_path = os.path.join(output_dir, model_name)
# Let's run performance test on PyTorch before updating environment variable.
with torch.no_grad():
outputs = model(input_ids=dummy_input_ids, past=dummy_past)
outputs = model(input_ids=dummy_input_ids, past=dummy_past, attention_mask=dummy_mask)
logger.debug(f"present_0 shape={outputs[1][0].shape}")
torch.onnx.export(model,
args=export_inputs,
args=(dummy_input_ids, tuple(dummy_past), dummy_mask) if use_attention_mask else (dummy_input_ids, tuple(dummy_past)),
f=export_model_path,
input_names=input_names,
input_names=['input_ids'] + past_names + (['attention_mask'] if use_attention_mask else []),
output_names=output_names,
example_outputs=outputs,
dynamic_axes=dynamic_axes,
@ -324,7 +338,7 @@ def main():
os.makedirs(output_dir)
use_torchscript = False
(model_class, tokenizer_class, model_name) = MODEL_CLASSES[args.model_type]
(model_class, tokenizer_class, model_name, use_LMHead, use_attention_mask) = MODEL_CLASSES[args.model_type]
config = AutoConfig.from_pretrained(model_name, torchscript=use_torchscript, cache_dir=cache_dir)
model = model_class.from_pretrained(model_name, config=config, cache_dir=cache_dir)
tokenizer = tokenizer_class.from_pretrained(model_name, cache_dir=cache_dir)
@ -332,7 +346,7 @@ def main():
# model = torch.jit.trace(model, (input_ids, past))
device = torch.device("cuda:0" if args.use_gpu else "cpu")
export_model_path = export_onnx(model, config, tokenizer, device, output_dir)
export_model_path = export_onnx(model, config, tokenizer, device, output_dir, use_LMHead, use_attention_mask)
# setup environment variables before importing onnxruntime.
setup_environment()
@ -379,7 +393,10 @@ def main():
]
# dummy last state
last_state_size = numpy.prod([max_batch_size, 1, config.hidden_size])
if use_LMHead:
last_state_size = numpy.prod([max_batch_size, 1, config.vocab_size])
else:
last_state_size = numpy.prod([max_batch_size, 1, config.hidden_size])
dummy_last_state = torch.empty(last_state_size).to(device)
for batch_size in args.batch_sizes:
@ -394,6 +411,7 @@ def main():
size=(batch_size, 1),
dtype=torch.int64,
device=device)
dummy_mask = torch.ones([batch_size, 1], dtype=torch.float32, device=device) if use_attention_mask else None
# Calculate the expected output shapes
last_state_shape = [batch_size, 1, config.hidden_size]
@ -406,6 +424,7 @@ def main():
session,
dummy_input_ids,
dummy_past,
dummy_mask,
dummy_last_state,
dummy_present,
last_state_shape,

9
onnxruntime/python/tools/transformers/bert_perf_test.py
View file

@ -411,7 +411,11 @@ def parse_arguments():
parser.add_argument('--input_ids_name', required=False, type=str, default=None, help="input name for input ids")
parser.add_argument('--segment_ids_name', required=False, type=str, default=None, help="input name for segment ids")
parser.add_argument('--input_mask_name', required=False, type=str, default=None, help="input name for attention mask")
parser.add_argument('--input_mask_name',
required=False,
type=str,
default=None,
help="input name for attention mask")
args = parser.parse_args()
return args
@ -430,7 +434,8 @@ def main():
if not min(batch_size_set) >= 1 and max(batch_size_set) <= 128:
raise Exception("batch_size not in range [1, 128]")
model_setting = ModelSetting(args.model, args.input_ids_name, args.segment_ids_name, args.input_mask_name, args.opt_level)
model_setting = ModelSetting(args.model, args.input_ids_name, args.segment_ids_name, args.input_mask_name,
args.opt_level)
for batch_size in batch_size_set:
test_setting = TestSetting(

6
onnxruntime/python/tools/transformers/bert_test_data.py
View file

@ -237,7 +237,11 @@ def parse_arguments():
parser.add_argument('--input_ids_name', required=False, type=str, default=None, help="input name for input ids")
parser.add_argument('--segment_ids_name', required=False, type=str, default=None, help="input name for segment ids")
parser.add_argument('--input_mask_name', required=False, type=str, default=None, help="input name for attention mask")
parser.add_argument('--input_mask_name',
required=False,
type=str,
default=None,
help="input name for attention mask")
parser.add_argument('--samples', required=False, type=int, default=1, help="number of test cases to be generated")

2
onnxruntime/python/tools/transformers/dev_benchmark.cmd
View file

@ -48,7 +48,7 @@ REM This script will generate a logs file with a list of commands used in tests.
>benchmark.log echo echo ort=%run_ort% torch=%run_torch% torchscript=%run_torchscript% gpu_fp32=%run_gpu_fp32% gpu_fp16=%run_gpu_fp16% cpu=%run_cpu% optimizer=%use_optimizer% batch="%batch_sizes%" sequence="%sequence_length%" models="%models_to_test%" input_counts="%input_counts%"
REM Set it to false to skip testing. You can use it to dry run this script with the benchmark.log file.
set run_tests=false
set run_tests=true
REM -------------------------------------------
if %run_cpu% == true if %run_gpu_fp32% == true echo cannot test cpu and gpu at same time & goto :EOF

7
onnxruntime/python/tools/transformers/fusion_attention.py
View file

@ -188,10 +188,8 @@ class FusionAttention(Fusion):
# Note that Cast might be removed by OnnxRuntime so we match two patterns here.
_, mask_nodes, _ = self.model.match_parent_paths(
add_qk,
[(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]),
(['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])],
output_name_to_node)
add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]),
(['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node)
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return
@ -215,4 +213,3 @@ class FusionAttention(Fusion):
# Use prune graph to remove mask nodes since they are shared by all attention nodes.
#self.nodes_to_remove.extend(mask_nodes)
self.prune_graph = True

17
onnxruntime/python/tools/transformers/fusion_base.py
View file

@ -10,17 +10,21 @@ logger = getLogger(__name__)
class Fusion:
def __init__(self, model: OnnxModel, name: str, search_op_types: Union[str, List[str]]):
def __init__(self,
model: OnnxModel,
fused_op_type: str,
search_op_types: Union[str, List[str]],
description: str = None):
self.search_op_types: List[str] = [search_op_types] if isinstance(search_op_types, str) else search_op_types
self.name: str = name
self.fused_op_type: str = fused_op_type
self.description: str = f"{fused_op_type}({description})" if description else fused_op_type
self.model: OnnxModel = model
self.nodes_to_remove: List = []
self.nodes_to_add: List = []
self.prune_graph: bool = False
def apply(self):
logger.debug(f"start {self.name} fusion...")
logger.debug(f"start {self.description} fusion...")
input_name_to_nodes = self.model.input_name_to_nodes()
output_name_to_node = self.model.output_name_to_node()
@ -29,7 +33,10 @@ class Fusion:
for node in self.model.get_nodes_by_op_type(search_op_type):
self.fuse(node, input_name_to_nodes, output_name_to_node)
logger.info(f"Fused {self.name} count: {len(self.nodes_to_add)}")
op_list = [node.op_type for node in self.nodes_to_add]
count = op_list.count(self.fused_op_type)
if count > 0:
logger.info(f"Fused {self.description} count: {count}")
self.model.remove_nodes(self.nodes_to_remove)
self.model.add_nodes(self.nodes_to_add)

2
onnxruntime/python/tools/transformers/fusion_biasgelu.py
View file

@ -14,7 +14,7 @@ logger = getLogger(__name__)
class FusionBiasGelu(Fusion):
def __init__(self, model: OnnxModel, is_fastgelu):
if is_fastgelu:
super().__init__(model, 'FastGelu(add bias)', 'FastGelu')
super().__init__(model, 'FastGelu', 'FastGelu', 'add bias')
else:
super().__init__(model, 'BiasGelu', 'Gelu')

12
onnxruntime/python/tools/transformers/fusion_embedlayer.py
View file

@ -38,9 +38,8 @@ class FusionEmbedLayerNoMask(Fusion):
"""
def __init__(self,
model: OnnxModel,
name: str = "EmbedLayerNormalization(no mask)",
search_op_types="SkipLayerNormalization"):
super().__init__(model, name, search_op_types)
description='no mask'):
super().__init__(model, "EmbedLayerNormalization", "SkipLayerNormalization", description)
self.utils = FusionUtils(model)
def fuse(self, node, input_name_to_nodes, output_name_to_node):
@ -141,10 +140,9 @@ class FusionEmbedLayerNoMask(Fusion):
# Cast might be removed by OnnxRuntime.
_, segment_id_path, _ = self.model.match_parent_paths(
segment_ids_cast_node,
segment_ids_cast_node,
[(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]),
(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])],
output_name_to_node)
(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node)
if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]:
logger.debug("Simplify semgent id path...")
@ -187,7 +185,7 @@ class FusionEmbedLayerNoMask(Fusion):
class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask):
def __init__(self, model: OnnxModel, mask_indice: Dict, mask_casted: Dict):
super().__init__(model, "EmbedLayerNormalization(with mask)", "SkipLayerNormalization")
super().__init__(model, "with mask")
self.mask_indice: Dict = mask_indice
self.mask_casted: Dict = mask_casted
self.mask_input_name = None

2
onnxruntime/python/tools/transformers/fusion_gelu_approximation.py
View file

@ -11,7 +11,7 @@ from fusion_base import Fusion
class FusionGeluApproximation(Fusion):
def __init__(self, model: OnnxModel):
super().__init__(model, 'FastGelu(GeluApproximation)', ['Gelu', 'BiasGelu'])
super().__init__(model, 'FastGelu', ['Gelu', 'BiasGelu'], 'GeluApproximation')
def fuse(self, node, input_name_to_nodes, output_name_to_node):
new_node = helper.make_node("FastGelu",

196
onnxruntime/python/tools/transformers/fusion_gpt_attention.py Normal file
View file

@ -0,0 +1,196 @@
#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
import numpy as np
from logging import getLogger
from onnx import helper, numpy_helper, TensorProto
from OnnxModel import OnnxModel
from fusion_base import Fusion
from fusion_utils import FusionUtils
logger = getLogger(__name__)
class FusionGptAttention(Fusion):
"""
Fuse GPT-2 Attention with past state subgraph into one Attention node.
This does not support attention_mask graph input right now.
"""
def __init__(self, model: OnnxModel, num_heads: int):
super().__init__(model, "Attention", "LayerNormalization", "with past")
self.num_heads = num_heads
def create_attention_node(self, gemm, gemm_qkv, past, present, input, output):
attention_node_name = self.model.create_node_name('GptAttention')
mask_index = ''
attention_node = helper.make_node('Attention',
inputs=[input, gemm.input[1], gemm.input[2], mask_index, past],
outputs=[attention_node_name + "_output", present],
name=attention_node_name)
attention_node.domain = "com.microsoft"
attention_node.attribute.extend(
[helper.make_attribute("num_heads", self.num_heads),
helper.make_attribute("unidirectional", 1)])
matmul_node = helper.make_node('MatMul',
inputs=[attention_node_name + "_output", gemm_qkv.input[1]],
outputs=[attention_node_name + "_matmul_output"],
name=attention_node_name + "_matmul")
add_node = helper.make_node('Add',
inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]],
outputs=[output],
name=attention_node_name + "_add")
self.nodes_to_add.extend([attention_node, matmul_node, add_node])
def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
past = None
present = None
return_indice = []
qkv_nodes = self.model.match_parent_path(
normalize_node,
['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'],
[0, None, 0, 0, 0, 0, 0],
output_name_to_node=output_name_to_node,
return_indice=return_indice
) # yapf: disable
if qkv_nodes is None:
return
(add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes
another_input = add_qkv.input[1 - return_indice[0]]
v_nodes = self.model.match_parent_path(
matmul_qkv,
['Concat', 'Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'],
[1, 1, 0, 0, 0, 0, 0]) # yapf: disable
if v_nodes is None:
logger.debug("fuse_attention: failed to match v path")
return
(concat_v, transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes
# concat <-- Gather(indices=1) <-- past
# |
# unsqueeze
# |
# concat --> present
gather_v = self.model.get_parent(concat_v, 0, output_name_to_node)
if gather_v.op_type != 'Gather':
logger.info("expect Gather for past")
return
if not self.model.find_constant_input(gather_v, 1) == 1:
logger.info("expect indices=1 for Gather of past")
return
past = gather_v.input[0]
if not self.model.find_graph_input(past):
logger.info("expect past to be graph input")
return
unsqueeze_present_v = self.model.find_first_child_by_type(concat_v,
'Unsqueeze',
input_name_to_nodes,
recursive=False)
if not unsqueeze_present_v:
logger.info("expect unsqueeze for present")
return
concat_present = self.model.find_first_child_by_type(unsqueeze_present_v,
'Concat',
input_name_to_nodes,
recursive=False)
if not concat_present:
logger.info("expect concat for present")
return
present = concat_present.output[0]
if not self.model.find_graph_output(present):
logger.info("expect present to be graph input")
return
layernorm_before_attention = self.model.get_parent(reshape_before_gemm, 0, output_name_to_node)
if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization':
logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}")
return
if not another_input in layernorm_before_attention.input:
logger.debug("Add and LayerNormalization shall have one same input")
return
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0])
if qk_nodes is not None:
(softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes
mask_nodes = self.model.match_parent_path(
sub_qk,
['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
[1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return
div_mask = mask_nodes[-1]
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
return
else:
# New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0.
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'Div', 'MatMul'], [0, 0, 1, 0])
if qk_nodes is None:
logger.debug("fuse_attention: failed to match qk path")
return
(softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes
mask_nodes = self.model.match_parent_path(
where_qk,
['Cast', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
[ 0, 0, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return
div_mask = mask_nodes[-1]
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
return
q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0])
if q_nodes is None:
logger.debug("fuse_attention: failed to match q path")
return
(transpose_q, reshape_q, split_q) = q_nodes
if split_v != split_q:
logger.debug("fuse_attention: skip since split_v != split_q")
return
k_nodes = self.model.match_parent_path(matmul_qk, ['Concat', 'Transpose', 'Reshape', 'Split'], [1, 1, 0, 0])
if k_nodes is None:
logger.debug("fuse_attention: failed to match k path")
return
(concat_k, transpose_k, reshape_k, split_k) = k_nodes
if split_v != split_k:
logger.debug("fuse_attention: skip since split_v != split_k")
return
# concat_k <-- Transpose (perm=0,1,3,2) <-- Gather(axes=0, indices=0) <-- past
# |
# Transpose (perm=0,1,3,2)
# |
# unsqueeze
# |
# concat --> present
past_k_nodes = self.model.match_parent_path(concat_k, ['Transpose', 'Gather'], [0, 0])
if past_k_nodes is None:
logger.debug("fuse_attention: failed to match past_k_nodes path")
return
gather_past_k = past_k_nodes[-1]
if not self.model.find_constant_input(gather_past_k, 0) == 1:
logger.info("expect indices=0 for Gather k of past")
return
past_k = gather_past_k.input[0]
if past != past_k:
logger.info("expect past to be same")
return
self.create_attention_node(gemm, gemm_qkv, past, present, layernorm_before_attention.output[0],
reshape_qkv.output[0])
# we rely on prune_graph() to clean old subgraph nodes:
# qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv]
self.prune_graph = True

137
onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py Normal file
View file

@ -0,0 +1,137 @@
#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
import numpy as np
from logging import getLogger
from onnx import helper, numpy_helper, TensorProto
from OnnxModel import OnnxModel
from fusion_base import Fusion
from fusion_utils import FusionUtils
logger = getLogger(__name__)
class FusionGptAttentionNoPast(Fusion):
"""
Fuse GPT-2 Attention without past state into one Attention node.
This does not support attention_mask graph input right now.
"""
def __init__(self, model: OnnxModel, num_heads: int):
super().__init__(model, "Attention", "LayerNormalization", "without past")
self.num_heads = num_heads
def create_attention_node(self, gemm, gemm_qkv, input, output):
attention_node_name = self.model.create_node_name('Attention')
attention_node = helper.make_node('Attention',
inputs=[input, gemm.input[1], gemm.input[2]],
outputs=[attention_node_name + "_output"],
name=attention_node_name)
attention_node.domain = "com.microsoft"
attention_node.attribute.extend(
[helper.make_attribute("num_heads", self.num_heads),
helper.make_attribute("unidirectional", 1)])
matmul_node = helper.make_node('MatMul',
inputs=[attention_node_name + "_output", gemm_qkv.input[1]],
outputs=[attention_node_name + "_matmul_output"],
name=attention_node_name + "_matmul")
add_node = helper.make_node('Add',
inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]],
outputs=[output],
name=attention_node_name + "_add")
self.nodes_to_add.extend([attention_node, matmul_node, add_node])
def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
return_indice = []
qkv_nodes = self.model.match_parent_path(
normalize_node,
['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'],
[0, None, 0, 0, 0, 0, 0],
output_name_to_node=output_name_to_node,
return_indice=return_indice
) # yapf: disable
if qkv_nodes is None:
return
(add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes
another_input = add_qkv.input[1 - return_indice[0]]
v_nodes = self.model.match_parent_path(
matmul_qkv,
['Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'],
[1, 0, 0, 0, 0, 0]) # yapf: disable
if v_nodes is None:
logger.debug("fuse_attention: failed to match v path")
return
(transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes
layernorm_before_attention = self.model.get_parent(reshape_before_gemm, 0, output_name_to_node)
if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization':
logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}")
return
if not another_input in layernorm_before_attention.input:
logger.debug("Add and LayerNormalization shall have one same input")
return
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0])
if qk_nodes is not None:
(softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes
mask_nodes = self.model.match_parent_path(
sub_qk,
['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
[1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return
div_mask = mask_nodes[-1]
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
return
else:
# New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0.
qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'Div', 'MatMul'], [0, 0, 1, 0])
if qk_nodes is None:
logger.debug("fuse_attention: failed to match qk path")
return
(softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes
mask_nodes = self.model.match_parent_path(
where_qk,
['Cast', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'],
[ 0, 0, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable
if mask_nodes is None:
logger.debug("fuse_attention: failed to match mask path")
return
div_mask = mask_nodes[-1]
if div_qk != div_mask:
logger.debug("fuse_attention: skip since div_qk != div_mask")
return
q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0])
if q_nodes is None:
logger.debug("fuse_attention: failed to match q path")
return
(transpose_q, reshape_q, split_q) = q_nodes
if split_v != split_q:
logger.debug("fuse_attention: skip since split_v != split_q")
return
k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [1, 0, 0])
if k_nodes is None:
logger.debug("fuse_attention: failed to match k path")
return
(transpose_k, reshape_k, split_k) = k_nodes
if split_v != split_k:
logger.debug("fuse_attention: skip since split_v != split_k")
return
self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0])
# we rely on prune_graph() to clean old subgraph nodes:
# qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv]
self.prune_graph = True

2
onnxruntime/python/tools/transformers/fusion_layernorm.py
View file

@ -117,7 +117,7 @@ class FusionLayerNormalization(Fusion):
class FusionLayerNormalizationTF(Fusion):
def __init__(self, model: OnnxModel):
super().__init__(model, "LayerNormalization", "Add")
super().__init__(model, "LayerNormalization", "Add", "TF")
def fuse(self, node, input_name_to_nodes: Dict, output_name_to_node: Dict):
"""

2
onnxruntime/python/tools/transformers/fusion_skiplayernorm.py
View file

@ -47,7 +47,7 @@ class FusionSkipLayerNormalization(Fusion):
class FusionBiasSkipLayerNormalization(Fusion):
def __init__(self, model: OnnxModel):
super().__init__(model, "SkipLayerNormalization(add bias)", "SkipLayerNormalization")
super().__init__(model, "SkipLayerNormalization", "SkipLayerNormalization", "add bias")
def fuse(self, node, input_name_to_nodes, output_name_to_node):
if len(node.input) != 4:

117
onnxruntime/python/tools/transformers/optimizer.py
View file

@ -6,22 +6,17 @@
# Convert Bert ONNX model converted from TensorFlow or exported from PyTorch to use Attention, Gelu,
# SkipLayerNormalization and EmbedLayerNormalization ops to optimize
# performance on NVidia GPU and CPU.
#
# For Bert model exported from PyTorch, OnnxRuntime has bert model optimization support internally.
# You can use the option --use_onnxruntime to use model optimization from OnnxRuntime package.
# You can use the option --use_onnxruntime to check optimizations from OnnxRuntime.
# For Bert model file like name.onnx, optimized model for GPU or CPU from OnnxRuntime will output as
# name_ort_gpu.onnx or name_ort_cpu.onnx in the same directory.
#
# This script is retained for experiment purpose. Useful senarios like the following:
# (1) Change model from fp32 to fp16.
# (1) Change model from fp32 to fp16 for mixed precision inference in GPU with Tensor Core.
# (2) Change input data type from int64 to int32.
# (3) Some model cannot be handled by OnnxRuntime, and you can modify this script to get optimized model.
# This script has been tested using the following models:
# (1) BertForSequenceClassification as in https://github.com/huggingface/transformers/blob/master/examples/run_glue.py
# PyTorch 1.2 or above, and exported to Onnx using opset version 10 or 11.
# (2) BertForQuestionAnswering as in https://github.com/huggingface/transformers/blob/master/examples/run_squad.py
# PyTorch 1.2 or above, and exported to Onnx using opset version 10 or 11.
import logging
import coloredlogs
import onnx
@ -29,6 +24,7 @@ import os
import sys
import argparse
import numpy as np
from typing import Dict
from collections import deque
from onnx import ModelProto, TensorProto, numpy_helper, load_model
from BertOnnxModel import BertOnnxModel, BertOptimizationOptions
@ -47,17 +43,21 @@ MODEL_CLASSES = {
}
def optimize_by_onnxruntime(onnx_model_path, use_gpu=False, optimized_model_path=None, opt_level=99):
def optimize_by_onnxruntime(onnx_model_path: str,
use_gpu: bool = False,
optimized_model_path: str = None,
opt_level: int = 99) -> str:
"""
Use onnxruntime package to optimize model. It could support models exported by PyTorch.
Use onnxruntime to optimize model.
Args:
onnx_model_path (str): th path of input onnx model.
onnx_model_path (str): the path of input onnx model.
use_gpu (bool): whether the optimized model is targeted to run in GPU.
optimized_model_path (str or None): the path of optimized model.
opt_level (int): graph optimization level.
Returns:
optimized_model_path: the path of optimized model
optimized_model_path (str): the path of optimized model
"""
import onnxruntime
@ -91,7 +91,22 @@ def optimize_by_onnxruntime(onnx_model_path, use_gpu=False, optimized_model_path
return optimized_model_path
def parse_arguments():
def get_fusion_statistics(optimized_model_path: str) -> Dict[str, int]:
"""
Get counter of fused operators in optimized model.
Args:
optimized_model_path (str): the path of onnx model.
Returns:
A dictionary with operator type as key, and count as value
"""
model = load_model(optimized_model_path, format=None, load_external_data=True)
optimizer = BertOnnxModel(model, num_heads=12, hidden_size=768)
return optimizer.get_fused_operator_statistics()
def _parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--input', required=True, type=str, help="input onnx model path")
@ -192,7 +207,7 @@ def parse_arguments():
return args
def get_optimization_options(args):
def _get_optimization_options(args):
optimization_options = BertOptimizationOptions(args.model_type)
if args.disable_gelu:
optimization_options.enable_gelu = False
@ -214,18 +229,38 @@ def get_optimization_options(args):
def optimize_model(input,
model_type,
num_heads,
hidden_size,
opt_level=0,
model_type='bert',
num_heads=12,
hidden_size=768,
optimization_options=None,
opt_level=0,
use_gpu=False,
only_onnxruntime=False):
""" Optimize Model by OnnxRuntime and/or offline fusion logic.
The following optimizes model by OnnxRuntime only, and no offline fusion logic:
optimize_model(input, opt_level=1, use_gpu=False, only_onnxruntime=True)
If you want to optimize model by offline fusion logic.
optimize_model(input, model_type, num_heads=12, hidden_size=768, optimization_options=your_options)
Args:
input (str): input model path.
model_type (str): model type - like bert, bert_tf, bert_keras or gpt2.
num_heads (int): number of attention heads.
hidden_size (int): hidden size.
optimization_options (OptimizationOptions or None): optimization options that can use to turn on/off some fusions.
opt_level (int): onnxruntime graph optimization level (0, 1, 2 or 99). When the level > 0, onnxruntime will be used to optimize model first.
use_gpu (bool): use gpu or not for onnxruntime.
only_onnxruntime (bool): only use onnxruntime to optimize model, and no offline fusion logic is used.
Returns:
object of an optimizer class.
"""
(optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type]
input_model_path = input
if opt_level > 1: # Optimization specified for an execution provider.
if opt_level > 1: # Optimization specified for an execution provider.
input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=use_gpu, opt_level=opt_level)
elif run_onnxruntime:
# Use Onnxruntime to do optimizations (like constant folding and cast elimation) that is not specified to exection provider.
@ -242,15 +277,15 @@ def optimize_model(input,
if optimization_options is None:
optimization_options = BertOptimizationOptions(model_type)
bert_model = optimizer_class(model, num_heads, hidden_size)
optimizer = optimizer_class(model, num_heads, hidden_size)
if not only_onnxruntime:
bert_model.optimize(optimization_options)
optimizer.optimize(optimization_options)
return bert_model
return optimizer
def setup_logger(verbose):
def _setup_logger(verbose):
if verbose:
coloredlogs.install(level='DEBUG', fmt='[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s')
else:
@ -258,33 +293,33 @@ def setup_logger(verbose):
def main():
args = parse_arguments()
args = _parse_arguments()
setup_logger(args.verbose)
_setup_logger(args.verbose)
optimization_options = get_optimization_options(args)
optimization_options = _get_optimization_options(args)
bert_model = optimize_model(args.input,
args.model_type,
args.num_heads,
args.hidden_size,
opt_level=args.opt_level,
optimization_options=optimization_options,
use_gpu=args.use_gpu,
only_onnxruntime=args.only_onnxruntime)
optimizer = optimize_model(args.input,
args.model_type,
args.num_heads,
args.hidden_size,
opt_level=args.opt_level,
optimization_options=optimization_options,
use_gpu=args.use_gpu,
only_onnxruntime=args.only_onnxruntime)
if args.float16:
bert_model.convert_model_float32_to_float16()
optimizer.convert_model_float32_to_float16()
if args.input_int32:
bert_model.change_input_to_int32()
optimizer.change_input_to_int32()
bert_model.save_model_to_file(args.output)
optimizer.save_model_to_file(args.output)
if bert_model.is_fully_optimized():
logger.info("The output model is fully optimized.")
if optimizer.is_fully_optimized():
logger.info("The model has been fully optimized.")
else:
logger.warning("The output model is not fully optimized. It might not be usable.")
logger.info("The model has been optimized.")
if __name__ == "__main__":

9
onnxruntime/python/tools/transformers/pytest.ini Normal file
View file

@ -0,0 +1,9 @@
[pytest]
log_cli = 1
log_cli_level = INFO
log_cli_format = %(message)s
log_file = pytest.log
log_file_level = DEBUG
log_file_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)
log_file_date_format=%Y-%m-%d %H:%M:%S

363
onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/generate_tiny_gpt2_model.py Normal file
View file

@ -0,0 +1,363 @@
#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
# This tool generates a tiny GPT2 model for testing fusion script.
# You can use benchmark_gpt2.py to get a gpt2 ONNX model as input of this tool.
import onnx
import onnx.utils
import sys
import argparse
import numpy as np
from onnx import ModelProto, TensorProto, numpy_helper
from OnnxModel import OnnxModel
import os
import onnxruntime
import random
from pathlib import Path
import timeit
DICT_SIZE = 20
SEQ_LEN = 5
""" This class creates a tiny bert model for test purpose. """
# parameters of input base model.
old_parameters = {
"seq_len": 5,
"hidden_size": 768,
"num_heads": 12,
"size_per_head": 64,
"word_dict_size": [50257], # list of supported dictionary size.
"max_word_position": 1024
}
# parameters of output tiny model.
new_parameters = {
"seq_len": SEQ_LEN,
"hidden_size": 4,
"num_heads": 2,
"size_per_head": 2,
"word_dict_size": DICT_SIZE,
"max_word_position": 8
}
class TinyGpt2Model(OnnxModel):
def __init__(self, model):
super(TinyGpt2Model, self).__init__(model)
self.resize_model()
def resize_weight(self, initializer_name, target_shape):
weight = self.get_initializer(initializer_name)
w = numpy_helper.to_array(weight)
target_w = w
if len(target_shape) == 1:
target_w = w[:target_shape[0]]
elif len(target_shape) == 2:
target_w = w[:target_shape[0], :target_shape[1]]
elif len(target_shape) == 3:
target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]]
elif len(target_shape) == 4:
target_w = w[:target_shape[0], :target_shape[1], :target_shape[2], :target_shape[3]]
else:
print("at most 3 dimensions")
tensor = onnx.helper.make_tensor(name=initializer_name + '_resize',
data_type=TensorProto.FLOAT,
dims=target_shape,
vals=target_w.flatten().tolist())
return tensor
def resize_model(self):
graph = self.model.graph
initializers = graph.initializer
for input in graph.input:
if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]):
print("input", input.name, input.type.tensor_type.shape)
input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"]
print("=>", input.type.tensor_type.shape)
reshapes = {}
for initializer in initializers:
tensor = numpy_helper.to_array(initializer)
if initializer.data_type == TensorProto.FLOAT:
dtype = np.float32
elif initializer.data_type == TensorProto.INT32:
dtype = np.int32
elif initializer.data_type == TensorProto.INT64:
dtype = np.int64
else:
print("data type not supported by this tool:", dtype)
if len(tensor.shape) == 1 and tensor.shape[0] == 1:
if tensor == old_parameters["num_heads"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]")
initializer.CopyFrom(
numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype),
initializer.name))
elif tensor == old_parameters["seq_len"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]")
initializer.CopyFrom(
numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name))
elif tensor == old_parameters["size_per_head"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]")
initializer.CopyFrom(
numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype),
initializer.name))
elif tensor == old_parameters["hidden_size"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]")
initializer.CopyFrom(
numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype),
initializer.name))
elif tensor == 4 * old_parameters["hidden_size"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]")
initializer.CopyFrom(
numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype),
initializer.name))
elif tensor == 3 * old_parameters["hidden_size"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
3 * old_parameters["hidden_size"], "=>[", 3 * new_parameters["hidden_size"], "]")
initializer.CopyFrom(
numpy_helper.from_array(np.asarray([3 * new_parameters["hidden_size"]], dtype=dtype),
initializer.name))
elif len(tensor.shape) == 0:
if tensor == old_parameters["num_heads"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["num_heads"], "=>", new_parameters["num_heads"])
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name))
elif tensor == old_parameters["seq_len"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["seq_len"], "=>", new_parameters["seq_len"])
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name))
elif tensor == old_parameters["size_per_head"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["size_per_head"], "=>", new_parameters["size_per_head"])
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype),
initializer.name))
elif tensor == old_parameters["hidden_size"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
old_parameters["hidden_size"], "=>", new_parameters["hidden_size"])
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype),
initializer.name))
elif tensor == 4 * old_parameters["hidden_size"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"])
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype),
initializer.name))
elif tensor == 3 * old_parameters["hidden_size"]:
print("initializer type={}".format(initializer.data_type), initializer.name,
3 * old_parameters["hidden_size"], "=>", 3 * new_parameters["hidden_size"])
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(3 * new_parameters["hidden_size"], dtype=dtype),
initializer.name))
elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]):
print("initializer type={}".format(initializer.data_type), initializer.name,
1.0 / np.sqrt(old_parameters["size_per_head"]), "=>",
1.0 / np.sqrt(new_parameters["size_per_head"]))
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype),
initializer.name))
elif tensor == np.sqrt(old_parameters["size_per_head"]):
print("initializer type={}".format(initializer.data_type), initializer.name,
np.sqrt(old_parameters["size_per_head"]), "=>", np.sqrt(new_parameters["size_per_head"]))
initializer.CopyFrom(
numpy_helper.from_array(np.asarray(np.sqrt(new_parameters["size_per_head"]), dtype=dtype),
initializer.name))
new_shape = []
shape_changed = False
for dim in tensor.shape:
if (dim == old_parameters["hidden_size"]):
new_shape.append(new_parameters["hidden_size"])
shape_changed = True
elif (dim == 4 * old_parameters["hidden_size"]):
new_shape.append(4 * new_parameters["hidden_size"])
shape_changed = True
elif (dim == 3 * old_parameters["hidden_size"]):
new_shape.append(3 * new_parameters["hidden_size"])
shape_changed = True
elif (dim in old_parameters["word_dict_size"]):
new_shape.append(new_parameters["word_dict_size"])
shape_changed = True
elif (dim == old_parameters["max_word_position"]):
new_shape.append(new_parameters["max_word_position"])
shape_changed = True
else:
new_shape.append(dim)
if shape_changed:
reshapes[initializer.name] = new_shape
print("initializer", initializer.name, tensor.shape, "=>", new_shape)
for initializer_name in reshapes:
self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize')
tensor = self.resize_weight(initializer_name, reshapes[initializer_name])
self.model.graph.initializer.extend([tensor])
# Add node name, replace split node attribute.
nodes_to_add = []
nodes_to_remove = []
for i, node in enumerate(graph.node):
if node.op_type == "Split":
nodes_to_add.append(
onnx.helper.make_node('Split',
node.input,
node.output,
name="Split_{}".format(i),
axis=2,
split=[
new_parameters["hidden_size"], new_parameters["hidden_size"],
new_parameters["hidden_size"]
]))
nodes_to_remove.append(node)
print("update split",
[new_parameters["hidden_size"], new_parameters["hidden_size"], new_parameters["hidden_size"]])
if node.op_type == "Constant":
for att in node.attribute:
if att.name == 'value':
if numpy_helper.to_array(att.t) == old_parameters["num_heads"]:
nodes_to_add.append(
onnx.helper.make_node('Constant',
inputs=node.input,
outputs=node.output,
value=onnx.helper.make_tensor(name=att.t.name,
data_type=TensorProto.INT64,
dims=[],
vals=[new_parameters["num_heads"]
])))
print("constant", att.t.name, old_parameters["num_heads"], "=>",
new_parameters["num_heads"])
if numpy_helper.to_array(att.t) == np.sqrt(old_parameters["size_per_head"]):
nodes_to_add.append(
onnx.helper.make_node('Constant',
inputs=node.input,
outputs=node.output,
value=onnx.helper.make_tensor(
name=att.t.name,
data_type=TensorProto.FLOAT,
dims=[],
vals=[np.sqrt(new_parameters["size_per_head"])])))
print("constant", att.t.name, np.sqrt(old_parameters["size_per_head"]), "=>",
np.sqrt(new_parameters["size_per_head"]))
else:
node.name = node.op_type + "_" + str(i)
for node in nodes_to_remove:
graph.node.remove(node)
graph.node.extend(nodes_to_add)
for i, input in enumerate(self.model.graph.input):
if i > 0:
dim_proto = input.type.tensor_type.shape.dim[2]
dim_proto.dim_value = new_parameters["num_heads"]
dim_proto = input.type.tensor_type.shape.dim[4]
dim_proto.dim_value = new_parameters["size_per_head"]
for i, output in enumerate(self.model.graph.output):
if i == 0:
dim_proto = output.type.tensor_type.shape.dim[2]
dim_proto.dim_value = new_parameters["hidden_size"]
if i > 0:
dim_proto = output.type.tensor_type.shape.dim[2]
dim_proto.dim_value = new_parameters["num_heads"]
dim_proto = output.type.tensor_type.shape.dim[4]
dim_proto.dim_value = new_parameters["size_per_head"]
def generate_test_data(onnx_file,
output_path,
batch_size=1,
use_cpu=True,
input_tensor_only=False,
dictionary_size=DICT_SIZE,
test_cases=1,
output_optimized_model=False):
for test_case in range(test_cases):
sequence_length = 3
input_1 = np.random.randint(dictionary_size, size=(batch_size, 1), dtype=np.int64)
tensor_1 = numpy_helper.from_array(input_1, 'input_ids')
path = os.path.join(output_path, 'test_data_set_' + str(test_case))
try:
os.mkdir(path)
except OSError:
print("Creation of the directory %s failed" % path)
else:
print("Successfully created the directory %s " % path)
sess_options = onnxruntime.SessionOptions()
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider'])
input1_name = sess.get_inputs()[0].name
output_names = [output.name for output in sess.get_outputs()]
inputs = {input1_name: input_1}
with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f:
f.write(tensor_1.SerializeToString())
for i in range(12):
input_name = f"past_{i}"
input = np.random.rand(2, batch_size, new_parameters["num_heads"], sequence_length,
new_parameters["size_per_head"]).astype(np.float32)
tensor = numpy_helper.from_array(input, input_name)
inputs.update({input_name: input})
with open(os.path.join(path, 'input_{}.pb'.format(1 + i)), 'wb') as f:
f.write(tensor.SerializeToString())
if input_tensor_only:
return
result = sess.run(output_names, inputs)
print("result 0 shape:", result[0].shape)
print("result 1 shape:", result[1].shape)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--input', required=True, type=str)
parser.add_argument('--output', required=True, type=str)
parser.add_argument('--output_optimized_model', required=False, action='store_true')
parser.set_defaults(output_optimized_model=False)
args = parser.parse_args()
model = ModelProto()
with open(args.input, "rb") as f:
model.ParseFromString(f.read())
bert_model = TinyGpt2Model(model)
bert_model.update_graph()
bert_model.remove_unused_constant()
print("opset verion", bert_model.model.opset_import[0].version)
with open(args.output, "wb") as out:
out.write(bert_model.model.SerializeToString())
p = Path(args.output)
data_path = p.parent
generate_test_data(args.output,
data_path,
batch_size=1,
use_cpu=True,
output_optimized_model=args.output_optimized_model)
if __name__ == "__main__":
main()

BIN
onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/gpt2_past.onnx Normal file
View file

Binary file not shown.

61
onnxruntime/python/tools/transformers/test_optimizer.py
View file

@ -12,6 +12,7 @@ import unittest
import os
import onnx
import onnxruntime
import pytest
from onnx import helper, TensorProto, ModelProto
from onnx.helper import make_node, make_tensor_value_info
import numpy as np
@ -20,15 +21,23 @@ from optimizer import optimize_model, optimize_by_onnxruntime
from OnnxModel import OnnxModel
BERT_TEST_MODELS = {
"bert_pytorch_0": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_0.onnx',
"bert_pytorch_1": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_1.onnx',
"bert_squad_pytorch1.4_opset10_fp32":
'test_data\\bert_squad_pytorch1.4_opset10_fp32\\BertForQuestionAnswering.onnx',
"bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx',
"bert_keras_squad": 'test_data\\bert_squad_tensorflow2.1_keras2onnx_opset11\\TFBertForQuestionAnswering.onnx',
"gpt2": 'test_data\\gpt2_pytorch1.4_opset11_no_past\\GPT2Model.onnx'
"bert_pytorch_0": ('bert_squad_pytorch1.4_opset11', 'BertForQuestionAnswering_0.onnx'),
"bert_pytorch_1": ('bert_squad_pytorch1.4_opset11', 'BertForQuestionAnswering_1.onnx'),
"bert_squad_pytorch1.4_opset10_fp32": ('bert_squad_pytorch1.4_opset10_fp32', 'BertForQuestionAnswering.onnx'),
"bert_keras_0": ('bert_mrpc_tensorflow2.1_opset10', 'TFBertForSequenceClassification_1.onnx'),
"bert_keras_squad": ('bert_squad_tensorflow2.1_keras2onnx_opset11', 'TFBertForQuestionAnswering.onnx'),
"gpt2": ('gpt2_pytorch1.4_opset11_no_past', 'GPT2Model.onnx'),
"gpt2_past": ('gpt2_pytorch1.5_opset11', 'gpt2_past.onnx'),
}
skip_on_ort_version = pytest.mark.skipif(onnxruntime.__version__.startswith('1.3.'),
reason="skip failed tests. TODO: fix them in 1.4.0.")
def _get_test_model_path(name):
sub_dir, file = BERT_TEST_MODELS[name]
return os.path.join('test_data', sub_dir, file)
class TestBertOptimization(unittest.TestCase):
def verify_node_count(self, bert_model, expected_node_count, test_name):
@ -39,8 +48,9 @@ class TestBertOptimization(unittest.TestCase):
print("{}: {} expected={}".format(op, len(bert_model.get_nodes_by_op_type(op)), counter))
self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count)
@skip_on_ort_version
def test_pytorch_model_0_cpu_onnxruntime(self):
input = BERT_TEST_MODELS['bert_pytorch_0']
input = _get_test_model_path('bert_pytorch_0')
output = 'temp.onnx'
optimize_by_onnxruntime(input, use_gpu=False, optimized_model_path=output)
model = ModelProto()
@ -58,12 +68,13 @@ class TestBertOptimization(unittest.TestCase):
}
self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_0_cpu_onnxruntime')
@skip_on_ort_version
def test_pytorch_model_0_gpu_onnxruntime(self):
if 'CUDAExecutionProvider' not in onnxruntime.get_available_providers():
print("skip test_pytorch_model_0_gpu_onnxruntime since no gpu found")
return
input = BERT_TEST_MODELS['bert_pytorch_0']
input = _get_test_model_path('bert_pytorch_0')
output = 'temp.onnx'
optimize_by_onnxruntime(input, use_gpu=True, optimized_model_path=output)
model = ModelProto()
@ -81,8 +92,9 @@ class TestBertOptimization(unittest.TestCase):
}
self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_0_gpu_onnxruntime')
@skip_on_ort_version
def test_pytorch_model_1_cpu_onnxruntime(self):
input = BERT_TEST_MODELS['bert_pytorch_1']
input = _get_test_model_path('bert_pytorch_1')
output = 'temp.onnx'
optimize_by_onnxruntime(input, use_gpu=False, optimized_model_path=output)
model = ModelProto()
@ -101,12 +113,13 @@ class TestBertOptimization(unittest.TestCase):
}
self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_1_cpu_onnxruntime')
@skip_on_ort_version
def test_pytorch_model_1_gpu_onnxruntime(self):
if 'CUDAExecutionProvider' not in onnxruntime.get_available_providers():
print("skip test_pytorch_model_1_gpu_onnxruntime since no gpu found")
return
input = BERT_TEST_MODELS['bert_pytorch_1']
input = _get_test_model_path('bert_pytorch_1')
output = 'temp.onnx'
optimize_by_onnxruntime(input, use_gpu=True, optimized_model_path=output)
model = ModelProto()
@ -125,8 +138,9 @@ class TestBertOptimization(unittest.TestCase):
}
self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_1_gpu_onnxruntime')
@skip_on_ort_version
def test_pytorch_model_0(self):
input = BERT_TEST_MODELS['bert_pytorch_0']
input = _get_test_model_path('bert_pytorch_0')
bert_model = optimize_model(input, 'bert', num_heads=2, hidden_size=8)
expected_node_count = {
@ -140,13 +154,13 @@ class TestBertOptimization(unittest.TestCase):
self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_0')
def test_pytorch_model_2(self):
input = BERT_TEST_MODELS['bert_squad_pytorch1.4_opset10_fp32']
input = _get_test_model_path('bert_squad_pytorch1.4_opset10_fp32')
bert_model = optimize_model(input, 'bert', num_heads=2, hidden_size=8)
print("fused_operator_statistics for test_pytorch_model_2", bert_model.get_fused_operator_statistics())
self.assertTrue(bert_model.is_fully_optimized())
def test_keras_model_1(self):
input = BERT_TEST_MODELS['bert_keras_0']
input = _get_test_model_path('bert_keras_0')
bert_model = optimize_model(input, 'bert_keras', num_heads=2, hidden_size=8)
@ -162,7 +176,7 @@ class TestBertOptimization(unittest.TestCase):
self.verify_node_count(bert_model, expected_node_count, 'test_keras_model_1')
def test_keras_squad_model(self):
input = BERT_TEST_MODELS['bert_keras_squad']
input = _get_test_model_path('bert_keras_squad')
bert_model = optimize_model(input, 'bert_keras', num_heads=2, hidden_size=8)
@ -171,7 +185,7 @@ class TestBertOptimization(unittest.TestCase):
self.assertTrue(bert_model.is_fully_optimized())
def test_gpt2(self):
input = BERT_TEST_MODELS['gpt2']
input = _get_test_model_path('gpt2')
bert_model = optimize_model(input, 'gpt2', num_heads=2, hidden_size=4)
expected_node_count = {
@ -185,6 +199,21 @@ class TestBertOptimization(unittest.TestCase):
}
self.verify_node_count(bert_model, expected_node_count, 'test_gpt2')
def test_gpt2_past(self):
input = _get_test_model_path('gpt2_past')
bert_model = optimize_model(input, 'gpt2', num_heads=2, hidden_size=4)
expected_node_count = {
'EmbedLayerNormalization': 0,
'Attention': 12,
'Gelu': 0,
'FastGelu': 12,
'BiasGelu': 0,
'LayerNormalization': 25,
'SkipLayerNormalization': 0
}
self.verify_node_count(bert_model, expected_node_count, 'test_gpt2_past')
if __name__ == '__main__':
unittest.main()

337
onnxruntime/test/contrib_ops/attention_op_test.cc
View file

@ -20,7 +20,12 @@ static void RunAttentionTest(
int hidden_size,
int number_of_heads,
bool use_float16 = false,
bool is_unidirectional = false) {
bool is_unidirectional = false,
bool use_past_state = false,
int past_sequence_length = 0,
int head_size = 0,
const std::vector<float>* past_data = nullptr,
const std::vector<float>* present_data = nullptr) {
int min_cuda_architecture = use_float16 ? 530 : 0;
bool enable_cuda = HasCudaEnvironment(min_cuda_architecture);
@ -35,6 +40,8 @@ static void RunAttentionTest(
std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};
std::vector<int64_t> bias_dims = {3 * hidden_size};
std::vector<int64_t> mask_index_dims = {batch_size};
std::vector<int64_t> past_dims = {2, batch_size, head_size, past_sequence_length, head_size};
std::vector<int64_t> present_dims = {2, batch_size, head_size, past_sequence_length + sequence_length, head_size};
std::vector<int64_t> output_dims = input_dims;
if (use_float16) {
@ -51,6 +58,23 @@ static void RunAttentionTest(
if (mask_index_data.size() > 0) { // mask index is optional.
tester.AddInput<int32_t>("mask_index", mask_index_dims, mask_index_data);
} else {
std::vector<int64_t> dims = {static_cast<int64_t>(mask_index_data.size())};
tester.AddInput<int32_t>("", dims, mask_index_data);
}
if (use_past_state) {
if (use_float16) {
if (past_sequence_length > 0) {
tester.AddInput<MLFloat16>("past", past_dims, ToFloat16(*past_data));
}
tester.AddOutput<MLFloat16>("present", present_dims, ToFloat16(*present_data));
} else {
if (past_sequence_length > 0) {
tester.AddInput<float>("past", past_dims, *past_data);
}
tester.AddOutput<float>("present", present_dims, *present_data);
}
}
if (enable_cuda) {
@ -256,8 +280,7 @@ TEST(AttentionTest, AttentionUnidirectional) {
std::vector<float> input_data = {
0.091099896f, -0.018294459f, -0.36594841f, 0.28410032f,
-0.12125026f, -0.0066160089f, 0.38809127f, -0.22455512f
};
-0.12125026f, -0.0066160089f, 0.38809127f, -0.22455512f};
std::vector<float> weight_data = {
-0.2659236192703247f,
@ -310,8 +333,7 @@ TEST(AttentionTest, AttentionUnidirectional) {
-0.09368397295475006f,
0.07878211885690689f,
0.2973634898662567f,
0.11210034042596817f
};
0.11210034042596817f};
std::vector<float> bias_data = {
-0.0540979839861393f,
@ -325,20 +347,317 @@ TEST(AttentionTest, AttentionUnidirectional) {
0.3670335114002228f,
0.028461361303925514f,
-0.08913630992174149f,
0.28048714995384216f
};
0.28048714995384216f};
// No mask_index
std::vector<int32_t> mask_index_data = {};
std::vector<float> output_data = {
0.28109729f, 0.069518551f, 0.0038009658f, 0.29213354f, 0.3692801f, 0.029495837f, -0.084964074f, 0.28169215f
};
0.28109729f, 0.069518551f, 0.0038009658f, 0.29213354f, 0.3692801f, 0.029495837f, -0.084964074f, 0.28169215f};
bool is_unidirectional = true;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional);
}
TEST(AttentionTest, AttentionEmptyPastState) {
int batch_size = 1;
int sequence_length = 2;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
0.091099896f, -0.018294459f, -0.36594841f, 0.28410032f,
-0.12125026f, -0.0066160089f, 0.38809127f, -0.22455512f};
std::vector<float> weight_data = {
-0.2659236192703247f,
0.02789675071835518f,
0.07280516624450684f,
0.050951678305864334f,
0.020417947322130203f,
-0.04751841351389885f,
0.043815530836582184f,
0.006015353370457888f,
-0.11496957391500473f,
-0.1773347705602646f,
0.30928605794906616f,
0.005648412741720676f,
0.08960387855768204f,
-0.27270448207855225f,
0.14847396314144135f,
-0.17960812151432037f,
0.01788954995572567f,
0.09993876516819f,
0.03943513706326485f,
-0.02484400011599064f,
-0.12958766520023346f,
0.220433309674263f,
0.1720484346151352f,
0.22024005651474f,
0.059368450194597244f,
0.1710093915462494f,
-0.3967452347278595f,
-0.1591450721025467f,
0.1446179747581482f,
-0.20505407452583313f,
0.12749597430229187f,
0.32139700651168823f,
0.139958456158638f,
-0.10619817674160004f,
0.04528557509183884f,
0.045598603785037994f,
-0.007152545265853405f,
0.109454445540905f,
-0.1582530289888382f,
-0.2646341919898987f,
0.0920850858092308f,
0.0701494812965393f,
-0.19062495231628418f,
-0.24360455572605133f,
-0.09368397295475006f,
0.07878211885690689f,
0.2973634898662567f,
0.11210034042596817f};
std::vector<float> bias_data = {
-0.0540979839861393f,
-0.06444740295410156f,
0.03112877532839775f,
-0.08288222551345825f,
0.07840359210968018f,
0.039143580943346024f,
-0.45591455698013306f,
-0.11876055598258972f,
0.3670335114002228f,
0.028461361303925514f,
-0.08913630992174149f,
0.28048714995384216f};
// No mask_index
std::vector<int32_t> mask_index_data = {};
std::vector<float> output_data = {
0.28109729f, 0.069518551f, 0.0038009658f, 0.29213354f, 0.3692801f, 0.029495837f, -0.084964074f, 0.28169215f};
std::vector<float> past_data = {};
std::vector<float> present_data = {
0.053175069391727448f, 0.12795503437519073f, 0.11125634610652924f, -0.0510881207883358f, -0.55345797538757324f, -0.3045809268951416f, -0.36920222640037537f, 0.060108467936515808f, 0.28109729290008545f, 0.069518551230430603f, 0.45718482136726379f, -0.010400654748082161f, 0.0038009658455848694f, 0.29213353991508484f, -0.17697516083717346f, 0.27086889743804932f};
bool is_unidirectional = true;
bool use_past_state = true;
int past_sequence_length = 0;
int head_size = 2;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, head_size, &past_data, &present_data);
}
TEST(AttentionTest, AttentionPastStateBatch1) {
int batch_size = 1;
int sequence_length = 1;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
-0.019333266f, -0.21813886f, 0.16212955f, -0.015626367f};
std::vector<float> weight_data = {
-0.4738484025001526f,
-0.2613658607006073f,
-0.0978037416934967f,
-0.34988933801651f,
0.2243240624666214f,
-0.0429205559194088f,
0.418695330619812f,
0.17441125214099884f,
-0.18825532495975494f,
0.18357256054878235f,
-0.5806483626365662f,
-0.02251487597823143f,
0.08742205798625946f,
0.14734269678592682f,
0.2387014478445053f,
0.2884027063846588f,
0.6490834355354309f,
0.16965825855731964f,
-0.06346885114908218f,
0.4073973298072815f,
-0.03070945478975773f,
0.4110257923603058f,
0.07896808534860611f,
0.16783113777637482f,
0.0038893644232302904f,
0.06946629285812378f,
0.36680519580841064f,
-0.07261059433221817f,
-0.14960581064224243f,
0.020944256335496902f,
-0.09378612786531448f,
-0.1336742341518402f,
0.06061394885182381f,
0.2205914407968521f,
-0.03519909828901291f,
-0.18405692279338837f,
0.22149960696697235f,
-0.1884360909461975f,
-0.014074507169425488f,
0.4252440333366394f,
0.24987126886844635f,
-0.31396418809890747f,
0.14036843180656433f,
0.2854192554950714f,
0.09709841012954712f,
0.09935075044631958f,
-0.012154420837759972f,
0.2575816512107849f};
std::vector<float> bias_data = {
0.4803391396999359f,
-0.5254325866699219f,
-0.42926454544067383f,
-0.2059524953365326f,
-0.12773379683494568f,
-0.09542735666036606f,
-0.35286077857017517f,
-0.07646317780017853f,
-0.04590314254164696f,
-0.03752850368618965f,
-0.013764488510787487f,
-0.18478283286094666f};
// No mask_index
std::vector<int32_t> mask_index_data = {};
std::vector<float> output_data = {
0.20141591f, 0.43005896f, 0.35745093f, 0.19957167f};
std::vector<float> past_data = {
0.55445826f, 0.10127074f, 0.71770734f, 0.15915526f, 0.13913247f, 0.77447522f, 0.66044068f, 0.27559045f, 0.35731629f, 0.62033528f, 0.24354559f, 0.22859341f,
0.45075402f, 0.85365993f, 0.097346395f, 0.28859729f, 0.26926181f, 0.65922296f, 0.8177433f, 0.4212271f, 0.34352475f, 0.059609573f, 0.46556228f, 0.7226882f};
std::vector<float> present_data = {
0.55445826f, 0.10127074f, 0.71770734f, 0.15915526f, 0.13913247f, 0.77447522f, -0.30182117f, -0.12330482f, 0.66044068f, 0.27559045f, 0.35731629f, 0.62033528f, 0.24354559f, 0.22859341f, -0.36450946f, -0.19483691f,
0.45075402f, 0.85365993f, 0.097346395f, 0.28859729f, 0.26926181f, 0.65922296f, -0.027254611f, -0.096526355f, 0.8177433f, 0.4212271f, 0.34352475f, 0.059609573f, 0.46556228f, 0.7226882f, -0.025281552f, -0.25482416f};
bool is_unidirectional = true;
bool use_past_state = true;
int past_sequence_length = 3;
int head_size = 2;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, head_size, &past_data, &present_data);
}
TEST(AttentionTest, AttentionPastStateBatch2) {
int batch_size = 2;
int sequence_length = 1;
int hidden_size = 4;
int number_of_heads = 2;
std::vector<float> input_data = {
-0.10902753f, 0.0041178204f, 0.1871525f, -0.20399982f,
0.027207348f, -0.25321805f, 0.12869114f, 0.023136809f};
std::vector<float> weight_data = {
-0.4738484025001526f,
-0.2613658607006073f,
-0.0978037416934967f,
-0.34988933801651f,
0.2243240624666214f,
-0.0429205559194088f,
0.418695330619812f,
0.17441125214099884f,
-0.18825532495975494f,
0.18357256054878235f,
-0.5806483626365662f,
-0.02251487597823143f,
0.08742205798625946f,
0.14734269678592682f,
0.2387014478445053f,
0.2884027063846588f,
0.6490834355354309f,
0.16965825855731964f,
-0.06346885114908218f,
0.4073973298072815f,
-0.03070945478975773f,
0.4110257923603058f,
0.07896808534860611f,
0.16783113777637482f,
0.0038893644232302904f,
0.06946629285812378f,
0.36680519580841064f,
-0.07261059433221817f,
-0.14960581064224243f,
0.020944256335496902f,
-0.09378612786531448f,
-0.1336742341518402f,
0.06061394885182381f,
0.2205914407968521f,
-0.03519909828901291f,
-0.18405692279338837f,
0.22149960696697235f,
-0.1884360909461975f,
-0.014074507169425488f,
0.4252440333366394f,
0.24987126886844635f,
-0.31396418809890747f,
0.14036843180656433f,
0.2854192554950714f,
0.09709841012954712f,
0.09935075044631958f,
-0.012154420837759972f,
0.2575816512107849f};
std::vector<float> bias_data = {
0.4803391396999359f,
-0.5254325866699219f,
-0.42926454544067383f,
-0.2059524953365326f,
-0.12773379683494568f,
-0.09542735666036606f,
-0.35286077857017517f,
-0.07646317780017853f,
-0.04590314254164696f,
-0.03752850368618965f,
-0.013764488510787487f,
-0.18478283286094666f};
// No mask_index
std::vector<int32_t> mask_index_data = {};
std::vector<float> output_data = {
0.14902574f, 0.62273371f, 0.43022552f, 0.12759127f,
0.26993567f, 0.23553593f, 0.43190649f, 0.086044826f};
std::vector<float> past_data = {
0.42028648f, 0.55855948f, 0.044569403f, 0.76525789f, 0.13962431f, 0.40977913f, 0.36911047f, 0.83399564f, 0.36905321f, 0.91414654f, 0.17300875f, 0.78793788f,
0.10279467f, 0.80501258f, 0.089550517f, 0.85371113f, 0.61801594f, 0.91222942f, 0.88626182f, 0.069776468f, 0.10591964f, 0.84836882f, 0.83520192f, 0.0098680854f,
0.3113814f, 0.63999802f, 0.28603253f, 0.98899829f, 0.044405211f, 0.95105386f, 0.81278932f, 0.63969064f, 0.14494057f, 0.11349615f, 0.87086016f, 0.20983537f,
0.35107401f, 0.90144604f, 0.68950737f, 0.18928574f, 0.18029204f, 0.074517399f, 0.70763874f, 0.48440042f, 0.58114725f, 0.1048766f, 0.73694098f, 0.17766342f};
std::vector<float> present_data = {
0.42028648f, 0.55855948f, 0.044569403f, 0.76525789f, 0.13962431f, 0.40977913f, -0.22849128f, -0.022080801f, 0.36911047f, 0.83399564f, 0.36905321f, 0.91414654f, 0.17300875f, 0.78793788f, -0.4449589f, -0.17704415f, 0.10279467f, 0.80501258f, 0.089550517f, 0.85371113f, 0.61801594f, 0.91222942f, -0.2994619f, -0.14412443f, 0.88626182f, 0.069776468f, 0.10591964f, 0.84836882f, 0.83520192f, 0.0098680854f, -0.33421949f, -0.18547727f,
0.3113814f, 0.63999802f, 0.28603253f, 0.98899829f, 0.044405211f, 0.95105386f, -0.033968594f, -0.034833729f, 0.81278932f, 0.63969064f, 0.14494057f, 0.11349615f, 0.87086016f, 0.20983537f, 0.045759238f, -0.26863033f, 0.35107401f, 0.90144604f, 0.68950737f, 0.18928574f, 0.18029204f, 0.074517399f, -0.033201858f, -0.10592631f, 0.70763874f, 0.48440042f, 0.58114725f, 0.1048766f, 0.73694098f, 0.17766342f, -0.054369561f, -0.24562015f};
bool is_unidirectional = true;
bool use_past_state = true;
int past_sequence_length = 3;
int head_size = 2;
RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional,
use_past_state, past_sequence_length, head_size, &past_data, &present_data);
}
} // namespace test
} // namespace onnxruntime