Add quantization support for GPT2 past state and use model to generate outputs in OpTester (#4340)

* Make quantization support GPT2 past state
* Make OpTester to be able to generate reference outputs with a model. With it, there is no need to compute outputs manually, which are impossible for some cases.
This commit is contained in:
Yufeng Li 2020-06-26 09:29:29 -07:00 committed by GitHub
parent ceedf126a2
commit fc5e65a22d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
22 changed files with 591 additions and 625 deletions

View file

@ -44,7 +44,7 @@ Status AttentionBase::CheckInputs(const Tensor* input,
const auto& dims = input->Shape().GetDims();
if (dims.size() != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 0 is expected to have 3 dimensions, got ",
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'input' is expected to have 3 dimensions, got ",
dims.size());
}
int batch_size = static_cast<int>(dims[0]);
@ -56,7 +56,7 @@ Status AttentionBase::CheckInputs(const Tensor* input,
const auto& weights_dims = weights->Shape().GetDims();
if (weights_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 1 is expected to have 2 dimensions, got ",
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' is expected to have 2 dimensions, got ",
weights_dims.size());
}
if (weights_dims[0] != dims[2]) {
@ -64,12 +64,12 @@ Status AttentionBase::CheckInputs(const Tensor* input,
"Input 1 dimension 0 should have same length as dimension 2 of input 0");
}
if (weights_dims[1] != 3 * weights_dims[0]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 1 dimension 1 should be 3 times of dimension 0");
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' dimension 1 should be 3 times of dimension 0");
}
const auto& bias_dims = bias->Shape().GetDims();
if (bias_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 2 is expected to have 1 dimension, got ",
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' is expected to have 1 dimension, got ",
bias_dims.size());
}
if (bias_dims[0] != weights_dims[1]) {
@ -80,40 +80,40 @@ Status AttentionBase::CheckInputs(const Tensor* input,
if (mask_index != nullptr) { // mask_index is optional
// unidirectional (like GPT2) does not need mask input. Here we do not allowed the input for unidirectional.
if (is_unidirectional_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 3 (mask_index) is not allowed for unidirectional");
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask_index' is not allowed for unidirectional");
}
const auto& mask_dims = mask_index->Shape().GetDims();
if (mask_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 3 is expected to have 1 dimension, got ",
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'mask_index' is expected to have 1 dimension, got ",
mask_dims.size());
}
if (static_cast<int>(mask_dims[0]) != batch_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 3 and 0 shall have same length at dimension 0");
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' and 'input' shall have same length at dimension 0");
}
}
if (past != nullptr) { // past is optional
if (!is_unidirectional_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 (past) is only allowed for unidirectional");
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input '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 ",
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past' 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");
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' 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");
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' 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_);
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' 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 ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'past' dimension 2 shall have length of ", hidden_size / num_heads_);
}
}
@ -147,7 +147,7 @@ Tensor* AttentionBase::GetPresent(OpKernelContext* context,
}
template <typename T>
Attention<T>::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionBase(info) {
Attention<T>::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) {
}
template <typename T>
@ -168,19 +168,15 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
Tensor* output = context->Output(0, 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;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));
auto* tp = context->GetOperatorThreadPool();
// STEP.1: gemm_data(BS, 3NH) = input(BS, NH) x weights(NH, 3NH) + bias(3NH)
// Compute Q, K, V
// gemm_data(BS, 3NH) = input(BS, NH) x weights(NH, 3NH) + bias(3NH)
auto gemm_data = allocator->Alloc(SafeInt<size_t>(batch_size) * sequence_length * 3 * hidden_size * element_size);
BufferUniquePtr gemm_buffer(gemm_data, BufferDeleter(allocator));
auto Q = reinterpret_cast<T*>(gemm_data);
@ -235,51 +231,15 @@ Status Attention<T>::Compute(OpKernelContext* context) const {
qkv_dest + qkv_offset, // C
head_size, // ldc
nullptr // use single-thread
);
);
}
});
}
// 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 * 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 * all_sequence_length * element_size;
} else if (is_unidirectional_) {
mask_data_bytes = SafeInt<size_t>(sequence_length) * all_sequence_length * element_size;
}
void* mask_data = nullptr;
if (mask_data_bytes > 0) {
mask_data = allocator->Alloc(mask_data_bytes);
memset(mask_data, 0, mask_data_bytes);
}
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, 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 =
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, past_sequence_length, head_size, num_heads_, hidden_size,
past_data, present_data, tp);
return Status::OK();
// Compute the attention score and apply the score to V
return ApplyAttention(Q, K, V, mask_index, past, output,
batch_size, sequence_length,
head_size, hidden_size, context);
}
} // namespace contrib

View file

