mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
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:
parent
ceedf126a2
commit
fc5e65a22d
22 changed files with 591 additions and 625 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
33
onnxruntime/contrib_ops/cpu/bert/attention_base.h
Normal file
33
onnxruntime/contrib_ops/cpu/bert/attention_base.h
Normal 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
|
||||
218
onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
Normal file
218
onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
|
|
|
|||
BIN
onnxruntime/test/testdata/attention_past_state.onnx
vendored
Normal file
BIN
onnxruntime/test/testdata/attention_past_state.onnx
vendored
Normal file
Binary file not shown.
BIN
onnxruntime/test/testdata/attention_past_state.u8s8.onnx
vendored
Normal file
BIN
onnxruntime/test/testdata/attention_past_state.u8s8.onnx
vendored
Normal file
Binary file not shown.
BIN
onnxruntime/test/testdata/attention_past_state.u8u8.onnx
vendored
Normal file
BIN
onnxruntime/test/testdata/attention_past_state.u8u8.onnx
vendored
Normal file
Binary file not shown.
Loading…
Reference in a new issue