diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.cc b/onnxruntime/contrib_ops/cpu/bert/attention.cc index d6bb58300d..e6ed0c644b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/attention.cc @@ -27,18 +27,21 @@ AttentionBase::AttentionBase(const OpKernelInfo& info) { int64_t num_heads = 0; ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); num_heads_ = static_cast(num_heads); + is_unidirectional_ = info.GetAttrOrDefault("unidirectional", 0) == 1; } Status AttentionBase::CheckInputs(const Tensor* input, const Tensor* weights, const Tensor* bias, - const Tensor* mask_index) const { - // Input and output shapes: + const Tensor* mask_index, + const Tensor* past) const { + // Input shapes: // input : (batch_size, sequence_length, hidden_size) // weights : (hidden_size, 3 * hidden_size) // bias : (3 * hidden_size) // mask_index : (batch_size) if presented + // past : (2, batch_size, num_heads, past_sequence_length, head_size) const auto dims = input->Shape().GetDims(); if (dims.size() != 3) { @@ -91,9 +94,59 @@ Status AttentionBase::CheckInputs(const Tensor* input, } } + if (past != nullptr) { // past is optional + if (!is_unidirectional_) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 (past) is only allowed for unidirectional"); + } + + const auto past_dims = past->Shape().GetDims(); + if (past_dims.size() != 5) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 4 is expected to have 5 dimension, got ", + past_dims.size()); + } + if (static_cast(past_dims[0]) != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 0 shall have length of 2"); + } + if (static_cast(past_dims[1]) != batch_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 1 shall have same length as dimension 0 of input 0"); + } + if (static_cast(past_dims[2]) != num_heads_) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 2 shall have length of num_heads", num_heads_); + } + if (static_cast(past_dims[4]) != hidden_size / num_heads_) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 4 dimension 2 shall have length of ", hidden_size / num_heads_); + } + } + return Status::OK(); } +Tensor* AttentionBase::GetPresent(OpKernelContext* context, + const Tensor* past, + int batch_size, + int head_size, + int sequence_length, + int& past_sequence_length) const { + // Input and output shapes: + // past : (2, batch_size, num_heads, past_sequence_length, head_size) + // present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size) + + std::vector present_dims{2, batch_size, num_heads_, sequence_length, head_size}; + if (nullptr != past) { + const auto past_dims = past->Shape().GetDims(); + past_sequence_length = static_cast(past_dims[3]); + present_dims[3] += past_dims[3]; + } + + TensorShape present_shape(present_dims); + Tensor* present = context->Output(1, present_shape); + if (nullptr != past && nullptr == present) { + ORT_THROW("Expect to have present state output when past state input is given"); + } + + return present; +} + template Attention::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionBase(info) { } @@ -104,7 +157,9 @@ Status Attention::Compute(OpKernelContext* context) const { const Tensor* weights = context->Input(1); const Tensor* bias = context->Input(2); const Tensor* mask_index = context->Input(3); - ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index)); + const Tensor* past = context->Input(4); + + ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index, past)); const auto dims = input->Shape().GetDims(); const int batch_size = static_cast(dims[0]); @@ -115,6 +170,12 @@ Status Attention::Compute(OpKernelContext* context) const { TensorShape output_shape(dims); Tensor* output = context->Output(0, output_shape); + int past_sequence_length = 0; + Tensor* present = GetPresent(context, past, batch_size, head_size, sequence_length, past_sequence_length); + + // Total sequence length including that of past state: S* = S' + S + const int all_sequence_length = past_sequence_length + sequence_length; + constexpr size_t element_size = sizeof(T); AllocatorPtr allocator; @@ -182,18 +243,18 @@ Status Attention::Compute(OpKernelContext* context) const { } // STEP.2: compute the attention score. It does 2 things: - // I. attention_probs(B, N, S, S) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S, H -> B, N, H, S) + - // 1 x mask_data(B, N, S, S) - // II.attention_probs(B, N, S, S) = Softmax(attention_probs) - size_t attention_probs_bytes = SafeInt(batch_size) * num_heads_ * sequence_length * sequence_length * element_size; + // I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) + + // 1 x mask_data(B, N, S, S*) + // II.attention_probs(B, N, S, S*) = Softmax(attention_probs) + size_t attention_probs_bytes = SafeInt(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(batch_size) * sequence_length * sequence_length * element_size; + mask_data_bytes = SafeInt(batch_size) * sequence_length * all_sequence_length * element_size; } else if (is_unidirectional_) { - mask_data_bytes = SafeInt(sequence_length) * sequence_length * element_size; + mask_data_bytes = SafeInt(sequence_length) * all_sequence_length * element_size; } void* mask_data = nullptr; @@ -204,17 +265,21 @@ Status Attention::Compute(OpKernelContext* context) const { BufferUniquePtr mask_data_buffer(mask_data, BufferDeleter(allocator)); const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data() : nullptr; + const T* past_data = past != nullptr ? past->template Data() : nullptr; + T* present_data = present != nullptr ? present->template MutableData() : nullptr; ComputeAttentionProbs(static_cast(attention_probs), Q, K, mask_index_data, static_cast(mask_data), - batch_size, sequence_length, head_size, num_heads_, is_unidirectional_, tp); + batch_size, sequence_length, past_sequence_length, head_size, num_heads_, is_unidirectional_, + past_data, present_data, tp); - // STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S) x V(B, N, S, H) + // STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S*) x V(B, N, S*, H) auto out_tmp_data = allocator->Alloc(SafeInt(batch_size) * num_heads_ * sequence_length * head_size * element_size); BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator)); ComputeVxAttentionScore(output->template MutableData(), static_cast(out_tmp_data), static_cast(attention_probs), V, - batch_size, sequence_length, head_size, num_heads_, hidden_size, tp); + batch_size, sequence_length, past_sequence_length, head_size, num_heads_, hidden_size, + past_data, present_data, tp); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.h b/onnxruntime/contrib_ops/cpu/bert/attention.h index e82d758cdb..fa686284ac 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention.h @@ -15,10 +15,18 @@ class AttentionBase { Status CheckInputs(const Tensor* input, const Tensor* weights, const Tensor* bias, - const Tensor* mask_index) const; + const Tensor* mask_index, + const Tensor* past) const; - int num_heads_; // number of attention heads - bool is_unidirectional_; // whether every token can only attend to previous tokens. + Tensor* GetPresent(OpKernelContext* context, + const Tensor* past, + int batch_size, + int head_size, + int sequence_length, + int& past_sequence_length) const; + + int num_heads_; // number of attention heads + bool is_unidirectional_; // whether every token can only attend to previous tokens. }; template diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h index bed0513ff6..4a94d0b45b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h @@ -59,111 +59,176 @@ inline void ComputeAttentionSoftmaxInplace(float* score, int N, int D, ThreadPoo MlasComputeSoftmax(score, score, N, D, false, tp); } -// Helper function to compute the attention probs. It does 2 things: -// I. attention_probs(B, N, S, S) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S, H -> B, N, H, S) + -// 1 x mask_data(B, N, S, S) -// II.attention_probs(B, N, S, S) = Softmax(attention_probs) template -void ComputeAttentionProbs(T* attention_probs, // output buffer for the attention probs. Its size is B*N*S*S - const T* Q, // Q data. Its size is B*N*S*H - const T* K, // k data. Its size is B*N*S*H +void PrepareMask(const int32_t* mask_index, + T* mask_data, + bool is_unidirectional, + int batch_size, + int sequence_length, + int past_sequence_length) { + const int all_sequence_length = past_sequence_length + sequence_length; + T* p_mask = mask_data; + if (is_unidirectional) { + // unidirectional mask has shape SxS* + for (int s_i = 0; s_i < sequence_length - 1; s_i++) { + for (int m_i = past_sequence_length + s_i + 1; m_i < all_sequence_length; m_i++) { + p_mask[s_i * all_sequence_length + m_i] = static_cast(-10000.0); + } + } + return; + } + + ORT_ENFORCE(mask_index, "mask index should not be null."); + for (int b_i = 0; b_i < batch_size; b_i++) { + // TODO: mask_index can be used in softmax to save some calculation. + // Convert mask_index to mask (-10000 means out of range, which will be 0 after softmax): B => BxS* + int valid_length = mask_index[b_i]; + for (int m_i = valid_length; m_i < all_sequence_length; m_i++) { + p_mask[m_i] = static_cast(-10000.0); + } + + // Broadcast mask from BxS* to BxSxS* + for (int s_i = 1; s_i < sequence_length; s_i++) { + memcpy(p_mask + s_i * all_sequence_length, p_mask, all_sequence_length * sizeof(T)); + } + p_mask += sequence_length * sequence_length; + } +} + +// Concatenate a past state chunk S'xH with input state chunk SxH into present state chunk S*xH +// Returns a pointer to the start of present state chunk. +template +T* ConcatStateChunk(const T* past, const T* chunk, T* present, size_t past_chunk_length, size_t present_chunk_length, std::ptrdiff_t i) { + T* start = present + i * present_chunk_length; + + T* p = start; + if (nullptr != past) { + const T* src_past = past + i * past_chunk_length; + memcpy(p, src_past, past_chunk_length * sizeof(T)); + p += past_chunk_length; + } + + memcpy(p, chunk, (present_chunk_length - past_chunk_length) * sizeof(T)); + return start; +} + +// Helper function to compute the attention probs. It does 2 things: +// I. attention_probs(B, N, S, S*) = 1/sqrt(H) x Q(B, N, S, H) x K'(B, N, S*, H -> B, N, H, S*) + +// 1 x mask_data(B, N, S, S*) +// II.attention_probs(B, N, S, S*) = Softmax(attention_probs) +template +void ComputeAttentionProbs(T* attention_probs, // output buffer for the attention probs. Its size is BxNxSxS + const T* Q, // Q data. Its size is BxNxSxH + const T* K, // k data. Its size is BxNxSxH const int32_t* mask_index, // mask index. nullptr if no mask or its size is B - T* mask_data, // buffer for mask data. Its size is: S*S if is_unidirectiona; B*S*S if mask_index; null otherwise + T* mask_data, // buffer for mask data. Its size is: SxS* if is_unidirectional; BxSxS* if mask_index; null otherwise int batch_size, // batch size of self-attention int sequence_length, // sequence length of self-attention + int past_sequence_length, // sequence length of past state int head_size, // head size of self-attention int num_heads, // number of heads of self-attention bool is_unidirectional, // indicate if it is unidrectional. + const T* past, // past state + T* present, // present state ThreadPool* tp) { + const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S + const size_t past_chunk_length = static_cast(past_sequence_length * head_size); // S' x H + const size_t input_chunk_length = static_cast(sequence_length * head_size); // S x H + const size_t present_chunk_length = past_chunk_length + input_chunk_length; // S* x H + { if (mask_data != nullptr) { - if (is_unidirectional) { - for (int s_i = 0; s_i < sequence_length - 1; s_i++) { - for (int m_i = s_i + 1; m_i < sequence_length; m_i++) { - mask_data[s_i * sequence_length + m_i] = static_cast(-10000.0); - } - } - } else { - ORT_ENFORCE(mask_index, "mask index should not be null."); - T* p_mask = mask_data; - for (int b_i = 0; b_i < batch_size; b_i++) { - // TODO: mask_index can be used in softmax to save some calculation. - // Convert mask_index to mask (-10000 means out of range, which will be 0 after softmax): B => BxS - int valid_length = mask_index[b_i]; - for (int m_i = valid_length; m_i < sequence_length; m_i++) { - p_mask[m_i] = static_cast(-10000.0); - } - - // Broadcast mask from BxS to BxSxS - for (int s_i = 1; s_i < sequence_length; s_i++) { - memcpy(p_mask + s_i * sequence_length, p_mask, sequence_length * sizeof(T)); - } - p_mask += sequence_length * sequence_length; - } - } + PrepareMask(mask_index, mask_data, is_unidirectional, batch_size, sequence_length, past_sequence_length); } else { // no any mask - memset(attention_probs, 0, batch_size * num_heads * sequence_length * sequence_length * sizeof(T)); + memset(attention_probs, 0, batch_size * num_heads * sequence_length * all_sequence_length * sizeof(T)); } const int loop_len = batch_size * num_heads; const float alpha = 1.0f / sqrt(static_cast(head_size)); // The cost of Gemm - const double cost = - static_cast(head_size) * static_cast(sequence_length) * static_cast(sequence_length); + const double cost = static_cast(head_size * sequence_length * all_sequence_length); + ThreadPool::TryParallelFor(tp, loop_len, cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t i = begin; i != end; ++i) { const std::ptrdiff_t batch_index = i / num_heads; - // broadcast mask data: SxS or (Bx)SxS -> (BxNx)SxS + // broadcast mask data: SxS* or (Bx)SxS* -> (BxNx)SxS* if (mask_data != nullptr) { - const T* broadcast_data_src = is_unidirectional ? reinterpret_cast(mask_data) : reinterpret_cast(mask_data) + batch_index * sequence_length * sequence_length; - T* broadcast_data_dest = reinterpret_cast(attention_probs) + sequence_length * sequence_length * i; - memcpy(broadcast_data_dest, broadcast_data_src, sequence_length * sequence_length * sizeof(T)); + const T* broadcast_data_src = is_unidirectional ? reinterpret_cast(mask_data) : reinterpret_cast(mask_data) + batch_index * sequence_length * all_sequence_length; + T* broadcast_data_dest = reinterpret_cast(attention_probs) + sequence_length * all_sequence_length * i; + memcpy(broadcast_data_dest, broadcast_data_src, sequence_length * all_sequence_length * sizeof(T)); + } + + const T* k = K + input_chunk_length * i; + if (nullptr != present) { + // concatenate past_K and K : (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH + k = ConcatStateChunk(past, k, present, past_chunk_length, present_chunk_length, i); } // gemm - - // original transposed iteration - // A: Q (BxNxSxH) (B.N.)S x H S x H - // B: K' (BxNxSxH) (B.N.)H x S H x S - // C: attention_probs (BxNxSxS) (B.N.)S x S S x S - - math::Gemm(CblasNoTrans, CblasTrans, sequence_length, sequence_length, head_size, alpha, - Q + sequence_length * head_size * i, K + sequence_length * head_size * i, 1.0, - reinterpret_cast(attention_probs) + sequence_length * sequence_length * i, nullptr); + // original transposed each iteration + // A: Q (B x N x) S x H (B x N x) S x H S x H + // B: K' (B x N x) S* x H (B x N x) H x S* H x S* + // C: attention_probs (B x N x) S x S* (B x N x) S x S* S x S* + math::Gemm(CblasNoTrans, CblasTrans, sequence_length, all_sequence_length, head_size, alpha, + Q + input_chunk_length * i, k, 1.0, + reinterpret_cast(attention_probs) + sequence_length * all_sequence_length * i, nullptr); } }); } - // attention_probs(B, N, S, S) = Softmax(attention_probs) + // attention_probs(B, N, S, S*) = Softmax(attention_probs) { const int N = batch_size * num_heads * sequence_length; - const int D = sequence_length; + const int D = all_sequence_length; ComputeAttentionSoftmaxInplace(attention_probs, N, D, tp); } } template -void ComputeVxAttentionScore(T* output, // buffer for the result with size B*S*N*H - T* tmp_buffer, // buffer for temp use with size is B*N*S*H - const T* attention_probs, // Attention probs with size B*N*S*S - const T* V, // V valuee with size B*N*S*H +void ComputeVxAttentionScore(T* output, // buffer for the result with size BxSxNxH + T* tmp_buffer, // buffer for temp use with size is BxNxSxH + const T* attention_probs, // Attention probs with size BxNxSxS* + const T* V, // V value with size BxNxSxH int batch_size, // batch size int sequence_length, // sequence length + int past_sequence_length, // sequence length in past state int head_size, // head size int num_heads, // number of heads int hidden_size, // hidden size + const T* past, // past state + T* present, // present state ThreadPool* tp) { + const int all_sequence_length = past_sequence_length + sequence_length; // S* = S' + S + const size_t past_chunk_length = static_cast(past_sequence_length * head_size); // S' x H + const size_t input_chunk_length = static_cast(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(sequence_length) * static_cast(head_size) * static_cast(sequence_length); + ThreadPool::TryParallelFor(tp, batch_size * num_heads, cost, [&](std::ptrdiff_t begin, std::ptrdiff_t end) { - const int sequence_length_mul_head_size = sequence_length * head_size; for (std::ptrdiff_t i = begin; i != end; ++i) { - T* current_tmp_data = tmp_buffer + sequence_length_mul_head_size * i; - math::MatMul(sequence_length, head_size, sequence_length, - attention_probs + sequence_length * sequence_length * i, - V + sequence_length_mul_head_size * i, current_tmp_data, nullptr); + + const T* v = V + input_chunk_length * i; + if (nullptr != present) { + // concatenate past_V and V: (BxNx)S'xH, (BxNx)SxH -> (BxNx)S*xH + v = ConcatStateChunk(past, v, present, past_chunk_length, present_chunk_length, i); + } + + T* current_tmp_data = reinterpret_cast(tmp_buffer) + input_chunk_length * i; + math::MatMul(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(i / num_heads); diff --git a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc index 8f186e23a4..145276c73c 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc @@ -59,7 +59,7 @@ Status QAttention::Compute(OpKernelContext* context) const { const Tensor* i_zp_tensor = context->Input(6); const Tensor* w_zp_tensor = context->Input(7); - ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index)); + ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr)); ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor), "input scale must be a scalar or 1D tensor of size 1"); @@ -193,8 +193,12 @@ Status QAttention::Compute(OpKernelContext* context) const { const int32_t* mask_index_data = mask_index != nullptr ? mask_index->template Data() : nullptr; + int past_sequence_length = 0; + const T* past_data = nullptr; + T* present_data = nullptr; ComputeAttentionProbs(static_cast(attention_probs), Q, K, mask_index_data, static_cast(mask_data), - batch_size, sequence_length, head_size, num_heads_, is_unidirectional_, tp); + batch_size, sequence_length, past_sequence_length, head_size, num_heads_, is_unidirectional_, + past_data, present_data, tp); // STEP.3: compute the attentionScore * Value. It does: out_tmp(B, N, S, H) = attention_probs(B, N, S, S) x V(B, N, S, H) auto out_tmp_data = @@ -202,7 +206,7 @@ Status QAttention::Compute(OpKernelContext* context) const { BufferUniquePtr out_tmp_buffer(out_tmp_data, BufferDeleter(allocator)); ComputeVxAttentionScore(output->template MutableData(), static_cast(out_tmp_data), static_cast(attention_probs), V, - batch_size, sequence_length, head_size, num_heads_, hidden_size, tp); + batch_size, sequence_length, past_sequence_length, head_size, num_heads_, hidden_size, past_data, present_data, tp); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.cc b/onnxruntime/contrib_ops/cuda/bert/attention.cc index 5d679fd335..b203a08345 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/attention.cc @@ -34,18 +34,16 @@ Attention::Attention(const OpKernelInfo& info) : CudaKernel(info), AttentionB template Status Attention::ComputeInternal(OpKernelContext* context) const { - // Input and output shapes: - // Input 0 - input : (batch_size, sequence_length, hidden_size) - // Input 1 - weights : (hidden_size, 3 * hidden_size) - // Input 2 - bias : (3 * hidden_size) - // Input 3 - mask_index : (batch_size) if presented - // Output : (batch_size, sequence_length, hidden_size) const Tensor* input = context->Input(0); const Tensor* weights = context->Input(1); const Tensor* bias = context->Input(2); const Tensor* mask_index = context->Input(3); - ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index)); + const Tensor* past = context->Input(4); + ORT_RETURN_IF_ERROR(CheckInputs(input, weights, bias, mask_index, past)); + // Input and output shapes: + // Input 0 - input : (batch_size, sequence_length, hidden_size) + // Output 0 - output : (batch_size, sequence_length, hidden_size) const auto dims = input->Shape().GetDims(); int batch_size = static_cast(dims[0]); int sequence_length = static_cast(dims[1]); @@ -55,8 +53,11 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { TensorShape output_shape(dims); Tensor* output = context->Output(0, output_shape); + int past_sequence_length = 0; + Tensor* present = GetPresent(context, past, batch_size, head_size, sequence_length, past_sequence_length); + cublasHandle_t cublas = CublasHandle(); - const size_t element_size = sizeof(T); + constexpr size_t element_size = sizeof(T); // Use GEMM for fully connection. int m = batch_size * sequence_length; @@ -84,7 +85,7 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { reinterpret_cast(input->template Data()), k, &one, reinterpret_cast(gemm_buffer.get()), n, device_prop)); - size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length); + size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, past_sequence_length); auto temp_buffer = GetScratchBuffer(workSpaceSize); if (!LaunchAttentionKernel( reinterpret_cast(gemm_buffer.get()), @@ -97,7 +98,11 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { temp_buffer.get(), cublas, element_size, - is_unidirectional_)) { + is_unidirectional_, + past_sequence_length, + nullptr == past ? nullptr : past->template Data(), + nullptr == present ? nullptr : present->template MutableData() + )) { // Get last error to reset it to cudaSuccess. CUDA_CALL(cudaGetLastError()); return Status(common::ONNXRUNTIME, common::FAIL); diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 6f583bf472..9e7a61b518 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -40,8 +40,8 @@ static size_t AlignTo(size_t a, size_t b) { return CeilDiv(a, b) * b; } -size_t ScratchSize(size_t element_size, int batch_size, int num_heads, int sequence_length) { - const size_t len = batch_size * num_heads * sequence_length * sequence_length; +size_t ScratchSize(size_t element_size, int batch_size, int num_heads, int sequence_length, int past_sequence_length) { + const size_t len = batch_size * num_heads * sequence_length * (sequence_length + past_sequence_length); const size_t bytes = len * element_size; const size_t alignment = 256; @@ -49,13 +49,19 @@ size_t ScratchSize(size_t element_size, int batch_size, int num_heads, int seque return bytesAligned; } -size_t GetAttentionWorkspaceSize(size_t element_size, int batch_size, int num_heads, int head_size, int sequence_length) { +size_t GetAttentionWorkspaceSize( + size_t element_size, + int batch_size, + int num_heads, + int head_size, + int sequence_length, + int past_sequence_length) { size_t qkv_size = 3 * batch_size * sequence_length * num_heads * head_size * element_size; - return qkv_size + 2 * ScratchSize(element_size, batch_size, num_heads, sequence_length); + return qkv_size + 2 * ScratchSize(element_size, batch_size, num_heads, sequence_length, past_sequence_length); } template -__device__ inline void Softmax(const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) { +__device__ inline void Softmax(const int past_sequence_length, const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) { using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage tmp_storage; @@ -64,13 +70,14 @@ __device__ inline void Softmax(const int sequence_length, const int valid_length float thread_data_max(-CUDART_INF_F); - const int num_valid = is_unidirectional ? (blockIdx.x % sequence_length) + 1 : valid_length; + const int num_valid = is_unidirectional ? past_sequence_length + (blockIdx.x % sequence_length) + 1 : valid_length; // e^x is represented as infinity if x is large enough, like 100.f. // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. // a math transform as below is leveraged to get a stable softmax: // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) - const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length; + const int all_sequence_length = past_sequence_length + sequence_length; + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; for (int i = threadIdx.x; i < num_valid; i += TPB) { const int index = offset + i; if (thread_data_max < float(input[index])) { @@ -99,7 +106,7 @@ __device__ inline void Softmax(const int sequence_length, const int valid_length } __syncthreads(); - for (int i = threadIdx.x; i < sequence_length; i += TPB) { + for (int i = threadIdx.x; i < all_sequence_length; i += TPB) { const int index = offset + i; const float val = (i < num_valid) ? expf(float(input[index]) - max_block) * sum_reverse_block : 0.f; output[index] = T(val); @@ -107,17 +114,19 @@ __device__ inline void Softmax(const int sequence_length, const int valid_length } template -__device__ inline void SoftmaxSmall(const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) { +__device__ inline void SoftmaxSmall(const int past_sequence_length, const int sequence_length, const int valid_length, const T* input, T* output, bool is_unidirectional) { using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage tmp_storage; __shared__ float sum_reverse_block; __shared__ float max_block; - const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length; + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + const int all_sequence_length = past_sequence_length + sequence_length; + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * all_sequence_length; const int index = offset + threadIdx.x; - const int num_valid = is_unidirectional ? (blockIdx.x % sequence_length) + 1 : valid_length; + const int num_valid = is_unidirectional ? past_sequence_length + (blockIdx.x % sequence_length) + 1 : valid_length; // e^x is represented as infinity if x is large enough, like 100.f. // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. @@ -146,43 +155,44 @@ __device__ inline void SoftmaxSmall(const int sequence_length, const int valid_l // Store max value if (threadIdx.x == 0) { - sum_reverse_block = (1.f) / sum; + sum_reverse_block = (num_valid == 0) ? 0.f : (1.f) / sum; } __syncthreads(); - if (threadIdx.x < sequence_length) { + // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. + if (threadIdx.x < all_sequence_length) { // this will be 0 for threadIdx.x >= num_valid output[index] = T(thread_data_exp * sum_reverse_block); } } template -__global__ void SoftmaxKernelSmall(const int sequence_length, const T* input, T* output, bool is_unidirectional) { - SoftmaxSmall(sequence_length, sequence_length, input, output, is_unidirectional); +__global__ void SoftmaxKernelSmall(const int past_sequence_length, const int sequence_length, const T* input, T* output, bool is_unidirectional) { + SoftmaxSmall(past_sequence_length, sequence_length, sequence_length, input, output, is_unidirectional); } template -__global__ void SoftmaxKernel(const int sequence_length, const T* input, T* output, bool is_unidirectional) { - Softmax(sequence_length, sequence_length, input, output, is_unidirectional); +__global__ void SoftmaxKernel(const int past_sequence_length, const int sequence_length, const T* input, T* output, bool is_unidirectional) { + Softmax(past_sequence_length, sequence_length, sequence_length, input, output, is_unidirectional); } template bool ComputeSoftmax( - cudaStream_t stream, const int sequence_length, const int batch_size, const int num_heads, + cudaStream_t stream, const int past_sequence_length, const int sequence_length, const int batch_size, const int num_heads, const T* input, T* output, bool is_unidirectional) { const dim3 grid(sequence_length * num_heads, batch_size, 1); if (sequence_length <= 32) { const int blockSize = 32; - SoftmaxKernelSmall<<>>(sequence_length, input, output, is_unidirectional); + SoftmaxKernelSmall<<>>(past_sequence_length, sequence_length, input, output, is_unidirectional); } else if (sequence_length <= 128) { const int blockSize = 128; - SoftmaxKernelSmall<<>>(sequence_length, input, output, is_unidirectional); + SoftmaxKernelSmall<<>>(past_sequence_length, sequence_length, input, output, is_unidirectional); } else if (sequence_length == 384) { const int blockSize = 384; - SoftmaxKernelSmall<<>>(sequence_length, input, output, is_unidirectional); + SoftmaxKernelSmall<<>>(past_sequence_length, sequence_length, input, output, is_unidirectional); } else { const int blockSize = 256; - SoftmaxKernel<<>>(sequence_length, input, output, is_unidirectional); + SoftmaxKernel<<>>(past_sequence_length, sequence_length, input, output, is_unidirectional); } return CUDA_CALL(cudaPeekAtLastError()); @@ -197,7 +207,7 @@ __global__ void MaskedSoftmaxKernelSmall(const int sequence_length, const int* m } __syncthreads(); - SoftmaxSmall(sequence_length, num_valid, input, output, false); + SoftmaxSmall(0, sequence_length, num_valid, input, output, false); } template @@ -209,7 +219,7 @@ __global__ void MaskedSoftmaxKernel(const int sequence_length, const int* mask_i } __syncthreads(); - Softmax(sequence_length, num_valid, input, output, false); + Softmax(0, sequence_length, num_valid, input, output, false); } template @@ -370,6 +380,89 @@ bool LaunchTransQkv(cudaStream_t stream, return CUDA_CALL(cudaPeekAtLastError()); } +template +__global__ void ConcatPastToPresent(const int sequence_length, + const T* past, + const T* k_v, + T* present) { + const int h = threadIdx.x; + const int n = threadIdx.y; + const int s = blockIdx.x; + const int b = blockIdx.y; + const int is_v = blockIdx.z; // 0 for k, 1 for v + + const int all_sequence_length = gridDim.x; + const int batch_size = gridDim.y; + const int num_heads = blockDim.y; + const int H = blockDim.x; + + // past: 2 x BxNxS'xH (past_k and past_v) + // k_v: 2 x BxNxSxH (k and v) + // present: 2 x BxNxS*xH (present_k and present_v) + const int past_sequence_length = all_sequence_length - sequence_length; + + const int present_SH = all_sequence_length * H; + const int present_NSH = num_heads * present_SH; + int out_offset = b * present_NSH + n * present_SH + s * H + h + is_v * (present_NSH * batch_size); + if (s < past_sequence_length) { + const int past_SH = past_sequence_length * H; + const int past_NSH = num_heads * past_SH; + const int in_offset = b * past_NSH + n * past_SH + s * H + h + is_v * (past_NSH * batch_size); + present[out_offset] = past[in_offset]; +} else if (s < all_sequence_length) { + const int SH = sequence_length * H; + const int NSH = num_heads * SH; + const int in_offset = b * NSH + n * SH + (s - past_sequence_length) * H + h + is_v * (NSH * batch_size); + present[out_offset] = k_v[in_offset]; + } +} + +bool LaunchConcatPastToPresent(cudaStream_t stream, + const int past_sequence_length, + const int sequence_length, + const int batch_size, + const int head_size, + const int num_heads, + const float* past, + const float* k_v, + float* present) { + const int all_sequence_length = past_sequence_length + sequence_length; + const dim3 grid(all_sequence_length, batch_size, 2); + if (0 == (head_size & 1)) { + const dim3 block(head_size / 2, num_heads, 1); + ConcatPastToPresent<<>>(sequence_length, reinterpret_cast(past), reinterpret_cast(k_v), reinterpret_cast(present)); + } else + { + const dim3 block(head_size, num_heads, 1); + ConcatPastToPresent<<>>(sequence_length, past, k_v, present); + } + return CUDA_CALL(cudaPeekAtLastError()); +} + +bool LaunchConcatPastToPresent(cudaStream_t stream, + const int past_sequence_length, + const int sequence_length, + const int batch_size, + const int head_size, + const int num_heads, + const half* past, + const half* k_v, + half* present) { + const int all_sequence_length = past_sequence_length + sequence_length; + const dim3 grid(all_sequence_length, batch_size, 2); + if (0 == (head_size % 4)) { + const dim3 block(head_size / 4, num_heads, 1); + ConcatPastToPresent<<>>(sequence_length, reinterpret_cast(past), reinterpret_cast(k_v), reinterpret_cast(present)); + } else if (0 == (head_size & 1)) { + const dim3 block(head_size / 2, num_heads, 1); + ConcatPastToPresent<<>>(sequence_length, reinterpret_cast(past), reinterpret_cast(k_v), reinterpret_cast(present)); + } else { // this should be an "odd" case. probably not worth catching it in the half2 kernel. + const dim3 block(head_size, num_heads, 1); + ConcatPastToPresent<<>>(sequence_length, past, k_v, present); + } + return CUDA_CALL(cudaPeekAtLastError()); +} + cublasStatus_t inline CublasGemmStridedBatched( cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const float alpha, @@ -394,8 +487,8 @@ bool QkvToContext( const int batch_size, const int sequence_length, const int num_heads, const int head_size, const size_t element_size, const T* input, T* output, T* workspace, const int* mask_index, - bool is_unidirectional) { - const size_t bytes = ScratchSize(element_size, batch_size, num_heads, sequence_length); + bool is_unidirectional, int past_sequence_length, const T* past, T* present) { + const size_t bytes = ScratchSize(element_size, batch_size, num_heads, sequence_length, past_sequence_length); T* scratch1 = workspace; T* scratch2 = scratch1 + (bytes / element_size); T* scratch3 = scratch2 + (bytes / element_size); @@ -409,7 +502,6 @@ bool QkvToContext( const int batches = batch_size * num_heads; const int size_per_batch = sequence_length * head_size; const int total_size = batches * size_per_batch; - const int temp_matrix_size = sequence_length * sequence_length; const T* q = scratch3; const T* k = q + total_size; @@ -418,29 +510,46 @@ bool QkvToContext( cublasSetStream(cublas, stream); CublasMathModeSetter helper(cublas, CUBLAS_TENSOR_OP_MATH); - // compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS + // Concat past (2xBxNxS'xH) to present (2xBxNxS*xH): + // past_k (BxNxS'xH) + k (BxNxSxH) => present_k (BxNxS*xH) + // past_v (BxNxS'xH) + v (BxNxSxH) => present_v (BxNxS*xH) + const int present_size_per_batch = (past_sequence_length + sequence_length) * head_size; + if (nullptr != present) { + if (!LaunchConcatPastToPresent(stream, past_sequence_length, sequence_length, batch_size, head_size, num_heads, past, k, present)) { + return false; + } + + // update pointers to present_k and present_v. + k = present; + v = present + batches * present_size_per_batch; + } + + // compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS* + // Q: BxNxSxH, K (present_k): BxNxS*xH, Q*K': BxNxSxS* const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); + const int all_sequence_length = past_sequence_length + sequence_length; + const int temp_matrix_size = sequence_length * all_sequence_length; if (!CUBLAS_CALL(CublasGemmStridedBatched( - cublas, CUBLAS_OP_T, CUBLAS_OP_N, sequence_length, sequence_length, head_size, rsqrt_head_size, k, head_size, size_per_batch, - q, head_size, size_per_batch, 0.f, scratch1, sequence_length, temp_matrix_size, batches))) { + cublas, CUBLAS_OP_T, CUBLAS_OP_N, all_sequence_length, sequence_length, head_size, rsqrt_head_size, k, head_size, present_size_per_batch, + q, head_size, size_per_batch, 0.f, scratch1, all_sequence_length, temp_matrix_size, batches))) { return false; } - // apply softmax and store result P to scratch2: BxNxSxS + // apply softmax and store result P to scratch2: BxNxSxS* if (nullptr != mask_index) { if (!ComputeMaskedSoftmax(stream, sequence_length, batch_size, num_heads, mask_index, scratch1, scratch2)) { return false; } } else { - if (!ComputeSoftmax(stream, sequence_length, batch_size, num_heads, scratch1, scratch2, is_unidirectional)) { + if (!ComputeSoftmax(stream, past_sequence_length, sequence_length, batch_size, num_heads, scratch1, scratch2, is_unidirectional)) { return false; } } // compute P*V (as V*P), and store in scratch3: BxNxSxH if (!CUBLAS_CALL(CublasGemmStridedBatched( - cublas, CUBLAS_OP_N, CUBLAS_OP_N, head_size, sequence_length, sequence_length, 1.f, v, head_size, size_per_batch, - scratch2, sequence_length, temp_matrix_size, 0.f, scratch3, head_size, size_per_batch, batches))) { + cublas, CUBLAS_OP_N, CUBLAS_OP_N, head_size, sequence_length, all_sequence_length, 1.f, v, head_size, present_size_per_batch, + scratch2, all_sequence_length, temp_matrix_size, 0.f, scratch3, head_size, size_per_batch, batches))) { return false; } @@ -459,7 +568,10 @@ bool LaunchAttentionKernel( void* workspace, cublasHandle_t& cublas, const size_t element_size, - bool is_unidirectional) { + bool is_unidirectional, + int past_sequence_length, + const void* past, + void* present) { // use default stream const cudaStream_t stream = nullptr; @@ -467,12 +579,14 @@ bool LaunchAttentionKernel( return QkvToContext(cublas, stream, batch_size, sequence_length, num_heads, head_size, element_size, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(workspace), - mask_index, is_unidirectional); + mask_index, is_unidirectional, + past_sequence_length, reinterpret_cast(past), reinterpret_cast(present)); } else { return QkvToContext(cublas, stream, batch_size, sequence_length, num_heads, head_size, element_size, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(workspace), - mask_index, is_unidirectional); + mask_index, is_unidirectional, + past_sequence_length, reinterpret_cast(past), reinterpret_cast(present)); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h index e106d8f266..6e58e73072 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h @@ -7,7 +7,13 @@ namespace onnxruntime { namespace contrib { namespace cuda { -size_t GetAttentionWorkspaceSize(size_t element_size, int batchsize, int num_heads, int head_size, int sequence_length); +size_t GetAttentionWorkspaceSize( + size_t element_size, + int batchsize, + int num_heads, + int head_size, + int sequence_length, + int past_sequence_length); bool LaunchAttentionKernel( const void* input, // Input tensor @@ -20,7 +26,10 @@ bool LaunchAttentionKernel( void* workspace, // Temporary buffer cublasHandle_t& cublas, // Cublas handle const size_t element_size, // Element size of input tensor - bool is_unidirectional // Whether there is unidirecitonal mask. + bool is_unidirectional, // Whether there is unidirecitonal mask. + int past_sequence_length, // Sequence length in past state + const void* past, // Past state input + void* present // Present state output ); } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc index 19b4d63723..6d4ed08d38 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc @@ -59,7 +59,7 @@ Status QAttention::CheckInputs(const Tensor* input, // Input 7 - weight_zero_point : scalar // Output : (batch_size, sequence_length, hidden_size) - ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index)); + ORT_RETURN_IF_ERROR(AttentionBase::CheckInputs(input, weights, bias, mask_index, nullptr)); ORT_RETURN_IF_NOT(IsScalarOr1ElementVector(input_scale_tensor), "input scale must be a scalar or 1D tensor of size 1"); @@ -161,7 +161,10 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { m, n); - size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length); + const int past_sequence_length = 0; + const T* past_data = nullptr; + T* present_data = nullptr; + size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length, past_sequence_length); auto temp_buffer = GetScratchBuffer(workSpaceSize); if (!LaunchAttentionKernel( reinterpret_cast(gemm_buffer.get()), @@ -174,7 +177,11 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { temp_buffer.get(), cublas, element_size, - is_unidirectional_)) { + is_unidirectional_, + past_sequence_length, + past_data, + present_data + )) { // Get last error to reset it to cudaSuccess. CUDA_CALL(cudaGetLastError()); return Status(common::ONNXRUNTIME, common::FAIL); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index fc0be49fa7..af85126923 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -309,10 +309,49 @@ mask_index shall not be provided.)DOC"; .Input(1, "weight", "2D input tensor with shape (hidden_size, 3 * hidden_size)", "T") .Input(2, "bias", "1D input tensor with shape (3 * hidden_size)", "T") .Input(3, "mask_index", "Attention mask index with shape (batch_size).", "M", OpSchema::Optional) - .Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T") + .Input(4, "past", "past state for key and value with shape (2, batch_size, num_heads, past_sequence_length, head_size).", "T", OpSchema::Optional) + .Output(0, "output", "3D output tensor with shape (batch_size, append_length, hidden_size)", "T") + .Output(1, "present", "present state for key and value with shape (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)", "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float tensors.") .TypeConstraint("M", {"tensor(int32)"}, "Constrain mask index to integer types") - .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput); + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (ctx.getNumOutputs() > 1) { + propagateElemTypeFromInputToOutput(ctx, 0, 1); + } + + if (hasInputShape(ctx, 0)) { + propagateShapeFromInputToOutput(ctx, 0, 0); + + if (ctx.getNumOutputs() > 1) { + auto& input_shape = getInputShape(ctx, 0); + auto& input_dims = input_shape.dim(); + if (input_dims.size() != 3) { + fail_shape_inference("Inputs 0 shall be 3 dimensions"); + } + + if (hasInputShape(ctx, 4)) { + auto& past_shape = getInputShape(ctx, 4); + auto& past_dims = past_shape.dim(); + if (past_dims.size() != 5) { + fail_shape_inference("Inputs 4 shall be 5 dimensions"); + } + + if (past_dims[3].has_dim_value() && input_dims[1].has_dim_value()) { + auto all_sequence_length = past_shape.dim(3).dim_value() + input_shape.dim(1).dim_value(); + + ONNX_NAMESPACE::TensorShapeProto present_shape; + for (auto& dim : past_dims) { + *present_shape.add_dim() = dim; + } + present_shape.mutable_dim(3)->set_dim_value(all_sequence_length); + + updateOutputShape(ctx, 1, present_shape); + } + } + } + } + }); ONNX_CONTRIB_OPERATOR_SCHEMA(QAttention) .SetDomain(kMSDomain) diff --git a/onnxruntime/python/tools/transformers/BertOnnxModel.py b/onnxruntime/python/tools/transformers/BertOnnxModel.py index 50ddd29263..1dcaad65f2 100644 --- a/onnxruntime/python/tools/transformers/BertOnnxModel.py +++ b/onnxruntime/python/tools/transformers/BertOnnxModel.py @@ -183,7 +183,7 @@ class BertOnnxModel(OnnxModel): for node in self.nodes(): # Before: # input_ids --> Shape --> Gather(indices=0) --> Unsqueeze ------+ - # | | + # | | # | v # +----> Shape --> Gather(indices=1) --> Unsqueeze---> Concat --> ConstantOfShape -->Cast --> EmbedLayerNormaliation/ReduceSum # After: @@ -292,8 +292,18 @@ class BertOnnxModel(OnnxModel): attention = op_count['Attention'] gelu = op_count['Gelu'] + op_count['BiasGelu'] + op_count['FastGelu'] layer_norm = op_count['LayerNormalization'] + op_count['SkipLayerNormalization'] - is_optimized = (embed > 0) and (attention > 0) and (attention == gelu) and (layer_norm >= 2 * attention) - logger.info( - f"EmbedLayer={embed}, Attention={attention}, Gelu={gelu}, LayerNormalization={layer_norm}, Successful={is_optimized}" - ) - return is_optimized + is_perfect = (embed > 0) and (attention > 0) and (attention == gelu) and (layer_norm >= 2 * attention) + + if layer_norm == 0: + logger.debug("Layer Normalization not fused") + + if gelu == 0: + logger.debug("Gelu/FastGelu not fused") + + if embed == 0: + logger.debug("Embed Layer not fused") + + if attention == 0: + logger.debug("Attention not fused") + + return is_perfect diff --git a/onnxruntime/python/tools/transformers/Gpt2OnnxModel.py b/onnxruntime/python/tools/transformers/Gpt2OnnxModel.py index e6922bce3b..02824451bf 100644 --- a/onnxruntime/python/tools/transformers/Gpt2OnnxModel.py +++ b/onnxruntime/python/tools/transformers/Gpt2OnnxModel.py @@ -10,6 +10,8 @@ import numpy as np from collections import deque from onnx import ModelProto, TensorProto, numpy_helper from BertOnnxModel import BertOnnxModel +from fusion_gpt_attention_no_past import FusionGptAttentionNoPast +from fusion_gpt_attention import FusionGptAttention logger = logging.getLogger(__name__) @@ -19,135 +21,12 @@ class Gpt2OnnxModel(BertOnnxModel): super().__init__(model, num_heads, hidden_size) def fuse_attention(self): - """ - Fuse Attention subgraph into one Attention node. - """ - logger.debug(f"start attention fusion...") - - input_name_to_nodes = self.input_name_to_nodes() - output_name_to_node = self.output_name_to_node() - - attention_count = 0 - - for normalize_node in self.get_nodes_by_op_type("LayerNormalization"): - return_indice = [] - qkv_nodes = self.match_parent_path( - normalize_node, - ['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'], - [0, None, 0, 0, 0, 0, 0], - output_name_to_node=output_name_to_node, - return_indice=return_indice - ) # yapf: disable - if qkv_nodes is None: - continue - (add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes - - another_input = add_qkv.input[1 - return_indice[0]] - - v_nodes = self.match_parent_path( - matmul_qkv, - ['Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'], - [1, 0, 0, 0, 0, 0]) # yapf: disable - if v_nodes is None: - logger.debug("fuse_attention: failed to match v path") - continue - (transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes - - layernorm_before_attention = self.get_parent(reshape_before_gemm, 0, output_name_to_node) - if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization': - logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}") - continue - - if not another_input in layernorm_before_attention.input: - logger.debug("Add and LayerNormalization shall have one same input") - continue - - qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0]) - if qk_nodes is not None: - (softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes - mask_nodes = self.match_parent_path( - sub_qk, - ['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], - [1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable - if mask_nodes is None: - logger.debug("fuse_attention: failed to match mask path") - continue - div_mask = mask_nodes[-1] - - if div_qk != div_mask: - logger.debug("fuse_attention: skip since div_qk != div_mask") - continue - else: - # New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0. - qk_nodes = self.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'Div', 'MatMul'], [0, 0, 1, 0]) - if qk_nodes is None: - logger.debug("fuse_attention: failed to match qk path") - continue - (softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes - mask_nodes = self.match_parent_path( - where_qk, - ['Cast', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], - [ 0, 0, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable - if mask_nodes is None: - logger.debug("fuse_attention: failed to match mask path") - continue - div_mask = mask_nodes[-1] - - if div_qk != div_mask: - logger.debug("fuse_attention: skip since div_qk != div_mask") - continue - - q_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0]) - if q_nodes is None: - logger.debug("fuse_attention: failed to match q path") - continue - (transpose_q, reshape_q, split_q) = q_nodes - if split_v != split_q: - logger.debug("fuse_attention: skip since split_v != split_q") - continue - - k_nodes = self.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [1, 0, 0]) - if k_nodes is None: - logger.debug("fuse_attention: failed to match k path") - continue - (transpose_k, reshape_k, split_k) = k_nodes - if split_v != split_k: - logger.debug("fuse_attention: skip since split_v != split_k") - continue - - self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0], - attention_count == 0) - # we rely on prune_graph() to clean old subgraph nodes: - # qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv] - attention_count += 1 - - self.prune_graph() - logger.info(f"Fused Attention count:{attention_count}") - - def create_attention_node(self, gemm, gemm_qkv, input, output, add_graph_input): - attention_node_name = self.create_node_name('Attention') - attention_node = onnx.helper.make_node('Attention', - inputs=[input, gemm.input[1], gemm.input[2]], - outputs=[attention_node_name + "_output"], - name=attention_node_name) - attention_node.domain = "com.microsoft" - attention_node.attribute.extend( - [onnx.helper.make_attribute("num_heads", self.num_heads), - onnx.helper.make_attribute("unidirectional", 1)]) - - matmul_node = onnx.helper.make_node('MatMul', - inputs=[attention_node_name + "_output", gemm_qkv.input[1]], - outputs=[attention_node_name + "_matmul_output"], - name=attention_node_name + "_matmul") - - add_node = onnx.helper.make_node('Add', - inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]], - outputs=[output], - name=attention_node_name + "_add") - - self.add_node(attention_node) - self.add_node(matmul_node) - self.add_node(add_node) + if len(self.model.graph.input) == 1 or len(self.model.graph.output) == 1: + fusion = FusionGptAttentionNoPast(self, self.num_heads) + fusion.apply() + else: + fusion = FusionGptAttention(self, self.num_heads) + fusion.apply() def postprocess(self): """ @@ -186,7 +65,12 @@ class Gpt2OnnxModel(BertOnnxModel): outputs=[add_node_name + "_output"], name=add_node_name) + self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output") + + # Link root node output with MatMul + self.replace_input_of_all_nodes(root_node.output[0], matmul_node_name + "_input") root_node.output[0] = matmul_node_name + "_input" + self.replace_input_of_all_nodes(reshape_after_gemm.output[0], add_node_name + "_output") self.add_node(matmul_node) diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py index ab16393a4d..4f321e93f0 100644 --- a/onnxruntime/python/tools/transformers/benchmark.py +++ b/onnxruntime/python/tools/transformers/benchmark.py @@ -222,12 +222,13 @@ def get_onnx_file_path(onnx_dir: str, model_name: str, input_count: int, optimiz def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwrite): if overwrite or not os.path.exists(ort_model_path): + from optimizer import optimize_model, get_fusion_statistics # Use onnxruntime to optimize model, which will be saved to *_ort.onnx opt_model = optimize_by_onnxruntime(onnx_model_path, use_gpu=use_gpu, optimized_model_path=ort_model_path, opt_level=99) - model_fusion_statistics[ort_model_path] = opt_model.get_fused_operator_statistics() + model_fusion_statistics[ort_model_path] = get_fusion_statistics(ort_model_path) else: logger.info(f"Skip optimization since model existed: {ort_model_path}") @@ -235,7 +236,7 @@ def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwri def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_attention_heads, hidden_size, use_gpu, fp16, overwrite): if overwrite or not os.path.exists(optimized_model_path): - from optimizer import optimize_model, optimize_by_onnxruntime + from optimizer import optimize_model from BertOnnxModel import BertOptimizationOptions optimization_options = BertOptimizationOptions(model_type) if fp16: diff --git a/onnxruntime/python/tools/transformers/benchmark_gpt2.py b/onnxruntime/python/tools/transformers/benchmark_gpt2.py index 2721f42b49..4722856f64 100644 --- a/onnxruntime/python/tools/transformers/benchmark_gpt2.py +++ b/onnxruntime/python/tools/transformers/benchmark_gpt2.py @@ -18,10 +18,10 @@ from transformers import GPT2Model, GPT2LMHeadModel, GPT2Tokenizer, AutoConfig logger = logging.getLogger('') -# Map alias to a tuple of Model Class and pretrained model name +# Map alias to a tuple of Model Class, Tokenizer, pretrained model name, use LMHead or not, use attention mask or not MODEL_CLASSES = { - "gpt2": (GPT2Model, GPT2Tokenizer, "gpt2"), - "distilgpt2": (GPT2LMHeadModel, GPT2Tokenizer, "distilgpt2"), + "gpt2": (GPT2Model, GPT2Tokenizer, "gpt2", False, False), + "distilgpt2": (GPT2LMHeadModel, GPT2Tokenizer, "distilgpt2", True, True), } @@ -44,12 +44,12 @@ def setup_environment(): dump_environment() -def pytorch_inference(model, input_ids, past=None, total_runs=100): +def pytorch_inference(model, input_ids, past=None, attention_mask=None, total_runs=100): latency = [] with torch.no_grad(): for _ in range(total_runs): start = time.time() - outputs = model(input_ids=input_ids, past=past) + outputs = model(input_ids=input_ids, past=past, attention_mask=attention_mask) latency.append(time.time() - start) average_latency = sum(latency) * 1000 / len(latency) @@ -57,10 +57,12 @@ def pytorch_inference(model, input_ids, past=None, total_runs=100): return outputs, average_latency -def onnxruntime_inference(ort_session, input_ids, past=None, total_runs=100): +def onnxruntime_inference(ort_session, input_ids, past=None, attention_mask=None, total_runs=100): ort_inputs = {'input_ids': numpy.ascontiguousarray(input_ids.cpu().numpy())} - # TODO: pass input tensor stored in GPU + if attention_mask is not None: + ort_inputs['attention_mask'] = numpy.ascontiguousarray(attention_mask.cpu().numpy()) + if past is not None: for i, past_i in enumerate(past): ort_inputs[f'past_{i}'] = numpy.ascontiguousarray(past[i].cpu().numpy()) @@ -82,6 +84,7 @@ def onnxruntime_inference_with_binded_io(ort_session, last_state, last_state_shape, past=None, + attention_mask=None, present=None, present_shape=None, total_runs=100): @@ -90,13 +93,17 @@ def onnxruntime_inference_with_binded_io(ort_session, # Bind inputs io_binding.bind_input('input_ids', input_ids.device.type, 0, numpy.longlong, list(input_ids.size()), input_ids.data_ptr()) + if attention_mask is not None: + io_binding.bind_input('attention_mask', attention_mask.device.type, 0, numpy.float32, list(attention_mask.size()), + attention_mask.data_ptr()) + if past is not None: for i, past_i in enumerate(past): io_binding.bind_input(f'past_{i}', past[i].device.type, 0, numpy.float32, list(past[i].size()), past[i].data_ptr()) # Bind outputs - io_binding.bind_output("last_state", last_state.device.type, 0, numpy.float32, last_state_shape, + io_binding.bind_output(ort_session.get_outputs()[0].name, last_state.device.type, 0, numpy.float32, last_state_shape, last_state.data_ptr()) if present is not None: for i, present_i in enumerate(present): @@ -125,6 +132,7 @@ def inference(model, ort_session, input_ids, past=None, + attention_mask=None, last_state=None, present=None, last_state_shape=None, @@ -132,12 +140,12 @@ def inference(model, total_runs=100, verify_outputs=True, with_io_binding=False): - outputs, torch_latency = pytorch_inference(model, input_ids, past, total_runs) - ort_outputs, ort_latency = onnxruntime_inference(ort_session, input_ids, past, total_runs) + outputs, torch_latency = pytorch_inference(model, input_ids, past, attention_mask, total_runs) + ort_outputs, ort_latency = onnxruntime_inference(ort_session, input_ids, past, attention_mask, total_runs) latencies = [torch_latency, ort_latency] if with_io_binding: ort_io_outputs, ort_io_latency = onnxruntime_inference_with_binded_io(ort_session, input_ids, last_state, - last_state_shape, past, present, + last_state_shape, past, attention_mask, present, present_shape, total_runs) latencies.append(ort_io_latency) if verify_outputs: @@ -203,7 +211,8 @@ def parse_arguments(): help='Use optimizer.py to optimize onnx model') parser.set_defaults(optimize_onnx=False) - parser.add_argument('--with_io_binding', + parser.add_argument('-i', + '--with_io_binding', required=False, action='store_true', help='Run ONNX Runtime with binded inputs and outputs. ') @@ -243,7 +252,9 @@ def setup_logger(verbose=True): logger.setLevel(logging_level) -def export_onnx(model, config, tokenizer, device, output_dir): +def export_onnx(model, config, tokenizer, device, output_dir, use_LMHead=False, use_attention_mask=False): + """ Export GPT-2 model with past state to ONNX model + """ model.to(device) inputs = tokenizer.encode_plus("Here is an example input for GPT2 model", @@ -251,56 +262,59 @@ def export_onnx(model, config, tokenizer, device, output_dir): return_tensors='pt') input_ids = inputs['input_ids'].to(device) logger.debug(f"input_ids={input_ids}") + + # Use example input to generate an example of past state. outputs = model(input_ids=input_ids, past=None) assert len(outputs) == 2 logger.debug(f"output 0 shape={outputs[0].shape}") logger.debug(f"outputs[1][0] shape={outputs[1][0].shape}") num_layer = model.config.n_layer - present_names = [f'present_{i}' for i in range(num_layer)] - output_names = ["last_state"] + present_names - - input_names = ['input_ids'] - - # input_ids has only one word for model with past state. - # Shape of input tensors: - # input_ids: (batch_size, 1) - # past_{i}: (2, batch_size, num_heads, seq_len, hidden_size/num_heads) - # Shape of output tensors: - # last_state: (batch_size, seq_len + 1, hidden_size) - # present_{i}: (2, batch_size, num_heads, seq_len + 1, hidden_size/num_heads) - dynamic_axes = {'input_ids': {0: 'batch_size'}, 'last_state': {0: 'batch_size', 1: 'seq_len_plus_1'}} - - for name in present_names: - dynamic_axes[name] = {1: 'batch_size', 3: 'seq_len_plus_1'} - past_names = [f'past_{i}' for i in range(num_layer)] - input_names = ['input_ids'] + past_names - dummy_past = [torch.zeros(list(outputs[1][0].shape), dtype=torch.float32, device=device) for _ in range(num_layer)] + present_names = [f'present_{i}' for i in range(num_layer)] + + # GPT2Model output last_state has shape (batch_size, all_seq_len, hidden_size) + # GPT2LMHeadModel output prediction_scores has shape (batch_size, all_seq_len, vocab_size) + # where all_seq_len = past_seq_len + seq_len + output_names = ["prediction_scores" if use_LMHead else "last_state"] + present_names + + # Shape of input tensors: + # input_ids: (batch_size, seq_len) + # past_{i}: (2, batch_size, num_heads, past_seq_len, hidden_size/num_heads) + # attention_mask: (batch_size, seq_len) + # Shape of output tensors: + # last_state: (batch_size, all_seq_len, hidden_size) + # present_{i}: (2, batch_size, num_heads, all_seq_len, hidden_size/num_heads) + dynamic_axes = {'input_ids': {0: 'batch_size', 1 : 'seq_len'}, output_names[0]: {0: 'batch_size', 1: 'all_seq_len'}} for name in past_names: - dynamic_axes[name] = {1: 'batch_size', 3: 'seq_len'} - logger.debug(f"vocab_size:{model.config.vocab_size}") + dynamic_axes[name] = {1: 'batch_size', 3: 'past_seq_len'} + for name in present_names: + dynamic_axes[name] = {1: 'batch_size', 3: 'all_seq_len'} + + if use_attention_mask: + dynamic_axes['attention_mask'] = {0: 'batch_size', 1: 'seq_len'} dummy_input_ids = torch.randint(low=0, high=model.config.vocab_size - 1, size=(1, 1), dtype=torch.int64, device=device) - logger.debug(f"dummy_input_ids={dummy_input_ids}") - export_inputs = (dummy_input_ids, tuple(dummy_past)) + # Use the example past state to create dummy past state inputs. + dummy_past = [torch.zeros(list(outputs[1][0].shape), dtype=torch.float32, device=device) for _ in range(num_layer)] - export_model_path = os.path.join(output_dir, 'gpt2_past.onnx') + dummy_mask = torch.ones([1, 1], dtype=torch.float32, device=device) if use_attention_mask else None + + model_name = "gpt2{}_past{}.onnx".format("_lm" if use_LMHead else "", "_mask" if use_attention_mask else "") + export_model_path = os.path.join(output_dir, model_name) - # Let's run performance test on PyTorch before updating environment variable. with torch.no_grad(): - outputs = model(input_ids=dummy_input_ids, past=dummy_past) - + outputs = model(input_ids=dummy_input_ids, past=dummy_past, attention_mask=dummy_mask) logger.debug(f"present_0 shape={outputs[1][0].shape}") torch.onnx.export(model, - args=export_inputs, + args=(dummy_input_ids, tuple(dummy_past), dummy_mask) if use_attention_mask else (dummy_input_ids, tuple(dummy_past)), f=export_model_path, - input_names=input_names, + input_names=['input_ids'] + past_names + (['attention_mask'] if use_attention_mask else []), output_names=output_names, example_outputs=outputs, dynamic_axes=dynamic_axes, @@ -324,7 +338,7 @@ def main(): os.makedirs(output_dir) use_torchscript = False - (model_class, tokenizer_class, model_name) = MODEL_CLASSES[args.model_type] + (model_class, tokenizer_class, model_name, use_LMHead, use_attention_mask) = MODEL_CLASSES[args.model_type] config = AutoConfig.from_pretrained(model_name, torchscript=use_torchscript, cache_dir=cache_dir) model = model_class.from_pretrained(model_name, config=config, cache_dir=cache_dir) tokenizer = tokenizer_class.from_pretrained(model_name, cache_dir=cache_dir) @@ -332,7 +346,7 @@ def main(): # model = torch.jit.trace(model, (input_ids, past)) device = torch.device("cuda:0" if args.use_gpu else "cpu") - export_model_path = export_onnx(model, config, tokenizer, device, output_dir) + export_model_path = export_onnx(model, config, tokenizer, device, output_dir, use_LMHead, use_attention_mask) # setup environment variables before importing onnxruntime. setup_environment() @@ -379,7 +393,10 @@ def main(): ] # dummy last state - last_state_size = numpy.prod([max_batch_size, 1, config.hidden_size]) + if use_LMHead: + last_state_size = numpy.prod([max_batch_size, 1, config.vocab_size]) + else: + last_state_size = numpy.prod([max_batch_size, 1, config.hidden_size]) dummy_last_state = torch.empty(last_state_size).to(device) for batch_size in args.batch_sizes: @@ -394,6 +411,7 @@ def main(): size=(batch_size, 1), dtype=torch.int64, device=device) + dummy_mask = torch.ones([batch_size, 1], dtype=torch.float32, device=device) if use_attention_mask else None # Calculate the expected output shapes last_state_shape = [batch_size, 1, config.hidden_size] @@ -406,6 +424,7 @@ def main(): session, dummy_input_ids, dummy_past, + dummy_mask, dummy_last_state, dummy_present, last_state_shape, diff --git a/onnxruntime/python/tools/transformers/bert_perf_test.py b/onnxruntime/python/tools/transformers/bert_perf_test.py index 124a750b3e..16d8a04939 100644 --- a/onnxruntime/python/tools/transformers/bert_perf_test.py +++ b/onnxruntime/python/tools/transformers/bert_perf_test.py @@ -411,7 +411,11 @@ def parse_arguments(): parser.add_argument('--input_ids_name', required=False, type=str, default=None, help="input name for input ids") parser.add_argument('--segment_ids_name', required=False, type=str, default=None, help="input name for segment ids") - parser.add_argument('--input_mask_name', required=False, type=str, default=None, help="input name for attention mask") + parser.add_argument('--input_mask_name', + required=False, + type=str, + default=None, + help="input name for attention mask") args = parser.parse_args() return args @@ -430,7 +434,8 @@ def main(): if not min(batch_size_set) >= 1 and max(batch_size_set) <= 128: raise Exception("batch_size not in range [1, 128]") - model_setting = ModelSetting(args.model, args.input_ids_name, args.segment_ids_name, args.input_mask_name, args.opt_level) + model_setting = ModelSetting(args.model, args.input_ids_name, args.segment_ids_name, args.input_mask_name, + args.opt_level) for batch_size in batch_size_set: test_setting = TestSetting( diff --git a/onnxruntime/python/tools/transformers/bert_test_data.py b/onnxruntime/python/tools/transformers/bert_test_data.py index 4dc7f74814..cb2ccb16e9 100644 --- a/onnxruntime/python/tools/transformers/bert_test_data.py +++ b/onnxruntime/python/tools/transformers/bert_test_data.py @@ -237,7 +237,11 @@ def parse_arguments(): parser.add_argument('--input_ids_name', required=False, type=str, default=None, help="input name for input ids") parser.add_argument('--segment_ids_name', required=False, type=str, default=None, help="input name for segment ids") - parser.add_argument('--input_mask_name', required=False, type=str, default=None, help="input name for attention mask") + parser.add_argument('--input_mask_name', + required=False, + type=str, + default=None, + help="input name for attention mask") parser.add_argument('--samples', required=False, type=int, default=1, help="number of test cases to be generated") diff --git a/onnxruntime/python/tools/transformers/dev_benchmark.cmd b/onnxruntime/python/tools/transformers/dev_benchmark.cmd index 7ef3f7dd00..192b5c8d4d 100644 --- a/onnxruntime/python/tools/transformers/dev_benchmark.cmd +++ b/onnxruntime/python/tools/transformers/dev_benchmark.cmd @@ -48,7 +48,7 @@ REM This script will generate a logs file with a list of commands used in tests. >benchmark.log echo echo ort=%run_ort% torch=%run_torch% torchscript=%run_torchscript% gpu_fp32=%run_gpu_fp32% gpu_fp16=%run_gpu_fp16% cpu=%run_cpu% optimizer=%use_optimizer% batch="%batch_sizes%" sequence="%sequence_length%" models="%models_to_test%" input_counts="%input_counts%" REM Set it to false to skip testing. You can use it to dry run this script with the benchmark.log file. -set run_tests=false +set run_tests=true REM ------------------------------------------- if %run_cpu% == true if %run_gpu_fp32% == true echo cannot test cpu and gpu at same time & goto :EOF diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py index 4ada0f690d..67a73b91df 100644 --- a/onnxruntime/python/tools/transformers/fusion_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_attention.py @@ -188,10 +188,8 @@ class FusionAttention(Fusion): # Note that Cast might be removed by OnnxRuntime so we match two patterns here. _, mask_nodes, _ = self.model.match_parent_paths( - add_qk, - [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]), - (['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], - output_name_to_node) + add_qk, [(['Mul', 'Sub', 'Cast', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0, 0]), + (['Mul', 'Sub', 'Unsqueeze', 'Unsqueeze'], [1, 0, 1, 0])], output_name_to_node) if mask_nodes is None: logger.debug("fuse_attention: failed to match mask path") return @@ -215,4 +213,3 @@ class FusionAttention(Fusion): # Use prune graph to remove mask nodes since they are shared by all attention nodes. #self.nodes_to_remove.extend(mask_nodes) self.prune_graph = True - \ No newline at end of file diff --git a/onnxruntime/python/tools/transformers/fusion_base.py b/onnxruntime/python/tools/transformers/fusion_base.py index c5f7ac1521..94b19eb7f7 100644 --- a/onnxruntime/python/tools/transformers/fusion_base.py +++ b/onnxruntime/python/tools/transformers/fusion_base.py @@ -10,17 +10,21 @@ logger = getLogger(__name__) class Fusion: - def __init__(self, model: OnnxModel, name: str, search_op_types: Union[str, List[str]]): + def __init__(self, + model: OnnxModel, + fused_op_type: str, + search_op_types: Union[str, List[str]], + description: str = None): self.search_op_types: List[str] = [search_op_types] if isinstance(search_op_types, str) else search_op_types - self.name: str = name + self.fused_op_type: str = fused_op_type + self.description: str = f"{fused_op_type}({description})" if description else fused_op_type self.model: OnnxModel = model self.nodes_to_remove: List = [] self.nodes_to_add: List = [] self.prune_graph: bool = False def apply(self): - logger.debug(f"start {self.name} fusion...") - + logger.debug(f"start {self.description} fusion...") input_name_to_nodes = self.model.input_name_to_nodes() output_name_to_node = self.model.output_name_to_node() @@ -29,7 +33,10 @@ class Fusion: for node in self.model.get_nodes_by_op_type(search_op_type): self.fuse(node, input_name_to_nodes, output_name_to_node) - logger.info(f"Fused {self.name} count: {len(self.nodes_to_add)}") + op_list = [node.op_type for node in self.nodes_to_add] + count = op_list.count(self.fused_op_type) + if count > 0: + logger.info(f"Fused {self.description} count: {count}") self.model.remove_nodes(self.nodes_to_remove) self.model.add_nodes(self.nodes_to_add) diff --git a/onnxruntime/python/tools/transformers/fusion_biasgelu.py b/onnxruntime/python/tools/transformers/fusion_biasgelu.py index 692c82f947..2d7fd3d45a 100644 --- a/onnxruntime/python/tools/transformers/fusion_biasgelu.py +++ b/onnxruntime/python/tools/transformers/fusion_biasgelu.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class FusionBiasGelu(Fusion): def __init__(self, model: OnnxModel, is_fastgelu): if is_fastgelu: - super().__init__(model, 'FastGelu(add bias)', 'FastGelu') + super().__init__(model, 'FastGelu', 'FastGelu', 'add bias') else: super().__init__(model, 'BiasGelu', 'Gelu') diff --git a/onnxruntime/python/tools/transformers/fusion_embedlayer.py b/onnxruntime/python/tools/transformers/fusion_embedlayer.py index f8f5c04839..92af1a749d 100644 --- a/onnxruntime/python/tools/transformers/fusion_embedlayer.py +++ b/onnxruntime/python/tools/transformers/fusion_embedlayer.py @@ -38,9 +38,8 @@ class FusionEmbedLayerNoMask(Fusion): """ def __init__(self, model: OnnxModel, - name: str = "EmbedLayerNormalization(no mask)", - search_op_types="SkipLayerNormalization"): - super().__init__(model, name, search_op_types) + description='no mask'): + super().__init__(model, "EmbedLayerNormalization", "SkipLayerNormalization", description) self.utils = FusionUtils(model) def fuse(self, node, input_name_to_nodes, output_name_to_node): @@ -141,10 +140,9 @@ class FusionEmbedLayerNoMask(Fusion): # Cast might be removed by OnnxRuntime. _, segment_id_path, _ = self.model.match_parent_paths( - segment_ids_cast_node, + segment_ids_cast_node, [(['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape', 'Cast'], [0, 0, 1, 0, 0, 0]), - (['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], - output_name_to_node) + (['ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather', 'Shape'], [0, 0, 1, 0, 0])], output_name_to_node) if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[0] == segment_id_path[-1].input[0]: logger.debug("Simplify semgent id path...") @@ -187,7 +185,7 @@ class FusionEmbedLayerNoMask(Fusion): class FusionEmbedLayerNormalization(FusionEmbedLayerNoMask): def __init__(self, model: OnnxModel, mask_indice: Dict, mask_casted: Dict): - super().__init__(model, "EmbedLayerNormalization(with mask)", "SkipLayerNormalization") + super().__init__(model, "with mask") self.mask_indice: Dict = mask_indice self.mask_casted: Dict = mask_casted self.mask_input_name = None diff --git a/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py b/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py index e336938920..32f2e0af5c 100644 --- a/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py +++ b/onnxruntime/python/tools/transformers/fusion_gelu_approximation.py @@ -11,7 +11,7 @@ from fusion_base import Fusion class FusionGeluApproximation(Fusion): def __init__(self, model: OnnxModel): - super().__init__(model, 'FastGelu(GeluApproximation)', ['Gelu', 'BiasGelu']) + super().__init__(model, 'FastGelu', ['Gelu', 'BiasGelu'], 'GeluApproximation') def fuse(self, node, input_name_to_nodes, output_name_to_node): new_node = helper.make_node("FastGelu", diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py new file mode 100644 index 0000000000..2d060fba82 --- /dev/null +++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py @@ -0,0 +1,196 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- +import numpy as np +from logging import getLogger +from onnx import helper, numpy_helper, TensorProto +from OnnxModel import OnnxModel +from fusion_base import Fusion +from fusion_utils import FusionUtils + +logger = getLogger(__name__) + + +class FusionGptAttention(Fusion): + """ + Fuse GPT-2 Attention with past state subgraph into one Attention node. + This does not support attention_mask graph input right now. + """ + def __init__(self, model: OnnxModel, num_heads: int): + super().__init__(model, "Attention", "LayerNormalization", "with past") + self.num_heads = num_heads + + def create_attention_node(self, gemm, gemm_qkv, past, present, input, output): + attention_node_name = self.model.create_node_name('GptAttention') + mask_index = '' + attention_node = helper.make_node('Attention', + inputs=[input, gemm.input[1], gemm.input[2], mask_index, past], + outputs=[attention_node_name + "_output", present], + name=attention_node_name) + attention_node.domain = "com.microsoft" + attention_node.attribute.extend( + [helper.make_attribute("num_heads", self.num_heads), + helper.make_attribute("unidirectional", 1)]) + + matmul_node = helper.make_node('MatMul', + inputs=[attention_node_name + "_output", gemm_qkv.input[1]], + outputs=[attention_node_name + "_matmul_output"], + name=attention_node_name + "_matmul") + + add_node = helper.make_node('Add', + inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]], + outputs=[output], + name=attention_node_name + "_add") + self.nodes_to_add.extend([attention_node, matmul_node, add_node]) + + def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node): + past = None + present = None + return_indice = [] + qkv_nodes = self.model.match_parent_path( + normalize_node, + ['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'], + [0, None, 0, 0, 0, 0, 0], + output_name_to_node=output_name_to_node, + return_indice=return_indice + ) # yapf: disable + if qkv_nodes is None: + return + (add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes + + another_input = add_qkv.input[1 - return_indice[0]] + + v_nodes = self.model.match_parent_path( + matmul_qkv, + ['Concat', 'Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'], + [1, 1, 0, 0, 0, 0, 0]) # yapf: disable + if v_nodes is None: + logger.debug("fuse_attention: failed to match v path") + return + (concat_v, transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes + + # concat <-- Gather(indices=1) <-- past + # | + # unsqueeze + # | + # concat --> present + gather_v = self.model.get_parent(concat_v, 0, output_name_to_node) + if gather_v.op_type != 'Gather': + logger.info("expect Gather for past") + return + if not self.model.find_constant_input(gather_v, 1) == 1: + logger.info("expect indices=1 for Gather of past") + return + past = gather_v.input[0] + if not self.model.find_graph_input(past): + logger.info("expect past to be graph input") + return + unsqueeze_present_v = self.model.find_first_child_by_type(concat_v, + 'Unsqueeze', + input_name_to_nodes, + recursive=False) + if not unsqueeze_present_v: + logger.info("expect unsqueeze for present") + return + concat_present = self.model.find_first_child_by_type(unsqueeze_present_v, + 'Concat', + input_name_to_nodes, + recursive=False) + if not concat_present: + logger.info("expect concat for present") + return + present = concat_present.output[0] + if not self.model.find_graph_output(present): + logger.info("expect present to be graph input") + return + + layernorm_before_attention = self.model.get_parent(reshape_before_gemm, 0, output_name_to_node) + if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization': + logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}") + return + + if not another_input in layernorm_before_attention.input: + logger.debug("Add and LayerNormalization shall have one same input") + return + + qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0]) + if qk_nodes is not None: + (softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes + mask_nodes = self.model.match_parent_path( + sub_qk, + ['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], + [1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable + if mask_nodes is None: + logger.debug("fuse_attention: failed to match mask path") + return + div_mask = mask_nodes[-1] + + if div_qk != div_mask: + logger.debug("fuse_attention: skip since div_qk != div_mask") + return + else: + # New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0. + qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'Div', 'MatMul'], [0, 0, 1, 0]) + if qk_nodes is None: + logger.debug("fuse_attention: failed to match qk path") + return + (softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes + mask_nodes = self.model.match_parent_path( + where_qk, + ['Cast', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], + [ 0, 0, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable + if mask_nodes is None: + logger.debug("fuse_attention: failed to match mask path") + return + div_mask = mask_nodes[-1] + + if div_qk != div_mask: + logger.debug("fuse_attention: skip since div_qk != div_mask") + return + + q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0]) + if q_nodes is None: + logger.debug("fuse_attention: failed to match q path") + return + (transpose_q, reshape_q, split_q) = q_nodes + if split_v != split_q: + logger.debug("fuse_attention: skip since split_v != split_q") + return + + k_nodes = self.model.match_parent_path(matmul_qk, ['Concat', 'Transpose', 'Reshape', 'Split'], [1, 1, 0, 0]) + if k_nodes is None: + logger.debug("fuse_attention: failed to match k path") + return + (concat_k, transpose_k, reshape_k, split_k) = k_nodes + if split_v != split_k: + logger.debug("fuse_attention: skip since split_v != split_k") + return + + # concat_k <-- Transpose (perm=0,1,3,2) <-- Gather(axes=0, indices=0) <-- past + # | + # Transpose (perm=0,1,3,2) + # | + # unsqueeze + # | + # concat --> present + past_k_nodes = self.model.match_parent_path(concat_k, ['Transpose', 'Gather'], [0, 0]) + if past_k_nodes is None: + logger.debug("fuse_attention: failed to match past_k_nodes path") + return + + gather_past_k = past_k_nodes[-1] + if not self.model.find_constant_input(gather_past_k, 0) == 1: + logger.info("expect indices=0 for Gather k of past") + return + past_k = gather_past_k.input[0] + if past != past_k: + logger.info("expect past to be same") + return + + self.create_attention_node(gemm, gemm_qkv, past, present, layernorm_before_attention.output[0], + reshape_qkv.output[0]) + + # we rely on prune_graph() to clean old subgraph nodes: + # qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv] + self.prune_graph = True diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py new file mode 100644 index 0000000000..9585e3739a --- /dev/null +++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention_no_past.py @@ -0,0 +1,137 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- +import numpy as np +from logging import getLogger +from onnx import helper, numpy_helper, TensorProto +from OnnxModel import OnnxModel +from fusion_base import Fusion +from fusion_utils import FusionUtils + +logger = getLogger(__name__) + + +class FusionGptAttentionNoPast(Fusion): + """ + Fuse GPT-2 Attention without past state into one Attention node. + This does not support attention_mask graph input right now. + """ + def __init__(self, model: OnnxModel, num_heads: int): + super().__init__(model, "Attention", "LayerNormalization", "without past") + self.num_heads = num_heads + + def create_attention_node(self, gemm, gemm_qkv, input, output): + attention_node_name = self.model.create_node_name('Attention') + attention_node = helper.make_node('Attention', + inputs=[input, gemm.input[1], gemm.input[2]], + outputs=[attention_node_name + "_output"], + name=attention_node_name) + attention_node.domain = "com.microsoft" + attention_node.attribute.extend( + [helper.make_attribute("num_heads", self.num_heads), + helper.make_attribute("unidirectional", 1)]) + + matmul_node = helper.make_node('MatMul', + inputs=[attention_node_name + "_output", gemm_qkv.input[1]], + outputs=[attention_node_name + "_matmul_output"], + name=attention_node_name + "_matmul") + + add_node = helper.make_node('Add', + inputs=[attention_node_name + "_matmul_output", gemm_qkv.input[2]], + outputs=[output], + name=attention_node_name + "_add") + + self.nodes_to_add.extend([attention_node, matmul_node, add_node]) + + def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node): + return_indice = [] + qkv_nodes = self.model.match_parent_path( + normalize_node, + ['Add', 'Reshape', 'Gemm', 'Reshape', 'Reshape', 'Transpose', 'MatMul'], + [0, None, 0, 0, 0, 0, 0], + output_name_to_node=output_name_to_node, + return_indice=return_indice + ) # yapf: disable + if qkv_nodes is None: + return + (add_qkv, reshape_qkv, gemm_qkv, reshape_1, reshape_2, transpose_qkv, matmul_qkv) = qkv_nodes + + another_input = add_qkv.input[1 - return_indice[0]] + + v_nodes = self.model.match_parent_path( + matmul_qkv, + ['Transpose', 'Reshape', 'Split', 'Reshape', 'Gemm', 'Reshape'], + [1, 0, 0, 0, 0, 0]) # yapf: disable + if v_nodes is None: + logger.debug("fuse_attention: failed to match v path") + return + (transpose_v, reshape_v, split_v, reshape_after_gemm, gemm, reshape_before_gemm) = v_nodes + + layernorm_before_attention = self.model.get_parent(reshape_before_gemm, 0, output_name_to_node) + if layernorm_before_attention is None or layernorm_before_attention.op_type != 'LayerNormalization': + logger.debug(f"failed to get layernorm before gemm. Got {layernorm_before_attention.op_type}") + return + + if not another_input in layernorm_before_attention.input: + logger.debug("Add and LayerNormalization shall have one same input") + return + + qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Sub', 'Mul', 'Div', 'MatMul'], [0, 0, 0, 0, 0]) + if qk_nodes is not None: + (softmax_qk, sub_qk, mul_qk, div_qk, matmul_qk) = qk_nodes + mask_nodes = self.model.match_parent_path( + sub_qk, + ['Mul', 'Sub', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], + [1, 0, 1, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable + if mask_nodes is None: + logger.debug("fuse_attention: failed to match mask path") + return + div_mask = mask_nodes[-1] + + if div_qk != div_mask: + logger.debug("fuse_attention: skip since div_qk != div_mask") + return + else: + # New pattern for gpt2 from PyTorch 1.5.0 and Transformers 2.9.0. + qk_nodes = self.model.match_parent_path(matmul_qkv, ['Softmax', 'Where', 'Div', 'MatMul'], [0, 0, 1, 0]) + if qk_nodes is None: + logger.debug("fuse_attention: failed to match qk path") + return + (softmax_qk, where_qk, div_qk, matmul_qk) = qk_nodes + mask_nodes = self.model.match_parent_path( + where_qk, + ['Cast', 'Slice', 'Slice', 'Unsqueeze', 'Sub', 'Squeeze', 'Slice', 'Shape', 'Div'], + [ 0, 0, 0, 1, 0, 0, 0, 0, 0]) # yapf: disable + if mask_nodes is None: + logger.debug("fuse_attention: failed to match mask path") + return + div_mask = mask_nodes[-1] + + if div_qk != div_mask: + logger.debug("fuse_attention: skip since div_qk != div_mask") + return + + q_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [0, 0, 0]) + if q_nodes is None: + logger.debug("fuse_attention: failed to match q path") + return + (transpose_q, reshape_q, split_q) = q_nodes + if split_v != split_q: + logger.debug("fuse_attention: skip since split_v != split_q") + return + + k_nodes = self.model.match_parent_path(matmul_qk, ['Transpose', 'Reshape', 'Split'], [1, 0, 0]) + if k_nodes is None: + logger.debug("fuse_attention: failed to match k path") + return + (transpose_k, reshape_k, split_k) = k_nodes + if split_v != split_k: + logger.debug("fuse_attention: skip since split_v != split_k") + return + + self.create_attention_node(gemm, gemm_qkv, layernorm_before_attention.output[0], reshape_qkv.output[0]) + + # we rely on prune_graph() to clean old subgraph nodes: + # qk_nodes + q_nodes + k_nodes + v_nodes + mask_nodes + [reshape_qkv, transpose_qkv, matmul_qkv] + self.prune_graph = True diff --git a/onnxruntime/python/tools/transformers/fusion_layernorm.py b/onnxruntime/python/tools/transformers/fusion_layernorm.py index efabdaea90..3be6747bf2 100644 --- a/onnxruntime/python/tools/transformers/fusion_layernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_layernorm.py @@ -117,7 +117,7 @@ class FusionLayerNormalization(Fusion): class FusionLayerNormalizationTF(Fusion): def __init__(self, model: OnnxModel): - super().__init__(model, "LayerNormalization", "Add") + super().__init__(model, "LayerNormalization", "Add", "TF") def fuse(self, node, input_name_to_nodes: Dict, output_name_to_node: Dict): """ diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 2693f96dd0..93e52d37ae 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -47,7 +47,7 @@ class FusionSkipLayerNormalization(Fusion): class FusionBiasSkipLayerNormalization(Fusion): def __init__(self, model: OnnxModel): - super().__init__(model, "SkipLayerNormalization(add bias)", "SkipLayerNormalization") + super().__init__(model, "SkipLayerNormalization", "SkipLayerNormalization", "add bias") def fuse(self, node, input_name_to_nodes, output_name_to_node): if len(node.input) != 4: diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py index b64a56e841..c4f1e0c07c 100644 --- a/onnxruntime/python/tools/transformers/optimizer.py +++ b/onnxruntime/python/tools/transformers/optimizer.py @@ -6,22 +6,17 @@ # Convert Bert ONNX model converted from TensorFlow or exported from PyTorch to use Attention, Gelu, # SkipLayerNormalization and EmbedLayerNormalization ops to optimize # performance on NVidia GPU and CPU. - +# # For Bert model exported from PyTorch, OnnxRuntime has bert model optimization support internally. -# You can use the option --use_onnxruntime to use model optimization from OnnxRuntime package. +# You can use the option --use_onnxruntime to check optimizations from OnnxRuntime. # For Bert model file like name.onnx, optimized model for GPU or CPU from OnnxRuntime will output as # name_ort_gpu.onnx or name_ort_cpu.onnx in the same directory. +# # This script is retained for experiment purpose. Useful senarios like the following: -# (1) Change model from fp32 to fp16. +# (1) Change model from fp32 to fp16 for mixed precision inference in GPU with Tensor Core. # (2) Change input data type from int64 to int32. # (3) Some model cannot be handled by OnnxRuntime, and you can modify this script to get optimized model. -# This script has been tested using the following models: -# (1) BertForSequenceClassification as in https://github.com/huggingface/transformers/blob/master/examples/run_glue.py -# PyTorch 1.2 or above, and exported to Onnx using opset version 10 or 11. -# (2) BertForQuestionAnswering as in https://github.com/huggingface/transformers/blob/master/examples/run_squad.py -# PyTorch 1.2 or above, and exported to Onnx using opset version 10 or 11. - import logging import coloredlogs import onnx @@ -29,6 +24,7 @@ import os import sys import argparse import numpy as np +from typing import Dict from collections import deque from onnx import ModelProto, TensorProto, numpy_helper, load_model from BertOnnxModel import BertOnnxModel, BertOptimizationOptions @@ -47,17 +43,21 @@ MODEL_CLASSES = { } -def optimize_by_onnxruntime(onnx_model_path, use_gpu=False, optimized_model_path=None, opt_level=99): +def optimize_by_onnxruntime(onnx_model_path: str, + use_gpu: bool = False, + optimized_model_path: str = None, + opt_level: int = 99) -> str: """ - Use onnxruntime package to optimize model. It could support models exported by PyTorch. + Use onnxruntime to optimize model. Args: - onnx_model_path (str): th path of input onnx model. + onnx_model_path (str): the path of input onnx model. use_gpu (bool): whether the optimized model is targeted to run in GPU. optimized_model_path (str or None): the path of optimized model. + opt_level (int): graph optimization level. Returns: - optimized_model_path: the path of optimized model + optimized_model_path (str): the path of optimized model """ import onnxruntime @@ -91,7 +91,22 @@ def optimize_by_onnxruntime(onnx_model_path, use_gpu=False, optimized_model_path return optimized_model_path -def parse_arguments(): +def get_fusion_statistics(optimized_model_path: str) -> Dict[str, int]: + """ + Get counter of fused operators in optimized model. + + Args: + optimized_model_path (str): the path of onnx model. + + Returns: + A dictionary with operator type as key, and count as value + """ + model = load_model(optimized_model_path, format=None, load_external_data=True) + optimizer = BertOnnxModel(model, num_heads=12, hidden_size=768) + return optimizer.get_fused_operator_statistics() + + +def _parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, type=str, help="input onnx model path") @@ -192,7 +207,7 @@ def parse_arguments(): return args -def get_optimization_options(args): +def _get_optimization_options(args): optimization_options = BertOptimizationOptions(args.model_type) if args.disable_gelu: optimization_options.enable_gelu = False @@ -214,18 +229,38 @@ def get_optimization_options(args): def optimize_model(input, - model_type, - num_heads, - hidden_size, - opt_level=0, + model_type='bert', + num_heads=12, + hidden_size=768, optimization_options=None, + opt_level=0, use_gpu=False, only_onnxruntime=False): + """ Optimize Model by OnnxRuntime and/or offline fusion logic. + + The following optimizes model by OnnxRuntime only, and no offline fusion logic: + optimize_model(input, opt_level=1, use_gpu=False, only_onnxruntime=True) + If you want to optimize model by offline fusion logic. + optimize_model(input, model_type, num_heads=12, hidden_size=768, optimization_options=your_options) + + Args: + input (str): input model path. + model_type (str): model type - like bert, bert_tf, bert_keras or gpt2. + num_heads (int): number of attention heads. + hidden_size (int): hidden size. + optimization_options (OptimizationOptions or None): optimization options that can use to turn on/off some fusions. + opt_level (int): onnxruntime graph optimization level (0, 1, 2 or 99). When the level > 0, onnxruntime will be used to optimize model first. + use_gpu (bool): use gpu or not for onnxruntime. + only_onnxruntime (bool): only use onnxruntime to optimize model, and no offline fusion logic is used. + + Returns: + object of an optimizer class. + """ (optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type] input_model_path = input - if opt_level > 1: # Optimization specified for an execution provider. + if opt_level > 1: # Optimization specified for an execution provider. input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=use_gpu, opt_level=opt_level) elif run_onnxruntime: # Use Onnxruntime to do optimizations (like constant folding and cast elimation) that is not specified to exection provider. @@ -242,15 +277,15 @@ def optimize_model(input, if optimization_options is None: optimization_options = BertOptimizationOptions(model_type) - bert_model = optimizer_class(model, num_heads, hidden_size) + optimizer = optimizer_class(model, num_heads, hidden_size) if not only_onnxruntime: - bert_model.optimize(optimization_options) + optimizer.optimize(optimization_options) - return bert_model + return optimizer -def setup_logger(verbose): +def _setup_logger(verbose): if verbose: coloredlogs.install(level='DEBUG', fmt='[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s') else: @@ -258,33 +293,33 @@ def setup_logger(verbose): def main(): - args = parse_arguments() + args = _parse_arguments() - setup_logger(args.verbose) + _setup_logger(args.verbose) - optimization_options = get_optimization_options(args) + optimization_options = _get_optimization_options(args) - bert_model = optimize_model(args.input, - args.model_type, - args.num_heads, - args.hidden_size, - opt_level=args.opt_level, - optimization_options=optimization_options, - use_gpu=args.use_gpu, - only_onnxruntime=args.only_onnxruntime) + optimizer = optimize_model(args.input, + args.model_type, + args.num_heads, + args.hidden_size, + opt_level=args.opt_level, + optimization_options=optimization_options, + use_gpu=args.use_gpu, + only_onnxruntime=args.only_onnxruntime) if args.float16: - bert_model.convert_model_float32_to_float16() + optimizer.convert_model_float32_to_float16() if args.input_int32: - bert_model.change_input_to_int32() + optimizer.change_input_to_int32() - bert_model.save_model_to_file(args.output) + optimizer.save_model_to_file(args.output) - if bert_model.is_fully_optimized(): - logger.info("The output model is fully optimized.") + if optimizer.is_fully_optimized(): + logger.info("The model has been fully optimized.") else: - logger.warning("The output model is not fully optimized. It might not be usable.") + logger.info("The model has been optimized.") if __name__ == "__main__": diff --git a/onnxruntime/python/tools/transformers/pytest.ini b/onnxruntime/python/tools/transformers/pytest.ini new file mode 100644 index 0000000000..0d98a076eb --- /dev/null +++ b/onnxruntime/python/tools/transformers/pytest.ini @@ -0,0 +1,9 @@ +[pytest] +log_cli = 1 +log_cli_level = INFO +log_cli_format = %(message)s + +log_file = pytest.log +log_file_level = DEBUG +log_file_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s) +log_file_date_format=%Y-%m-%d %H:%M:%S \ No newline at end of file diff --git a/onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/generate_tiny_gpt2_model.py b/onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/generate_tiny_gpt2_model.py new file mode 100644 index 0000000000..3ee3147235 --- /dev/null +++ b/onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/generate_tiny_gpt2_model.py @@ -0,0 +1,363 @@ +#------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +#-------------------------------------------------------------------------- +# This tool generates a tiny GPT2 model for testing fusion script. +# You can use benchmark_gpt2.py to get a gpt2 ONNX model as input of this tool. + +import onnx +import onnx.utils +import sys +import argparse +import numpy as np +from onnx import ModelProto, TensorProto, numpy_helper +from OnnxModel import OnnxModel +import os +import onnxruntime +import random +from pathlib import Path +import timeit + +DICT_SIZE = 20 +SEQ_LEN = 5 +""" This class creates a tiny bert model for test purpose. """ + +# parameters of input base model. +old_parameters = { + "seq_len": 5, + "hidden_size": 768, + "num_heads": 12, + "size_per_head": 64, + "word_dict_size": [50257], # list of supported dictionary size. + "max_word_position": 1024 +} + +# parameters of output tiny model. +new_parameters = { + "seq_len": SEQ_LEN, + "hidden_size": 4, + "num_heads": 2, + "size_per_head": 2, + "word_dict_size": DICT_SIZE, + "max_word_position": 8 +} + + +class TinyGpt2Model(OnnxModel): + def __init__(self, model): + super(TinyGpt2Model, self).__init__(model) + self.resize_model() + + def resize_weight(self, initializer_name, target_shape): + weight = self.get_initializer(initializer_name) + w = numpy_helper.to_array(weight) + + target_w = w + if len(target_shape) == 1: + target_w = w[:target_shape[0]] + elif len(target_shape) == 2: + target_w = w[:target_shape[0], :target_shape[1]] + elif len(target_shape) == 3: + target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]] + elif len(target_shape) == 4: + target_w = w[:target_shape[0], :target_shape[1], :target_shape[2], :target_shape[3]] + else: + print("at most 3 dimensions") + + tensor = onnx.helper.make_tensor(name=initializer_name + '_resize', + data_type=TensorProto.FLOAT, + dims=target_shape, + vals=target_w.flatten().tolist()) + + return tensor + + def resize_model(self): + graph = self.model.graph + initializers = graph.initializer + + for input in graph.input: + if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]): + print("input", input.name, input.type.tensor_type.shape) + input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"] + print("=>", input.type.tensor_type.shape) + + reshapes = {} + for initializer in initializers: + tensor = numpy_helper.to_array(initializer) + if initializer.data_type == TensorProto.FLOAT: + dtype = np.float32 + elif initializer.data_type == TensorProto.INT32: + dtype = np.int32 + elif initializer.data_type == TensorProto.INT64: + dtype = np.int64 + else: + print("data type not supported by this tool:", dtype) + + if len(tensor.shape) == 1 and tensor.shape[0] == 1: + if tensor == old_parameters["num_heads"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype), + initializer.name)) + elif tensor == old_parameters["seq_len"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name)) + elif tensor == old_parameters["size_per_head"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype), + initializer.name)) + elif tensor == old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype), + initializer.name)) + elif tensor == 4 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype), + initializer.name)) + elif tensor == 3 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 3 * old_parameters["hidden_size"], "=>[", 3 * new_parameters["hidden_size"], "]") + initializer.CopyFrom( + numpy_helper.from_array(np.asarray([3 * new_parameters["hidden_size"]], dtype=dtype), + initializer.name)) + elif len(tensor.shape) == 0: + if tensor == old_parameters["num_heads"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["num_heads"], "=>", new_parameters["num_heads"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name)) + elif tensor == old_parameters["seq_len"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["seq_len"], "=>", new_parameters["seq_len"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name)) + elif tensor == old_parameters["size_per_head"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["size_per_head"], "=>", new_parameters["size_per_head"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype), + initializer.name)) + elif tensor == old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + old_parameters["hidden_size"], "=>", new_parameters["hidden_size"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype), + initializer.name)) + elif tensor == 4 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype), + initializer.name)) + elif tensor == 3 * old_parameters["hidden_size"]: + print("initializer type={}".format(initializer.data_type), initializer.name, + 3 * old_parameters["hidden_size"], "=>", 3 * new_parameters["hidden_size"]) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(3 * new_parameters["hidden_size"], dtype=dtype), + initializer.name)) + elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]): + print("initializer type={}".format(initializer.data_type), initializer.name, + 1.0 / np.sqrt(old_parameters["size_per_head"]), "=>", + 1.0 / np.sqrt(new_parameters["size_per_head"])) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype), + initializer.name)) + elif tensor == np.sqrt(old_parameters["size_per_head"]): + print("initializer type={}".format(initializer.data_type), initializer.name, + np.sqrt(old_parameters["size_per_head"]), "=>", np.sqrt(new_parameters["size_per_head"])) + initializer.CopyFrom( + numpy_helper.from_array(np.asarray(np.sqrt(new_parameters["size_per_head"]), dtype=dtype), + initializer.name)) + + new_shape = [] + shape_changed = False + for dim in tensor.shape: + if (dim == old_parameters["hidden_size"]): + new_shape.append(new_parameters["hidden_size"]) + shape_changed = True + elif (dim == 4 * old_parameters["hidden_size"]): + new_shape.append(4 * new_parameters["hidden_size"]) + shape_changed = True + elif (dim == 3 * old_parameters["hidden_size"]): + new_shape.append(3 * new_parameters["hidden_size"]) + shape_changed = True + elif (dim in old_parameters["word_dict_size"]): + new_shape.append(new_parameters["word_dict_size"]) + shape_changed = True + elif (dim == old_parameters["max_word_position"]): + new_shape.append(new_parameters["max_word_position"]) + shape_changed = True + else: + new_shape.append(dim) + if shape_changed: + reshapes[initializer.name] = new_shape + print("initializer", initializer.name, tensor.shape, "=>", new_shape) + + for initializer_name in reshapes: + self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize') + tensor = self.resize_weight(initializer_name, reshapes[initializer_name]) + self.model.graph.initializer.extend([tensor]) + + # Add node name, replace split node attribute. + nodes_to_add = [] + nodes_to_remove = [] + for i, node in enumerate(graph.node): + if node.op_type == "Split": + nodes_to_add.append( + onnx.helper.make_node('Split', + node.input, + node.output, + name="Split_{}".format(i), + axis=2, + split=[ + new_parameters["hidden_size"], new_parameters["hidden_size"], + new_parameters["hidden_size"] + ])) + nodes_to_remove.append(node) + print("update split", + [new_parameters["hidden_size"], new_parameters["hidden_size"], new_parameters["hidden_size"]]) + if node.op_type == "Constant": + for att in node.attribute: + if att.name == 'value': + if numpy_helper.to_array(att.t) == old_parameters["num_heads"]: + nodes_to_add.append( + onnx.helper.make_node('Constant', + inputs=node.input, + outputs=node.output, + value=onnx.helper.make_tensor(name=att.t.name, + data_type=TensorProto.INT64, + dims=[], + vals=[new_parameters["num_heads"] + ]))) + print("constant", att.t.name, old_parameters["num_heads"], "=>", + new_parameters["num_heads"]) + if numpy_helper.to_array(att.t) == np.sqrt(old_parameters["size_per_head"]): + nodes_to_add.append( + onnx.helper.make_node('Constant', + inputs=node.input, + outputs=node.output, + value=onnx.helper.make_tensor( + name=att.t.name, + data_type=TensorProto.FLOAT, + dims=[], + vals=[np.sqrt(new_parameters["size_per_head"])]))) + print("constant", att.t.name, np.sqrt(old_parameters["size_per_head"]), "=>", + np.sqrt(new_parameters["size_per_head"])) + else: + node.name = node.op_type + "_" + str(i) + for node in nodes_to_remove: + graph.node.remove(node) + graph.node.extend(nodes_to_add) + + for i, input in enumerate(self.model.graph.input): + if i > 0: + dim_proto = input.type.tensor_type.shape.dim[2] + dim_proto.dim_value = new_parameters["num_heads"] + dim_proto = input.type.tensor_type.shape.dim[4] + dim_proto.dim_value = new_parameters["size_per_head"] + + for i, output in enumerate(self.model.graph.output): + if i == 0: + dim_proto = output.type.tensor_type.shape.dim[2] + dim_proto.dim_value = new_parameters["hidden_size"] + if i > 0: + dim_proto = output.type.tensor_type.shape.dim[2] + dim_proto.dim_value = new_parameters["num_heads"] + dim_proto = output.type.tensor_type.shape.dim[4] + dim_proto.dim_value = new_parameters["size_per_head"] + + +def generate_test_data(onnx_file, + output_path, + batch_size=1, + use_cpu=True, + input_tensor_only=False, + dictionary_size=DICT_SIZE, + test_cases=1, + output_optimized_model=False): + + for test_case in range(test_cases): + sequence_length = 3 + input_1 = np.random.randint(dictionary_size, size=(batch_size, 1), dtype=np.int64) + tensor_1 = numpy_helper.from_array(input_1, 'input_ids') + + path = os.path.join(output_path, 'test_data_set_' + str(test_case)) + try: + os.mkdir(path) + except OSError: + print("Creation of the directory %s failed" % path) + else: + print("Successfully created the directory %s " % path) + + sess_options = onnxruntime.SessionOptions() + sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider']) + + input1_name = sess.get_inputs()[0].name + output_names = [output.name for output in sess.get_outputs()] + inputs = {input1_name: input_1} + + with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f: + f.write(tensor_1.SerializeToString()) + + for i in range(12): + input_name = f"past_{i}" + input = np.random.rand(2, batch_size, new_parameters["num_heads"], sequence_length, + new_parameters["size_per_head"]).astype(np.float32) + tensor = numpy_helper.from_array(input, input_name) + inputs.update({input_name: input}) + + with open(os.path.join(path, 'input_{}.pb'.format(1 + i)), 'wb') as f: + f.write(tensor.SerializeToString()) + + if input_tensor_only: + return + + result = sess.run(output_names, inputs) + print("result 0 shape:", result[0].shape) + print("result 1 shape:", result[1].shape) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--input', required=True, type=str) + parser.add_argument('--output', required=True, type=str) + parser.add_argument('--output_optimized_model', required=False, action='store_true') + parser.set_defaults(output_optimized_model=False) + args = parser.parse_args() + + model = ModelProto() + with open(args.input, "rb") as f: + model.ParseFromString(f.read()) + + bert_model = TinyGpt2Model(model) + + bert_model.update_graph() + bert_model.remove_unused_constant() + + print("opset verion", bert_model.model.opset_import[0].version) + + with open(args.output, "wb") as out: + out.write(bert_model.model.SerializeToString()) + + p = Path(args.output) + data_path = p.parent + + generate_test_data(args.output, + data_path, + batch_size=1, + use_cpu=True, + output_optimized_model=args.output_optimized_model) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/gpt2_past.onnx b/onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/gpt2_past.onnx new file mode 100644 index 0000000000..350ec55f89 Binary files /dev/null and b/onnxruntime/python/tools/transformers/test_data/gpt2_pytorch1.5_opset11/gpt2_past.onnx differ diff --git a/onnxruntime/python/tools/transformers/test_optimizer.py b/onnxruntime/python/tools/transformers/test_optimizer.py index 5fe97edfaa..0d1288380f 100644 --- a/onnxruntime/python/tools/transformers/test_optimizer.py +++ b/onnxruntime/python/tools/transformers/test_optimizer.py @@ -12,6 +12,7 @@ import unittest import os import onnx import onnxruntime +import pytest from onnx import helper, TensorProto, ModelProto from onnx.helper import make_node, make_tensor_value_info import numpy as np @@ -20,15 +21,23 @@ from optimizer import optimize_model, optimize_by_onnxruntime from OnnxModel import OnnxModel BERT_TEST_MODELS = { - "bert_pytorch_0": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_0.onnx', - "bert_pytorch_1": 'test_data\\bert_squad_pytorch1.4_opset11\\BertForQuestionAnswering_1.onnx', - "bert_squad_pytorch1.4_opset10_fp32": - 'test_data\\bert_squad_pytorch1.4_opset10_fp32\\BertForQuestionAnswering.onnx', - "bert_keras_0": 'test_data\\bert_mrpc_tensorflow2.1_opset10\\TFBertForSequenceClassification_1.onnx', - "bert_keras_squad": 'test_data\\bert_squad_tensorflow2.1_keras2onnx_opset11\\TFBertForQuestionAnswering.onnx', - "gpt2": 'test_data\\gpt2_pytorch1.4_opset11_no_past\\GPT2Model.onnx' + "bert_pytorch_0": ('bert_squad_pytorch1.4_opset11', 'BertForQuestionAnswering_0.onnx'), + "bert_pytorch_1": ('bert_squad_pytorch1.4_opset11', 'BertForQuestionAnswering_1.onnx'), + "bert_squad_pytorch1.4_opset10_fp32": ('bert_squad_pytorch1.4_opset10_fp32', 'BertForQuestionAnswering.onnx'), + "bert_keras_0": ('bert_mrpc_tensorflow2.1_opset10', 'TFBertForSequenceClassification_1.onnx'), + "bert_keras_squad": ('bert_squad_tensorflow2.1_keras2onnx_opset11', 'TFBertForQuestionAnswering.onnx'), + "gpt2": ('gpt2_pytorch1.4_opset11_no_past', 'GPT2Model.onnx'), + "gpt2_past": ('gpt2_pytorch1.5_opset11', 'gpt2_past.onnx'), } +skip_on_ort_version = pytest.mark.skipif(onnxruntime.__version__.startswith('1.3.'), + reason="skip failed tests. TODO: fix them in 1.4.0.") + + +def _get_test_model_path(name): + sub_dir, file = BERT_TEST_MODELS[name] + return os.path.join('test_data', sub_dir, file) + class TestBertOptimization(unittest.TestCase): def verify_node_count(self, bert_model, expected_node_count, test_name): @@ -39,8 +48,9 @@ class TestBertOptimization(unittest.TestCase): print("{}: {} expected={}".format(op, len(bert_model.get_nodes_by_op_type(op)), counter)) self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count) + @skip_on_ort_version def test_pytorch_model_0_cpu_onnxruntime(self): - input = BERT_TEST_MODELS['bert_pytorch_0'] + input = _get_test_model_path('bert_pytorch_0') output = 'temp.onnx' optimize_by_onnxruntime(input, use_gpu=False, optimized_model_path=output) model = ModelProto() @@ -58,12 +68,13 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_0_cpu_onnxruntime') + @skip_on_ort_version def test_pytorch_model_0_gpu_onnxruntime(self): if 'CUDAExecutionProvider' not in onnxruntime.get_available_providers(): print("skip test_pytorch_model_0_gpu_onnxruntime since no gpu found") return - input = BERT_TEST_MODELS['bert_pytorch_0'] + input = _get_test_model_path('bert_pytorch_0') output = 'temp.onnx' optimize_by_onnxruntime(input, use_gpu=True, optimized_model_path=output) model = ModelProto() @@ -81,8 +92,9 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_0_gpu_onnxruntime') + @skip_on_ort_version def test_pytorch_model_1_cpu_onnxruntime(self): - input = BERT_TEST_MODELS['bert_pytorch_1'] + input = _get_test_model_path('bert_pytorch_1') output = 'temp.onnx' optimize_by_onnxruntime(input, use_gpu=False, optimized_model_path=output) model = ModelProto() @@ -101,12 +113,13 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_1_cpu_onnxruntime') + @skip_on_ort_version def test_pytorch_model_1_gpu_onnxruntime(self): if 'CUDAExecutionProvider' not in onnxruntime.get_available_providers(): print("skip test_pytorch_model_1_gpu_onnxruntime since no gpu found") return - input = BERT_TEST_MODELS['bert_pytorch_1'] + input = _get_test_model_path('bert_pytorch_1') output = 'temp.onnx' optimize_by_onnxruntime(input, use_gpu=True, optimized_model_path=output) model = ModelProto() @@ -125,8 +138,9 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_1_gpu_onnxruntime') + @skip_on_ort_version def test_pytorch_model_0(self): - input = BERT_TEST_MODELS['bert_pytorch_0'] + input = _get_test_model_path('bert_pytorch_0') bert_model = optimize_model(input, 'bert', num_heads=2, hidden_size=8) expected_node_count = { @@ -140,13 +154,13 @@ class TestBertOptimization(unittest.TestCase): self.verify_node_count(bert_model, expected_node_count, 'test_pytorch_model_0') def test_pytorch_model_2(self): - input = BERT_TEST_MODELS['bert_squad_pytorch1.4_opset10_fp32'] + input = _get_test_model_path('bert_squad_pytorch1.4_opset10_fp32') bert_model = optimize_model(input, 'bert', num_heads=2, hidden_size=8) print("fused_operator_statistics for test_pytorch_model_2", bert_model.get_fused_operator_statistics()) self.assertTrue(bert_model.is_fully_optimized()) def test_keras_model_1(self): - input = BERT_TEST_MODELS['bert_keras_0'] + input = _get_test_model_path('bert_keras_0') bert_model = optimize_model(input, 'bert_keras', num_heads=2, hidden_size=8) @@ -162,7 +176,7 @@ class TestBertOptimization(unittest.TestCase): self.verify_node_count(bert_model, expected_node_count, 'test_keras_model_1') def test_keras_squad_model(self): - input = BERT_TEST_MODELS['bert_keras_squad'] + input = _get_test_model_path('bert_keras_squad') bert_model = optimize_model(input, 'bert_keras', num_heads=2, hidden_size=8) @@ -171,7 +185,7 @@ class TestBertOptimization(unittest.TestCase): self.assertTrue(bert_model.is_fully_optimized()) def test_gpt2(self): - input = BERT_TEST_MODELS['gpt2'] + input = _get_test_model_path('gpt2') bert_model = optimize_model(input, 'gpt2', num_heads=2, hidden_size=4) expected_node_count = { @@ -185,6 +199,21 @@ class TestBertOptimization(unittest.TestCase): } self.verify_node_count(bert_model, expected_node_count, 'test_gpt2') + def test_gpt2_past(self): + input = _get_test_model_path('gpt2_past') + bert_model = optimize_model(input, 'gpt2', num_heads=2, hidden_size=4) + + expected_node_count = { + 'EmbedLayerNormalization': 0, + 'Attention': 12, + 'Gelu': 0, + 'FastGelu': 12, + 'BiasGelu': 0, + 'LayerNormalization': 25, + 'SkipLayerNormalization': 0 + } + self.verify_node_count(bert_model, expected_node_count, 'test_gpt2_past') + if __name__ == '__main__': unittest.main() diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 87970a9341..6f7a15ab64 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -20,7 +20,12 @@ static void RunAttentionTest( int hidden_size, int number_of_heads, bool use_float16 = false, - bool is_unidirectional = false) { + bool is_unidirectional = false, + bool use_past_state = false, + int past_sequence_length = 0, + int head_size = 0, + const std::vector* past_data = nullptr, + const std::vector* present_data = nullptr) { int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); @@ -35,6 +40,8 @@ static void RunAttentionTest( std::vector weights_dims = {hidden_size, 3 * hidden_size}; std::vector bias_dims = {3 * hidden_size}; std::vector mask_index_dims = {batch_size}; + std::vector past_dims = {2, batch_size, head_size, past_sequence_length, head_size}; + std::vector present_dims = {2, batch_size, head_size, past_sequence_length + sequence_length, head_size}; std::vector output_dims = input_dims; if (use_float16) { @@ -51,6 +58,23 @@ static void RunAttentionTest( if (mask_index_data.size() > 0) { // mask index is optional. tester.AddInput("mask_index", mask_index_dims, mask_index_data); + } else { + std::vector dims = {static_cast(mask_index_data.size())}; + tester.AddInput("", dims, mask_index_data); + } + + if (use_past_state) { + if (use_float16) { + if (past_sequence_length > 0) { + tester.AddInput("past", past_dims, ToFloat16(*past_data)); + } + tester.AddOutput("present", present_dims, ToFloat16(*present_data)); + } else { + if (past_sequence_length > 0) { + tester.AddInput("past", past_dims, *past_data); + } + tester.AddOutput("present", present_dims, *present_data); + } } if (enable_cuda) { @@ -256,8 +280,7 @@ TEST(AttentionTest, AttentionUnidirectional) { std::vector input_data = { 0.091099896f, -0.018294459f, -0.36594841f, 0.28410032f, - -0.12125026f, -0.0066160089f, 0.38809127f, -0.22455512f - }; + -0.12125026f, -0.0066160089f, 0.38809127f, -0.22455512f}; std::vector weight_data = { -0.2659236192703247f, @@ -310,8 +333,7 @@ TEST(AttentionTest, AttentionUnidirectional) { -0.09368397295475006f, 0.07878211885690689f, 0.2973634898662567f, - 0.11210034042596817f -}; + 0.11210034042596817f}; std::vector bias_data = { -0.0540979839861393f, @@ -325,20 +347,317 @@ TEST(AttentionTest, AttentionUnidirectional) { 0.3670335114002228f, 0.028461361303925514f, -0.08913630992174149f, - 0.28048714995384216f - }; + 0.28048714995384216f}; // No mask_index std::vector mask_index_data = {}; std::vector output_data = { - 0.28109729f, 0.069518551f, 0.0038009658f, 0.29213354f, 0.3692801f, 0.029495837f, -0.084964074f, 0.28169215f - }; + 0.28109729f, 0.069518551f, 0.0038009658f, 0.29213354f, 0.3692801f, 0.029495837f, -0.084964074f, 0.28169215f}; bool is_unidirectional = true; RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional); } +TEST(AttentionTest, AttentionEmptyPastState) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.091099896f, -0.018294459f, -0.36594841f, 0.28410032f, + -0.12125026f, -0.0066160089f, 0.38809127f, -0.22455512f}; + + std::vector weight_data = { + -0.2659236192703247f, + 0.02789675071835518f, + 0.07280516624450684f, + 0.050951678305864334f, + 0.020417947322130203f, + -0.04751841351389885f, + 0.043815530836582184f, + 0.006015353370457888f, + -0.11496957391500473f, + -0.1773347705602646f, + 0.30928605794906616f, + 0.005648412741720676f, + + 0.08960387855768204f, + -0.27270448207855225f, + 0.14847396314144135f, + -0.17960812151432037f, + 0.01788954995572567f, + 0.09993876516819f, + 0.03943513706326485f, + -0.02484400011599064f, + -0.12958766520023346f, + 0.220433309674263f, + 0.1720484346151352f, + 0.22024005651474f, + + 0.059368450194597244f, + 0.1710093915462494f, + -0.3967452347278595f, + -0.1591450721025467f, + 0.1446179747581482f, + -0.20505407452583313f, + 0.12749597430229187f, + 0.32139700651168823f, + 0.139958456158638f, + -0.10619817674160004f, + 0.04528557509183884f, + 0.045598603785037994f, + + -0.007152545265853405f, + 0.109454445540905f, + -0.1582530289888382f, + -0.2646341919898987f, + 0.0920850858092308f, + 0.0701494812965393f, + -0.19062495231628418f, + -0.24360455572605133f, + -0.09368397295475006f, + 0.07878211885690689f, + 0.2973634898662567f, + 0.11210034042596817f}; + + std::vector bias_data = { + -0.0540979839861393f, + -0.06444740295410156f, + 0.03112877532839775f, + -0.08288222551345825f, + 0.07840359210968018f, + 0.039143580943346024f, + -0.45591455698013306f, + -0.11876055598258972f, + 0.3670335114002228f, + 0.028461361303925514f, + -0.08913630992174149f, + 0.28048714995384216f}; + + // No mask_index + std::vector mask_index_data = {}; + + std::vector output_data = { + 0.28109729f, 0.069518551f, 0.0038009658f, 0.29213354f, 0.3692801f, 0.029495837f, -0.084964074f, 0.28169215f}; + + std::vector past_data = {}; + + std::vector present_data = { + 0.053175069391727448f, 0.12795503437519073f, 0.11125634610652924f, -0.0510881207883358f, -0.55345797538757324f, -0.3045809268951416f, -0.36920222640037537f, 0.060108467936515808f, 0.28109729290008545f, 0.069518551230430603f, 0.45718482136726379f, -0.010400654748082161f, 0.0038009658455848694f, 0.29213353991508484f, -0.17697516083717346f, 0.27086889743804932f}; + + bool is_unidirectional = true; + bool use_past_state = true; + int past_sequence_length = 0; + int head_size = 2; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, + use_past_state, past_sequence_length, head_size, &past_data, &present_data); +} + +TEST(AttentionTest, AttentionPastStateBatch1) { + int batch_size = 1; + int sequence_length = 1; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + -0.019333266f, -0.21813886f, 0.16212955f, -0.015626367f}; + + std::vector weight_data = { + -0.4738484025001526f, + -0.2613658607006073f, + -0.0978037416934967f, + -0.34988933801651f, + 0.2243240624666214f, + -0.0429205559194088f, + 0.418695330619812f, + 0.17441125214099884f, + -0.18825532495975494f, + 0.18357256054878235f, + -0.5806483626365662f, + -0.02251487597823143f, + + 0.08742205798625946f, + 0.14734269678592682f, + 0.2387014478445053f, + 0.2884027063846588f, + 0.6490834355354309f, + 0.16965825855731964f, + -0.06346885114908218f, + 0.4073973298072815f, + -0.03070945478975773f, + 0.4110257923603058f, + 0.07896808534860611f, + 0.16783113777637482f, + + 0.0038893644232302904f, + 0.06946629285812378f, + 0.36680519580841064f, + -0.07261059433221817f, + -0.14960581064224243f, + 0.020944256335496902f, + -0.09378612786531448f, + -0.1336742341518402f, + 0.06061394885182381f, + 0.2205914407968521f, + -0.03519909828901291f, + -0.18405692279338837f, + + 0.22149960696697235f, + -0.1884360909461975f, + -0.014074507169425488f, + 0.4252440333366394f, + 0.24987126886844635f, + -0.31396418809890747f, + 0.14036843180656433f, + 0.2854192554950714f, + 0.09709841012954712f, + 0.09935075044631958f, + -0.012154420837759972f, + 0.2575816512107849f}; + + std::vector bias_data = { + 0.4803391396999359f, + -0.5254325866699219f, + -0.42926454544067383f, + -0.2059524953365326f, + -0.12773379683494568f, + -0.09542735666036606f, + -0.35286077857017517f, + -0.07646317780017853f, + -0.04590314254164696f, + -0.03752850368618965f, + -0.013764488510787487f, + -0.18478283286094666f}; + + // No mask_index + std::vector mask_index_data = {}; + + std::vector output_data = { + 0.20141591f, 0.43005896f, 0.35745093f, 0.19957167f}; + + std::vector past_data = { + 0.55445826f, 0.10127074f, 0.71770734f, 0.15915526f, 0.13913247f, 0.77447522f, 0.66044068f, 0.27559045f, 0.35731629f, 0.62033528f, 0.24354559f, 0.22859341f, + 0.45075402f, 0.85365993f, 0.097346395f, 0.28859729f, 0.26926181f, 0.65922296f, 0.8177433f, 0.4212271f, 0.34352475f, 0.059609573f, 0.46556228f, 0.7226882f}; + + std::vector present_data = { + 0.55445826f, 0.10127074f, 0.71770734f, 0.15915526f, 0.13913247f, 0.77447522f, -0.30182117f, -0.12330482f, 0.66044068f, 0.27559045f, 0.35731629f, 0.62033528f, 0.24354559f, 0.22859341f, -0.36450946f, -0.19483691f, + 0.45075402f, 0.85365993f, 0.097346395f, 0.28859729f, 0.26926181f, 0.65922296f, -0.027254611f, -0.096526355f, 0.8177433f, 0.4212271f, 0.34352475f, 0.059609573f, 0.46556228f, 0.7226882f, -0.025281552f, -0.25482416f}; + + bool is_unidirectional = true; + bool use_past_state = true; + int past_sequence_length = 3; + int head_size = 2; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, + use_past_state, past_sequence_length, head_size, &past_data, &present_data); +} + +TEST(AttentionTest, AttentionPastStateBatch2) { + int batch_size = 2; + int sequence_length = 1; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + -0.10902753f, 0.0041178204f, 0.1871525f, -0.20399982f, + 0.027207348f, -0.25321805f, 0.12869114f, 0.023136809f}; + + std::vector weight_data = { + -0.4738484025001526f, + -0.2613658607006073f, + -0.0978037416934967f, + -0.34988933801651f, + 0.2243240624666214f, + -0.0429205559194088f, + 0.418695330619812f, + 0.17441125214099884f, + -0.18825532495975494f, + 0.18357256054878235f, + -0.5806483626365662f, + -0.02251487597823143f, + + 0.08742205798625946f, + 0.14734269678592682f, + 0.2387014478445053f, + 0.2884027063846588f, + 0.6490834355354309f, + 0.16965825855731964f, + -0.06346885114908218f, + 0.4073973298072815f, + -0.03070945478975773f, + 0.4110257923603058f, + 0.07896808534860611f, + 0.16783113777637482f, + + 0.0038893644232302904f, + 0.06946629285812378f, + 0.36680519580841064f, + -0.07261059433221817f, + -0.14960581064224243f, + 0.020944256335496902f, + -0.09378612786531448f, + -0.1336742341518402f, + 0.06061394885182381f, + 0.2205914407968521f, + -0.03519909828901291f, + -0.18405692279338837f, + + 0.22149960696697235f, + -0.1884360909461975f, + -0.014074507169425488f, + 0.4252440333366394f, + 0.24987126886844635f, + -0.31396418809890747f, + 0.14036843180656433f, + 0.2854192554950714f, + 0.09709841012954712f, + 0.09935075044631958f, + -0.012154420837759972f, + 0.2575816512107849f}; + + std::vector bias_data = { + 0.4803391396999359f, + -0.5254325866699219f, + -0.42926454544067383f, + -0.2059524953365326f, + -0.12773379683494568f, + -0.09542735666036606f, + -0.35286077857017517f, + -0.07646317780017853f, + -0.04590314254164696f, + -0.03752850368618965f, + -0.013764488510787487f, + -0.18478283286094666f}; + + // No mask_index + std::vector mask_index_data = {}; + + std::vector output_data = { + 0.14902574f, 0.62273371f, 0.43022552f, 0.12759127f, + 0.26993567f, 0.23553593f, 0.43190649f, 0.086044826f}; + + std::vector past_data = { + 0.42028648f, 0.55855948f, 0.044569403f, 0.76525789f, 0.13962431f, 0.40977913f, 0.36911047f, 0.83399564f, 0.36905321f, 0.91414654f, 0.17300875f, 0.78793788f, + 0.10279467f, 0.80501258f, 0.089550517f, 0.85371113f, 0.61801594f, 0.91222942f, 0.88626182f, 0.069776468f, 0.10591964f, 0.84836882f, 0.83520192f, 0.0098680854f, + 0.3113814f, 0.63999802f, 0.28603253f, 0.98899829f, 0.044405211f, 0.95105386f, 0.81278932f, 0.63969064f, 0.14494057f, 0.11349615f, 0.87086016f, 0.20983537f, + 0.35107401f, 0.90144604f, 0.68950737f, 0.18928574f, 0.18029204f, 0.074517399f, 0.70763874f, 0.48440042f, 0.58114725f, 0.1048766f, 0.73694098f, 0.17766342f}; + + std::vector present_data = { + 0.42028648f, 0.55855948f, 0.044569403f, 0.76525789f, 0.13962431f, 0.40977913f, -0.22849128f, -0.022080801f, 0.36911047f, 0.83399564f, 0.36905321f, 0.91414654f, 0.17300875f, 0.78793788f, -0.4449589f, -0.17704415f, 0.10279467f, 0.80501258f, 0.089550517f, 0.85371113f, 0.61801594f, 0.91222942f, -0.2994619f, -0.14412443f, 0.88626182f, 0.069776468f, 0.10591964f, 0.84836882f, 0.83520192f, 0.0098680854f, -0.33421949f, -0.18547727f, + 0.3113814f, 0.63999802f, 0.28603253f, 0.98899829f, 0.044405211f, 0.95105386f, -0.033968594f, -0.034833729f, 0.81278932f, 0.63969064f, 0.14494057f, 0.11349615f, 0.87086016f, 0.20983537f, 0.045759238f, -0.26863033f, 0.35107401f, 0.90144604f, 0.68950737f, 0.18928574f, 0.18029204f, 0.074517399f, -0.033201858f, -0.10592631f, 0.70763874f, 0.48440042f, 0.58114725f, 0.1048766f, 0.73694098f, 0.17766342f, -0.054369561f, -0.24562015f}; + + bool is_unidirectional = true; + bool use_past_state = true; + int past_sequence_length = 3; + int head_size = 2; + RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, false, is_unidirectional, + use_past_state, past_sequence_length, head_size, &past_data, &present_data); +} + } // namespace test } // namespace onnxruntime