@ -3,34 +3,15 @@
#pragma once
#include "attention_cpu_base.h"
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
class AttentionBase {
protected:
AttentionBase(const OpKernelInfo& info);
Status CheckInputs(const Tensor* input,
const Tensor* weights,
const Tensor* bias,
const Tensor* mask_index,
const Tensor* past) const;
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>
class Attention : public OpKernel, public AttentionBase {
class Attention : public OpKernel, public AttentionCPUBase {
public:
explicit Attention(const OpKernelInfo& info);
Status Compute(OpKernelContext* context) const override;

View file

@ -0,0 +1,33 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
class AttentionBase {
protected:
AttentionBase(const OpKernelInfo& info);
Status CheckInputs(const Tensor* input,
const Tensor* weights,
const Tensor* bias,
const Tensor* mask_index,
const Tensor* past) const;
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.
};
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,218 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "attention_base.h"
#include "attention_helper.h"
#include "core/common/common.h"
#include "core/common/safeint.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
class AttentionCPUBase : public AttentionBase {
protected:
AttentionCPUBase(const OpKernelInfo& info) : AttentionBase(info) {}
template <typename T>
Status ApplyAttention(const T* Q, // Q data. Its size is BxNxSxH
const T* K, // K data. Its size is BxNxSxH
const T* V, // V value with size BxNxSxH
const Tensor* mask_index, // mask index. nullptr if no mask or its size is B
const Tensor* past, // past state
Tensor* output, // output tensor
int batch_size, // batch size
int sequence_length, // sequence length
int head_size, // head size
int hidden_size, // hidden size
OpKernelContext* context) const {
AllocatorPtr allocator;
ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));
auto* tp = context->GetOperatorThreadPool();
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;
// 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 * all_sequence_length * sizeof(T);
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 * all_sequence_length * sizeof(T);
} else if (is_unidirectional_) {
mask_data_bytes = SafeInt<size_t>(sequence_length) * all_sequence_length * sizeof(T);
}
void* mask_data = nullptr;
if (mask_data_bytes > 0) {
mask_data = allocator->Alloc(mask_data_bytes);
memset(mask_data, 0, mask_data_bytes);
}
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, past_sequence_length, head_size,
past_data, present_data, tp);
// 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 * sizeof(T));
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, past_sequence_length, head_size, hidden_size,
past_data, present_data, tp);
return Status::OK();
}
private:
// 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: 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
const T* past, // past state
T* present, // present state
ThreadPool* tp) const {
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) {
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 * 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 * 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*
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 * 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 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)
{
const int N = batch_size * num_heads_ * 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 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 hidden_size, // hidden size
const T* past, // past state
T* present, // present state
ThreadPool* tp) const {
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) {
for (std::ptrdiff_t i = begin; i != end; ++i) {
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_);
const int head_index = static_cast<int>(i % num_heads_);
T* src = current_tmp_data;
T* dest = output + (batch_index * sequence_length * num_heads_ + head_index) * head_size;
const auto bytes_to_copy = SafeInt<size_t>(head_size) * sizeof(T);
for (int j = 0; j < sequence_length; j++) {
memcpy(dest, src, bytes_to_copy);
src += head_size;
dest += hidden_size;
}
}
});
}
};
} // namespace contrib
} // namespace onnxruntime

View file

@ -68,7 +68,7 @@ void PrepareMask(const int32_t* mask_index,
int past_sequence_length) {
const int all_sequence_length = past_sequence_length + sequence_length;
T* p_mask = mask_data;
if (is_unidirectional) {
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++) {
@ -97,7 +97,7 @@ void PrepareMask(const int32_t* mask_index,
// 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>
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;
@ -112,138 +112,5 @@ T* ConcatStateChunk(const T* past, const T* chunk, T* present, size_t past_chunk
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: 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) {
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 * 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 * 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*
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 * 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 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)
{
const int N = batch_size * num_heads * 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 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) {
for (std::ptrdiff_t i = begin; i != end; ++i) {
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);
const int head_index = static_cast<int>(i % num_heads);
T* src = current_tmp_data;
T* dest = output + (batch_index * sequence_length * num_heads + head_index) * head_size;
const auto bytes_to_copy = SafeInt<size_t>(head_size) * sizeof(T);
for (int j = 0; j < sequence_length; j++) {
memcpy(dest, src, bytes_to_copy);
src += head_size;
dest += hidden_size;
}
}
});
}
} // namespace contrib
} // namespace onnxruntime

View file

@ -34,21 +34,23 @@ REGISTER_KERNEL_TYPED(float, uint8_t, int8_t)
REGISTER_KERNEL_TYPED(float, uint8_t, uint8_t)
template <typename T, typename QInput, typename QWeight>
QAttention<T, QInput, QWeight>::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionBase(info) {
QAttention<T, QInput, QWeight>::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) {
}
template <typename T, typename QInput, typename QWeight>
Status QAttention<T, QInput, QWeight>::Compute(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 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 5 - mask_index : (batch_size)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Output : (batch_size, sequence_length, hidden_size)
// 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 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 5 - mask_index : (batch_size)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Input 8 - past : (2, batch_size, num_heads, past_sequence_length, head_size)
// Output 0 : (batch_size, sequence_length, hidden_size)
// Output 1 - present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)
// ORT_RETURN_IF_ERROR(CheckInputs(context));
const Tensor* input = context->Input<Tensor>(0);
const Tensor* weights = context->Input<Tensor>(1);
@ -58,8 +60,9 @@ Status QAttention<T, QInput, QWeight>::Compute(OpKernelContext* context) const {
const Tensor* mask_index = context->Input<Tensor>(5);
const Tensor* i_zp_tensor = context->Input<Tensor>(6);
const Tensor* w_zp_tensor = context->Input<Tensor>(7);
const Tensor* past_tensor = context->Input<Tensor>(8);
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr));
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, past_tensor));
ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor),
"input scale must be a scalar or 1D tensor of size 1");
@ -146,51 +149,15 @@ Status QAttention<T, QInput, QWeight>::Compute(OpKernelContext* context) const {
&dequant_scale, // output scale
bias_data + weights_offset, // bias
nullptr // use single-thread
);
);
}
});
}
// 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;
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;
} else if (is_unidirectional_) {
mask_data_bytes = SafeInt<size_t>(sequence_length) * sequence_length * element_size;
}
void* mask_data = nullptr;
if (mask_data_bytes > 0) {
mask_data = allocator->Alloc(mask_data_bytes);
memset(mask_data, 0, mask_data_bytes);
}
BufferUniquePtr mask_data_buffer(mask_data, BufferDeleter(allocator));
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, 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 =
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, past_sequence_length, head_size, num_heads_, hidden_size, past_data, present_data, tp);
return Status::OK();
// Compute the attention score and apply the score to V
return ApplyAttention(Q, K, V, mask_index, past_tensor, output,
batch_size, sequence_length,
head_size, hidden_size, context);
}
} // namespace contrib

View file

@ -2,12 +2,12 @@
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "contrib_ops/cpu/bert/attention.h"
#include "contrib_ops/cpu/bert/attention_cpu_base.h"
namespace onnxruntime {
namespace contrib {
template <typename T, typename QInput, typename QWeight>
class QAttention : public OpKernel, public AttentionBase {
class QAttention : public OpKernel, public AttentionCPUBase {
public:
QAttention(const OpKernelInfo& info);

View file

@ -6,7 +6,7 @@
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/providers/cuda/cuda_common.h"
#include "contrib_ops/cpu/bert/attention.h"
#include "contrib_ops/cpu/bert/attention_base.h"
namespace onnxruntime {
namespace contrib {

View file

@ -47,7 +47,8 @@ Status QAttention<T, int8_t>::CheckInputs(const Tensor* input,
const Tensor* weight_scale_tensor,
const Tensor* mask_index,
const Tensor* i_zp_tensor,
const Tensor* w_zp_tensor) const {
const Tensor* w_zp_tensor,
const Tensor* past_tensor) const {
// Input and output shapes:
// Input 0 - input : (batch_size, sequence_length, hidden_size)
// Input 1 - weights : (hidden_size, 3 * hidden_size)
@ -59,7 +60,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, nullptr));
ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, past_tensor));
ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor),
"input scale must be a scalar or 1D tensor of size 1");
@ -88,15 +89,17 @@ Status QAttention<T, int8_t>::CheckInputs(const Tensor* input,
template <typename T>
Status QAttention<T, int8_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 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 5 - mask_index : (batch_size)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Output : (batch_size, sequence_length, hidden_size)
// 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 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 5 - mask_index : (batch_size)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Input 8 - past : (2, batch_size, num_heads, past_sequence_length, head_size)
// Output 0 - output : (batch_size, sequence_length, hidden_size)
// Output 1 - present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)
// ORT_RETURN_IF_ERROR(CheckInputs(context));
const Tensor* input = context->Input<Tensor>(0);
const Tensor* weights = context->Input<Tensor>(1);
@ -106,6 +109,7 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
const Tensor* mask_index = context->Input<Tensor>(5);
const Tensor* i_zp_tensor = context->Input<Tensor>(6);
const Tensor* w_zp_tensor = context->Input<Tensor>(7);
const Tensor* past_tensor = context->Input<Tensor>(8);
ORT_RETURN_IF_ERROR(CheckInputs(input,
weights,
@ -114,7 +118,8 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
weight_scale_tensor,
mask_index,
i_zp_tensor,
w_zp_tensor));
w_zp_tensor,
past_tensor));
const auto& shape = input->Shape();
int batch_size = static_cast<int>(shape[0]);
@ -160,9 +165,9 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
m,
n);
const int past_sequence_length = 0;
const T* past_data = nullptr;
T* present_data = nullptr;
int past_sequence_length = 0;
Tensor* present_tensor = GetPresent(context, past_tensor, batch_size, head_size, sequence_length, past_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(
@ -178,9 +183,8 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
element_size,
is_unidirectional_,
past_sequence_length,
past_data,
present_data
)) {
nullptr == past_tensor ? nullptr : past_tensor->template Data<T>(),
nullptr == present_tensor ? nullptr : present_tensor->template MutableData<T>())) {
// Get last error to reset it to cudaSuccess.
CUDA_CALL(cudaGetLastError());
return Status(common::ONNXRUNTIME, common::FAIL);

View file

@ -36,7 +36,8 @@ class QAttention<T, int8_t> final : public CudaKernel, public AttentionBase {
const Tensor* weight_scale_tensor,
const Tensor* mask_index,
const Tensor* i_zp_tensor,
const Tensor* w_zp_tensor) const;
const Tensor* w_zp_tensor,
const Tensor* past_tensor) const;
};
} // namespace cuda

View file

@ -406,11 +406,23 @@ mask_index shall not be provided.)DOC";
"zero point of quantized weight tensor. It's a scalar, which means a per-tensor/layer quantization.",
"T2",
OpSchema::Optional)
.Input(
8,
"past",
"past state for key and value with shape (2, batch_size, num_heads, past_sequence_length, head_size).",
"T3",
OpSchema::Optional)
.Output(
0,
"output",
"3D output tensor with shape (batch_size, sequence_length, hidden_size)",
"T3")
.Output(
1,
"present",
"present state for key and value with shape (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)",
"T3",
OpSchema::Optional)
.TypeConstraint("T1", {"tensor(int8)", "tensor(uint8)"}, "Constrain input and output types to int8 tensors.")
.TypeConstraint("T2", {"tensor(int8)", "tensor(uint8)"}, "Constrain input and output types to int8 tensors.")
.TypeConstraint("T3", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.")
@ -2616,7 +2628,9 @@ Example 4:
.Attr("axis",
"The first normalization dimension: normalization will be performed along dimensions axis : rank(inputs).",
AttributeProto::INT, static_cast<int64_t>(-1))
.Attr("epsilon", "The epsilon value to use to avoid division by zero.", AttributeProto::FLOAT, 1e-5f)
.Attr("epsilon",
"The epsilon value to use to avoid division by zero.",
AttributeProto::FLOAT, 1e-5f)
.AllowUncheckedAttributes()
.Input(0, "X", "Input data tensor from the previous layer.", "T")
.Input(1, "scale", "Scale tensor.", "T")

View file

@ -1361,8 +1361,9 @@ class ONNXQuantizer:
inputs.extend(quantized_input_names)
inputs.extend([node.input[2]])
inputs.extend(scale_names)
inputs.extend([node.input[3]])
inputs.extend([node.input[3] if len(node.input) > 3 else ""])
inputs.extend(zero_point_names)
inputs.extend([node.input[4] if len(node.input) > 4 else ""])
kwargs = {}
for attribute in node.attribute:

View file

@ -57,6 +57,49 @@ class RandomValueGenerator {
return val;
}
// Gaussian distribution for float
template <typename TFloat>
typename std::enable_if<
std::is_floating_point<TFloat>::value,
std::vector<TFloat>>::type
Gaussian(const std::vector<int64_t>& dims, TFloat mean, TFloat stddev) {
std::vector<TFloat> val(detail::SizeFromDims(dims));
std::normal_distribution<TFloat> distribution(mean, stddev);
for (size_t i = 0; i < val.size(); ++i) {
val[i] = distribution(generator_);
}
return val;
}
// Gaussian distribution for Integer
template <typename TInt>
typename std::enable_if<
std::is_integral<TInt>::value,
std::vector<TInt>>::type
Gaussian(const std::vector<int64_t>& dims, TInt mean, TInt stddev) {
std::vector<TInt> val(detail::SizeFromDims(dims));
std::normal_distribution<float> distribution(static_cast<float>(mean), static_cast<float>(stddev));
for (size_t i = 0; i < val.size(); ++i) {
val[i] = static_cast<TInt>(std::round(distribution(generator_)));
}
return val;
}
// Gaussian distribution for Integer and Clamp to [min, max]
template <typename TInt>
typename std::enable_if<
std::is_integral<TInt>::value,
std::vector<TInt>>::type
Gaussian(const std::vector<int64_t>& dims, TInt mean, TInt stddev, TInt min, TInt max) {
std::vector<TInt> val(detail::SizeFromDims(dims));
std::normal_distribution<float> distribution(static_cast<float>(mean), static_cast<float>(stddev));
for (size_t i = 0; i < val.size(); ++i) {
int64_t round_val = static_cast<int64_t>(std::round(distribution(generator_)));
val[i] = static_cast<TInt>(std::min<int64_t>(std::max<int64_t>(round_val, min), max));
}
return val;
}
template <class T>
inline std::vector<T> OneHot(const std::vector<int64_t>& dims, int64_t stride) {
std::vector<T> val(detail::SizeFromDims(dims), T(0));

View file

@ -659,5 +659,34 @@ TEST(AttentionTest, AttentionPastStateBatch2) {
use_past_state, past_sequence_length, head_size, &past_data, &present_data);
}
TEST(AttentionTest, AttentionPastState_dynamic) {
// create rand inputs
RandomValueGenerator random{};
std::vector<int64_t> input_dims{2, 5, 768};
std::vector<float> input_data = random.Gaussian<float>(input_dims, 0.0f, 0.3f);
std::vector<int64_t> weight_dims{768, 2304};
std::vector<float> weight_data = random.Gaussian<float>(weight_dims, 0.0f, 0.3f);
std::vector<int64_t> bias_dims{2304};
std::vector<float> bias_data = random.Gaussian<float>(bias_dims, 0.0f, 0.3f);
std::vector<int64_t> past_dims{2, 2, 12, 15, 64};
std::vector<float> past_data = random.Gaussian<float>(past_dims, 0.0f, 0.3f);
OpTester test("Attention", 1, onnxruntime::kMSDomain);
test.AddAttribute<int64_t>("num_heads", 12);
test.AddAttribute<int64_t>("unidirectional", 1);
test.AddInput<float>("input", input_dims, input_data);
test.AddInput<float>("weight", weight_dims, weight_data);
test.AddInput<float>("bias", bias_dims, bias_data);
test.AddMissingOptionalInput<int32_t>();
test.AddInput<float>("past", past_dims, past_data);
test.AddReferenceOutputs("testdata/attention_past_state.onnx");
test.Run();
}
} // namespace test
} // namespace onnxruntime

View file

@ -21,142 +21,40 @@ namespace onnxruntime {
namespace test {
template <typename T>
class DynamicQuantizeMatMulOpTester : public OpTester {
public:
DynamicQuantizeMatMulOpTester(const char* op,
const std::vector<int64_t>& A_dims,
const std::vector<int64_t>& B_dims,
const std::vector<int64_t>& Y_dims,
string&& model_file,
int opset_version = 1,
const char* domain = onnxruntime::kMSDomain) : OpTester(op, opset_version, domain),
A_dims_(A_dims),
B_dims_(B_dims),
Y_dims_(Y_dims),
model_file_(model_file) {
Init();
}
void Init() {
// create rand inputs
RandomValueGenerator random{};
A_data_ = random.Uniform<float>(A_dims_, -1.0f, 1.0f);
void TestDynamicQuantizeMatMul(const std::vector<int64_t>& A_dims,
std::vector<int64_t> B_dims,
const std::string& reference_model) {
// create rand inputs
RandomValueGenerator random{};
std::vector<int> tmp_B_data = random.Uniform<int32_t>(B_dims_, std::numeric_limits<T>::min(), std::numeric_limits<T>::max());
std::transform(tmp_B_data.begin(), tmp_B_data.end(), std::back_inserter(B_data_), [](int32_t v) -> T {
return static_cast<T>(v);
});
std::vector<float> A_data = random.Uniform<float>(A_dims, -1.0f, 1.0f);
B_zero_point_ = {static_cast<T>(random.Uniform<int32_t>({1}, std::numeric_limits<T>::min(), std::numeric_limits<T>::max())[0])};
B_scale_ = random.Uniform<float>({1}, -0.1f, 0.1f);
std::vector<T> B_data;
std::vector<int> tmp_B_data = random.Uniform<int32_t>(B_dims, std::numeric_limits<T>::min(), std::numeric_limits<T>::max());
std::transform(tmp_B_data.begin(), tmp_B_data.end(), std::back_inserter(B_data), [](int32_t v) -> T {
return static_cast<T>(v);
});
const int64_t Y_size = std::accumulate(Y_dims_.cbegin(), Y_dims_.cend(), static_cast<int64_t>(1), std::multiplies<int64_t>{});
Y_data_.resize(Y_size);
std::vector<float> B_scale = random.Uniform<float>({1}, -0.1f, 0.1f);
std::vector<T> B_zero_point = {static_cast<T>(random.Uniform<int32_t>({1}, std::numeric_limits<T>::min(), std::numeric_limits<T>::max())[0])};
AddInput<float>("A", A_dims_, A_data_);
AddInput<T>("B", B_dims_, B_data_);
AddInput<float>("b_scale", {1}, B_scale_);
AddInput<T>("b_zero_point", {1}, B_zero_point_);
OpTester test("DynamicQuantizeMatMul", 1, onnxruntime::kMSDomain);
test.AddInput<float>("A", A_dims, A_data);
test.AddInput<T>("B", B_dims, B_data);
test.AddInput<float>("b_scale", {1}, B_scale);
test.AddInput<T>("b_zero_point", {1}, B_zero_point);
AddOutput<float>("Y", Y_dims_, Y_data_);
}
void Run() {
#ifndef NDEBUG
// run_called_ to true to avoid a complaining in the destructor of OpTester
run_called_ = true;
#endif
std::vector<MLValue> cpu_fetches;
std::vector<MLValue> subgraph_fetches;
Compute(cpu_fetches);
ComputeOriginalSubgraph(subgraph_fetches);
// Compare CPU with original subgraph on the output(0) - Y
ASSERT_TRUE(cpu_fetches.size() >= subgraph_fetches.size());
for (size_t i = 0; i < subgraph_fetches.size(); i++) {
if (cpu_fetches[i].IsTensor() && subgraph_fetches[i].IsTensor()) {
VLOGS_DEFAULT(1) << "Checking tensor " << i;
CheckTensor(subgraph_fetches[i].Get<Tensor>(), cpu_fetches[i].Get<Tensor>(), 1e-3, 1e-3);
}
}
}
private:
void Compute(std::vector<MLValue>& cpu_fetches) {
auto p_model = BuildGraph();
auto& graph = p_model->MainGraph();
Status status = graph.Resolve();
ASSERT_TRUE(status.IsOK()) << status;
// Hookup the inputs and outputs
NameMLValMap feeds;
std::vector<std::string> output_names;
FillFeedsAndOutputNames(feeds, output_names);
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
RunOptions run_options;
run_options.run_tag = op_;
run_options.run_log_verbosity_level = 1;
// run with DynamicQuantizeMatMul
InferenceSession session_object{so, GetEnvironment()};
std::string s1;
p_model->ToProto().SerializeToString(&s1);
std::istringstream str(s1);
ASSERT_TRUE((status = session_object.Load(str)).IsOK()) << status;
ASSERT_TRUE((status = session_object.Initialize()).IsOK()) << status;
ASSERT_TRUE((status = session_object.Run(run_options, feeds, output_names, &cpu_fetches)).IsOK());
}
void ComputeOriginalSubgraph(std::vector<MLValue>& subgraph_fetches) {
NameMLValMap feeds;
OrtValue ml_value;
std::vector<std::string> output_names;
FillFeedsAndOutputNames(feeds, output_names);
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
so.graph_optimization_level = TransformerLevel::Level1;
RunOptions run_options;
run_options.run_tag = op_;
run_options.run_log_verbosity_level = 1;
Status status;
InferenceSession subgraph_session_object{so, GetEnvironment()};
ASSERT_TRUE((status = subgraph_session_object.Load(model_file_)).IsOK()) << status;
ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status;
ASSERT_TRUE((status = subgraph_session_object.Run(run_options, feeds, output_names, &subgraph_fetches)).IsOK()) << status;
}
private:
std::vector<int64_t> A_dims_;
std::vector<int64_t> B_dims_;
std::vector<int64_t> Y_dims_;
std::vector<float> A_data_;
std::vector<T> B_data_;
std::vector<float> B_scale_;
std::vector<T> B_zero_point_;
std::vector<float> Y_data_;
std::string model_file_;
};
test.AddReferenceOutputs(reference_model);
test.Run();
}
TEST(DynamicQuantizeMatMul, Int8_test) {
#ifdef MLAS_SUPPORTS_GEMM_U8X8
std::vector<int64_t> A_dims{4, 128};
std::vector<int64_t> B_dims{128, 128};
std::vector<int64_t> Y_dims{4, 128};
DynamicQuantizeMatMulOpTester<int8_t> test("DynamicQuantizeMatMul",
A_dims,
B_dims,
Y_dims,
"testdata/dynamic_quantize_matmul_int8.onnx");
test.Run();
TestDynamicQuantizeMatMul<int8_t>(A_dims, B_dims, "testdata/dynamic_quantize_matmul_int8.onnx");
#endif
}
@ -164,12 +62,8 @@ TEST(DynamicQuantizeMatMul, UInt8_test) {
std::vector<int64_t> A_dims{4, 128};
std::vector<int64_t> B_dims{128, 128};
std::vector<int64_t> Y_dims{4, 128};
DynamicQuantizeMatMulOpTester<uint8_t> test("DynamicQuantizeMatMul",
A_dims,
B_dims,
Y_dims,
"testdata/dynamic_quantize_matmul_uint8.onnx");
test.Run();
TestDynamicQuantizeMatMul<uint8_t>(A_dims, B_dims, "testdata/dynamic_quantize_matmul_uint8.onnx");
}
} // namespace test

View file

@ -18,206 +18,29 @@ using namespace std;
namespace onnxruntime {
namespace test {
class LayerNormOpTester : public OpTester {
public:
LayerNormOpTester(const char* op,
const std::vector<int64_t>& X_dims,
const std::vector<int64_t>& scale_dims,
const std::vector<int64_t>& B_dims,
const std::vector<int64_t>& Y_dims,
float epsilon,
int64_t axis = -1,
int64_t keep_dims = 1,
int opset_version = 1,
const char* domain = onnxruntime::kOnnxDomain) : OpTester(op, opset_version, domain),
X_dims_(X_dims),
scale_dims_(scale_dims),
B_dims_(B_dims),
Y_dims_(Y_dims),
epsilon_(epsilon),
axis_(axis),
keep_dims_(keep_dims) {
Init();
}
void Init() {
AddAttribute("axis", axis_);
AddAttribute("keep_dims", keep_dims_);
AddAttribute("epsilon", epsilon_);
// create rand inputs
RandomValueGenerator random{};
X_data_ = random.Uniform<float>(X_dims_, 0.0f, 1.0f);
scale_data_ = random.Uniform<float>(scale_dims_, 0.0f, 1.0f);
B_data_ = random.Uniform<float>(B_dims_, 0.0f, 1.0f);
const int64_t Y_size = std::accumulate(Y_dims_.cbegin(), Y_dims_.cend(), static_cast<int64_t>(1), std::multiplies<int64_t>{});
Y_data_.resize(Y_size);
AddInput<float>("X", X_dims_, X_data_);
AddInput<float>("scale", scale_dims_, scale_data_, true);
AddInput<float>("B", B_dims_, B_data_, true);
AddOutput<float>("output", Y_dims_, Y_data_);
AddOutput<float>("mean", mean_);
AddOutput<float>("inv_std_var", inv_std_var_);
}
void Run() {
#ifndef NDEBUG
run_called_ = true;
#endif
std::vector<MLValue> cpu_fetches;
std::vector<MLValue> cuda_fetches;
std::vector<MLValue> subgraph_fetches;
ComputeWithCPU(cpu_fetches);
ComputeWithCUDA(cuda_fetches);
ComputeOriSubgraphWithCPU(subgraph_fetches);
// Compare CPU with original subgraph on the output(0) - Y
ASSERT_TRUE(cpu_fetches.size() >= subgraph_fetches.size());
for (size_t i = 0; i < subgraph_fetches.size(); i++) {
if (cpu_fetches[i].IsTensor() && subgraph_fetches[i].IsTensor()) {
VLOGS_DEFAULT(1) << "Checking tensor " << i;
CheckTensor(subgraph_fetches[i].Get<Tensor>(), cpu_fetches[i].Get<Tensor>(), 1e-3, 1e-3);
}
}
// Compare GPU with original subgraph on the output(0) - Y
if (DefaultCudaExecutionProvider()) {
ASSERT_TRUE(cuda_fetches.size() >= subgraph_fetches.size());
for (size_t i = 0; i < subgraph_fetches.size(); i++) {
if (cuda_fetches[i].IsTensor() && subgraph_fetches[i].IsTensor()) {
VLOGS_DEFAULT(1) << "Checking tensor " << i;
CheckTensor(subgraph_fetches[i].Get<Tensor>(), cuda_fetches[i].Get<Tensor>(), 1e-3, 1e-3);
}
}
}
}
private:
void ComputeWithCPU(std::vector<MLValue>& cpu_fetches);
void ComputeWithCUDA(std::vector<MLValue>& cuda_fetches);
void ComputeOriSubgraphWithCPU(std::vector<MLValue>& subgraph_fetches);
private:
std::vector<int64_t> X_dims_;
std::vector<int64_t> scale_dims_;
std::vector<int64_t> B_dims_;
std::vector<int64_t> Y_dims_;
std::vector<float> X_data_;
std::vector<float> scale_data_;
std::vector<float> B_data_;
std::vector<float> Y_data_;
float mean_;
float inv_std_var_;
float epsilon_;
int64_t axis_;
int64_t keep_dims_;
};
void LayerNormOpTester::ComputeWithCPU(std::vector<MLValue>& cpu_fetches) {
auto p_model = BuildGraph();
auto& graph = p_model->MainGraph();
Status status = graph.Resolve();
ASSERT_TRUE(status.IsOK()) << status;
// Hookup the inputs and outputs
std::unordered_map<std::string, MLValue> feeds;
std::vector<std::string> output_names;
FillFeedsAndOutputNames(feeds, output_names);
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
RunOptions run_options;
run_options.run_tag = op_;
run_options.run_log_verbosity_level = 1;
// run with LayerNormalization
InferenceSession layernorm_session_object{so, GetEnvironment()};
std::string s1;
p_model->ToProto().SerializeToString(&s1);
std::istringstream str(s1);
ASSERT_TRUE((status = layernorm_session_object.Load(str)).IsOK()) << status;
ASSERT_TRUE((status = layernorm_session_object.Initialize()).IsOK()) << status;
ASSERT_TRUE((status = layernorm_session_object.Run(run_options, feeds, output_names, &cpu_fetches)).IsOK());
}
void LayerNormOpTester::ComputeWithCUDA(std::vector<MLValue>& cuda_fetches) {
if (DefaultCudaExecutionProvider() == nullptr) {
return;
}
auto p_model = BuildGraph();
auto& graph = p_model->MainGraph();
Status status = graph.Resolve();
ASSERT_TRUE(status.IsOK()) << status;
// Hookup the inputs and outputs
std::unordered_map<std::string, MLValue> feeds;
std::vector<std::string> output_names;
FillFeedsAndOutputNames(feeds, output_names);
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
RunOptions run_options;
run_options.run_tag = op_;
run_options.run_log_verbosity_level = 1;
auto cuda_execution_provider = DefaultCudaExecutionProvider();
InferenceSession cuda_session_object{so, GetEnvironment()};
EXPECT_TRUE(cuda_session_object.RegisterExecutionProvider(std::move(cuda_execution_provider)).IsOK());
std::string s;
p_model->ToProto().SerializeToString(&s);
std::istringstream str2(s);
EXPECT_TRUE((status = cuda_session_object.Load(str2)).IsOK()) << status;
EXPECT_TRUE((status = cuda_session_object.Initialize()).IsOK()) << status;
EXPECT_TRUE((status = cuda_session_object.Run(run_options, feeds, output_names, &cuda_fetches)).IsOK()) << status;
}
void LayerNormOpTester::ComputeOriSubgraphWithCPU(std::vector<MLValue>& subgraph_fetches) {
NameMLValMap feeds;
OrtValue ml_value;
std::vector<std::string> output_names{"Y"};
CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), X_dims_, X_data_, &ml_value);
feeds.insert(std::make_pair("X", ml_value));
CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), scale_dims_, scale_data_, &ml_value);
feeds.insert(std::make_pair("Scale", ml_value));
CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), B_dims_, B_data_, &ml_value);
feeds.insert(std::make_pair("B", ml_value));
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
RunOptions run_options;
run_options.run_tag = op_;
run_options.run_log_verbosity_level = 1;
Status status;
InferenceSession subgraph_session_object{so, GetEnvironment()};
ASSERT_TRUE((status = subgraph_session_object.Load("testdata/layernorm.onnx")).IsOK()) << status;
ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status;
ASSERT_TRUE((status = subgraph_session_object.Run(run_options, feeds, output_names, &subgraph_fetches)).IsOK()) << status;
}
TEST(LayerNormTest, BERTLayerNorm) {
float epsilon = 1e-12f;
OpTester tester("LayerNormalization", 1 /*opset_version*/);
tester.AddAttribute<int64_t>("axis", -1);
tester.AddAttribute<float>("epsilon", 1e-12f);
// create rand inputs
RandomValueGenerator random{};
std::vector<int64_t> X_dims{4, 128};
std::vector<float> X_data = random.Uniform<float>(X_dims, 0.0f, 1.0f);
tester.AddInput<float>("X", X_dims, X_data);
std::vector<int64_t> scale_dims{128};
std::vector<float> scale_data = random.Uniform<float>(scale_dims, 0.0f, 1.0f);
tester.AddInput<float>("Scale", scale_dims, scale_data);
std::vector<int64_t> B_dims{128};
std::vector<int64_t> Y_dims{4, 128};
LayerNormOpTester test("LayerNormalization", X_dims, scale_dims, B_dims, Y_dims, epsilon, -1, 1, 9);
test.Run();
std::vector<float> B_data = random.Uniform<float>(B_dims, 0.0f, 1.0f);
tester.AddInput<float>("B", B_dims, B_data);
tester.AddReferenceOutputs("testdata/layernorm.onnx");
tester.Run();
}
} // namespace test

View file

@ -705,5 +705,70 @@ TEST(QAttentionTest, QAttentionUnidirectional_CUDA) {
true /*is_unidirectional*/);
}
template <typename InputT, typename WeightT>
void TestQuantizedAttentionPastState(int64_t batch,
int64_t seq_len,
int64_t past_seq_len,
int64_t hidden_size,
int64_t head_number,
int64_t head_size,
const std::string& reference_model) {
// create rand inputs
RandomValueGenerator random{};
constexpr InputT input_min = std::numeric_limits<InputT>::min();
constexpr InputT input_max = std::numeric_limits<InputT>::max();
constexpr int32_t input_range = input_max - input_min;
InputT input_mean = (input_min + input_max) / 2 + 1;
std::vector<InputT> input_zero_point{input_mean};
std::vector<int64_t> input_dims{batch, seq_len, hidden_size};
std::vector<InputT> input_data = random.Gaussian<InputT>(input_dims, input_mean, static_cast<InputT>(input_range / 6), input_min, input_max);
constexpr WeightT weight_min = std::numeric_limits<WeightT>::min();
constexpr WeightT weight_max = std::numeric_limits<WeightT>::max();
constexpr int32_t weight_range = weight_max - weight_min;
WeightT weight_mean = (weight_min + weight_max) / 2 + 1;
std::vector<WeightT> weight_zero_point{weight_mean};
std::vector<int64_t> weight_dims{hidden_size, 3 * hidden_size};
std::vector<WeightT> weight_data = random.Gaussian<WeightT>(weight_dims, weight_mean, static_cast<WeightT>(weight_range / 6), weight_min, weight_max);
std::vector<int64_t> bias_dims{3 * hidden_size};
std::vector<float> bias_data = random.Gaussian<float>(bias_dims, 0.0f, 0.3f);
std::vector<float> input_scale{0.005f};
std::vector<float> weight_scale{0.005f};
std::vector<int64_t> past_dims{2, batch, head_number, past_seq_len, head_size};
std::vector<float> past_data = random.Gaussian<float>(past_dims, 0.0f, 0.3f);
OpTester test("QAttention", 1, onnxruntime::kMSDomain);
test.AddAttribute<int64_t>("num_heads", head_number);
test.AddAttribute<int64_t>("unidirectional", 1);
test.AddInput<InputT>("input", input_dims, input_data);
test.AddInput<WeightT>("weight", weight_dims, weight_data);
test.AddInput<float>("bias", bias_dims, bias_data);
test.AddInput<float>("input_scale", {1}, input_scale);
test.AddInput<float>("weight_scale", {1}, weight_scale);
test.AddMissingOptionalInput<int32_t>();
test.AddInput<InputT>("input_zero_point", {1}, input_zero_point);
test.AddInput<WeightT>("weight_zero_point", {1}, weight_zero_point);
test.AddInput<float>("past", past_dims, past_data);
test.AddReferenceOutputs(reference_model);
test.Run();
}
TEST(QAttentionTest, QAttentionPastState_u8u8) {
TestQuantizedAttentionPastState<uint8_t, uint8_t>(2, 5, 15, 768, 12, 64, "testdata/attention_past_state.u8u8.onnx");
}
TEST(QAttentionTest, QAttentionPastState_u8s8) {
TestQuantizedAttentionPastState<uint8_t, int8_t>(2, 5, 15, 768, 12, 64, "testdata/attention_past_state.u8s8.onnx");
}
} // namespace test
} // namespace onnxruntime

View file

@ -393,6 +393,16 @@ void OpTester::FillFeedsAndOutputNames(
}
}
void OpTester::FillFeeds(std::unordered_map<std::string, OrtValue>& feeds) {
for (size_t i = 0; i < input_data_.size(); ++i) {
if (std::find(initializer_index_.begin(), initializer_index_.end(), i) ==
initializer_index_.end() &&
input_data_[i].def_.Exists()) {
feeds[input_data_[i].def_.Name()] = input_data_[i].data_;
}
}
}
void OpTester::SetOutputAbsErr(const char* name, float v) {
auto it =
std::find_if(output_data_.begin(), output_data_.end(),
@ -524,7 +534,7 @@ std::vector<MLValue> OpTester::ExecuteModel(
// Disable expected_failure_string checks for OpenVINO EP
if (provider_type != kOpenVINOExecutionProvider) {
EXPECT_THAT(status.ErrorMessage(),
testing::HasSubstr(expected_failure_string));
testing::HasSubstr(expected_failure_string));
}
} else {
LOGS_DEFAULT(ERROR) << "Initialize failed with status: "
@ -804,7 +814,7 @@ void OpTester::Run(
valid = false;
for (auto& custom_session_registry : custom_session_registries_) {
if (KernelRegistry::HasImplementationOf(*custom_session_registry->GetKernelRegistry(),
node, execution_provider->Type())) {
node, execution_provider->Type())) {
valid = true;
break;
}
@ -843,6 +853,56 @@ void OpTester::Run(
}
}
void OpTester::AddReferenceOutputs(const std::string& model_path) {
SessionOptions so;
so.session_logid = op_;
so.session_log_verbosity_level = 1;
RunOptions run_options;
run_options.run_tag = op_;
run_options.run_log_verbosity_level = 1;
Status status;
InferenceSession subgraph_session_object{so, GetEnvironment()};
ASSERT_TRUE((status = subgraph_session_object.Load(model_path)).IsOK()) << status;
ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status;
// Retrieve output names
auto model_outputs = subgraph_session_object.GetModelOutputs();
ASSERT_TRUE(model_outputs.first.IsOK());
std::vector<std::string> output_names;
std::transform(model_outputs.second->begin(),
model_outputs.second->end(),
std::back_inserter(output_names),
[](const onnxruntime::NodeArg* node_arg) -> std::string { return node_arg->Name(); });
NameMLValMap feeds;
FillFeeds(feeds);
std::vector<MLValue> subgraph_fetches;
ASSERT_TRUE((status = subgraph_session_object.Run(run_options, feeds, output_names, &subgraph_fetches)).IsOK()) << status;
for (size_t out_idx = 0; out_idx < subgraph_fetches.size(); out_idx++) {
// Retrieve TypeProto
ASSERT_TRUE(subgraph_fetches[out_idx].Type()->IsTensorType()) << status;
const Tensor& t = subgraph_fetches[out_idx].Get<Tensor>();
const TensorTypeBase* tensor_type = DataTypeImpl::TensorTypeFromONNXEnum(t.GetElementType());
// Construct a temp TypeProto with shape information
ONNX_NAMESPACE::TypeProto tmp_type_proto(*(tensor_type->GetTypeProto()));
auto mutable_shape = tmp_type_proto.mutable_tensor_type()->mutable_shape();
for (auto i : t.Shape().GetDims()) {
auto* mutable_dim = mutable_shape->add_dim();
mutable_dim->set_dim_value(i);
}
output_data_.push_back(Data(NodeArg(output_names[out_idx], &tmp_type_proto),
std::move(subgraph_fetches[out_idx]),
optional<float>(),
optional<float>()));
}
}
#ifdef ENABLE_TRAINING
template std::vector<MLValue> OpTester::ExecuteModel<training::TrainingSession>(
Model& model, training::TrainingSession& session_object,

View file

@ -210,7 +210,8 @@ const SequenceTensorTypeProto<ElemType> SequenceTensorType<ElemType>::s_sequence
// 1. Create one with the op name
// 2. Call AddAttribute with any attributes
// 3. Call AddInput for all the inputs
// 4. Call AddOutput with all expected outputs
// 4. Call AddOutput with all expected outputs,
// Or call AddReferenceOutputs to compute reference outputs with the model
// 5. Call Run
// Not all tensor types and output types are added, if a new input type is used, add it to the TypeToDataType list
// above for new output types, add a new specialization for Check<> See current usage for an example, should be self
@ -382,6 +383,9 @@ class OpTester {
optional<float>(), optional<float>()));
}
// Generate the reference outputs with the model file
void AddReferenceOutputs(const std::string& model_path);
void AddCustomOpRegistry(std::shared_ptr<CustomRegistry> registry) {
custom_schema_registries_.push_back(registry->GetOpschemaRegistry());
custom_session_registries_.push_back(registry);
@ -483,6 +487,8 @@ class OpTester {
void FillFeedsAndOutputNames(std::unordered_map<std::string, OrtValue>& feeds,
std::vector<std::string>& output_names);
void FillFeeds(std::unordered_map<std::string, OrtValue>& feeds);
template <class SessionType>
std::vector<MLValue> ExecuteModel(Model& model,
SessionType& session_object,
@ -545,7 +551,7 @@ class OpTester {
const static std::unordered_set<std::string> reserved_symbolic{"batch", "seq"};
for (size_t i = 0; i < dim_params_data.size(); ++i) {
if (reserved_symbolic.find(dim_params_data[i])!= reserved_symbolic.end()) {
if (reserved_symbolic.find(dim_params_data[i]) != reserved_symbolic.end()) {
new_shape.add_dim()->set_dim_param(dim_params_data[i]);
} else {
ASSERT_TRUE(std::stoi(dim_params_data[i]) == dims[i]);

Binary file not shown.

Binary file not shown.

Binary file not shown.