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

Zhalei/fix seqoutput type (#18765)

Browse source 
After refactoring beamsearch, all scores become fp32. Yet it need
support fp16 according to original specs.
This commit is contained in:
Zhang Lei 2024-01-22 10:40:48 -08:00 committed by GitHub
parent 21034a2c37
commit 373ebac167
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 220 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h
View file

@ -397,12 +397,8 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
output_sequences_scores);
// Output per token scores
if (output_scores) {
gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
gsl::span<const float> source = beam_state.scores;
assert(target.size() == source.size());
ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice));
}
gsl::span<const float> per_token_scores = beam_state.scores;
this->beam_scorer_->OutputScores(per_token_scores, output_scores);
return status;
}

8
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h
View file

@ -404,12 +404,8 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
output_sequences_scores);
// Output per token scores
if (output_scores) {
gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
gsl::span<const float> source = beam_state.scores;
assert(target.size() == source.size());
ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice));
}
gsl::span<const float> per_token_scores = beam_state.scores;
this->beam_scorer_->OutputScores(per_token_scores, output_scores);
return status;
}

8
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h
View file

@ -500,12 +500,8 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
output_sequences_scores);
// Output per token scores
if (output_scores) {
gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
gsl::span<const float> source = beam_state.scores;
assert(target.size() == source.size());
ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice));
}
gsl::span<const float> per_token_scores = beam_state.scores;
this->beam_scorer_->OutputScores(per_token_scores, output_scores);
return status;
}

82
onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc
View file

@ -50,11 +50,12 @@ bool BeamHypotheses::CanImprove(float best_sum_logprobs, int current_length) con
return beams_.back().score < current_score;
}
template <typename T>
void BeamHypotheses::Output(
int top_k,
int max_length,
gsl::span<int32_t>& sequences, // buffer filled with pad token ID, shape (num_return_sequences, max_length)
gsl::span<float>& sequences_scores) // buffer of shape (num_return_sequences) or empty
gsl::span<int32_t>& sequences, // buffer filled with pad token ID, shape (num_return_sequences, max_length)
gsl::span<T>& sequences_scores) // buffer of shape (num_return_sequences) or empty
{
// Copy the top_k beams into the sequences
ORT_ENFORCE(top_k <= beams_used_);
@ -67,7 +68,7 @@ void BeamHypotheses::Output(
gsl::copy(item.hypothesis, target);
if (!sequences_scores.empty())
sequences_scores[index] = item.score;
sequences_scores[index] = (T)item.score;
}
}
@ -181,21 +182,21 @@ void BeamSearchScorer::Process(ISequences& sequences,
}
}
void BeamSearchScorer::Finalize(ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) {
ORT_ENFORCE(output_sequences != nullptr);
template <typename T>
void OutputSequenceScores(BeamSearchScorer* scorer,
ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) {
// Finalize all open beam hypotheses and add to generated hypotheses.
for (size_t batch_index = 0; batch_index < batch_size_; batch_index++) {
BeamHypotheses& beam_hyp = beam_hyps_[batch_index];
for (size_t batch_index = 0; batch_index < scorer->batch_size_; batch_index++) {
BeamHypotheses& beam_hyp = scorer->beam_hyps_[batch_index];
if (beam_hyp.done_) {
continue;
}
for (size_t beam_index = 0; beam_index < num_beams_; beam_index++) {
size_t batch_beam_index = batch_index * num_beams_ + beam_index;
for (size_t beam_index = 0; beam_index < scorer->num_beams_; beam_index++) {
size_t batch_beam_index = batch_index * scorer->num_beams_ + beam_index;
float final_score = final_beam_scores[batch_beam_index];
auto final_tokens = sequences.GetSequence(narrow<int>(batch_beam_index));
beam_hyp.Add(final_tokens, final_score);
@ -206,26 +207,59 @@ void BeamSearchScorer::Finalize(ISequences& sequences,
gsl::span<int32_t> output = output_sequences->MutableDataAsSpan<int32_t>();
// Fill output sequences with pad token ID so that we do not need append it later.
std::fill_n(output.data(), output.size(), pad_token_id_);
std::fill_n(output.data(), output.size(), scorer->pad_token_id_);
// Score of each sequence, with shape (batch_size * num_return_sequences).
gsl::span<float> sequence_scores;
gsl::span<T> sequence_scores;
if (output_sequence_scores) {
sequence_scores = output_sequence_scores->MutableDataAsSpan<float>();
sequence_scores = output_sequence_scores->MutableDataAsSpan<T>();
}
// Select the best hypotheses according to number of sequences to return.
for (size_t batch_index = 0; batch_index < batch_size_; batch_index++) {
BeamHypotheses& beam_hyp = beam_hyps_[batch_index];
for (size_t batch_index = 0; batch_index < scorer->batch_size_; batch_index++) {
BeamHypotheses& beam_hyp = scorer->beam_hyps_[batch_index];
auto batch_output = output.subspan(batch_index * num_return_sequences_ * max_length_,
num_return_sequences_ * max_length_);
gsl::span<float> sequence_scores_buffer;
auto batch_output = output.subspan(batch_index * scorer->num_return_sequences_ * scorer->max_length_,
scorer->num_return_sequences_ * scorer->max_length_);
gsl::span<T> sequence_scores_buffer;
if (!sequence_scores.empty())
sequence_scores_buffer = sequence_scores.subspan(batch_index * num_return_sequences_, num_return_sequences_);
sequence_scores_buffer = sequence_scores.subspan(batch_index * scorer->num_return_sequences_, scorer->num_return_sequences_);
beam_hyp.Output(narrow<int>(num_return_sequences_), narrow<int>(max_length_), batch_output,
sequence_scores_buffer);
beam_hyp.template Output<T>(narrow<int>(scorer->num_return_sequences_), narrow<int>(scorer->max_length_), batch_output,
sequence_scores_buffer);
}
}
void BeamSearchScorer::Finalize(ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) {
ORT_ENFORCE(output_sequences != nullptr);
if (output_sequence_scores == nullptr || output_sequence_scores->IsDataType<float>()) {
OutputSequenceScores<float>(this, sequences, final_beam_scores, output_sequences, output_sequence_scores);
} else {
ORT_ENFORCE(output_sequence_scores->IsDataType<MLFloat16>());
OutputSequenceScores<MLFloat16>(this, sequences, final_beam_scores, output_sequences, output_sequence_scores);
}
}
void BeamSearchScorer::OutputScores(gsl::span<const float>& final_scores, Tensor* output_scores) {
if (output_scores) {
if (output_scores->IsDataType<float>()) {
gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
ORT_ENFORCE(target.size() == final_scores.size());
std::copy_n(final_scores.data(), final_scores.size(), target.data());
} else {
ORT_ENFORCE(output_scores->IsDataType<MLFloat16>());
gsl::span<MLFloat16> target = output_scores->MutableDataAsSpan<MLFloat16>();
ORT_ENFORCE(target.size() == final_scores.size());
const float* src = final_scores.data();
MLFloat16* dst = target.data();
for (size_t i = 0; i < target.size(); i++) {
dst[i] = MLFloat16(src[i]);
}
}
}
}

12
onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h
View file

@ -35,10 +35,11 @@ struct BeamHypotheses {
bool CanImprove(float best_sum_logprobs, int current_length) const;
// Output results
void Output(int top_k, // number of sequences to return
int max_length, // max sequence length
gsl::span<int32_t>& sequences, // buffer with pad token, shape (num_return_sequences, max_length)
gsl::span<float>& sequences_scores); // buffer for sequence scores, with shape (num_return_sequences)
template <typename T>
void Output(int top_k, // number of sequences to return
int max_length, // max sequence length
gsl::span<int32_t>& sequences, // buffer with pad token, shape (num_return_sequences, max_length)
gsl::span<T>& sequences_scores); // buffer for sequence scores, with shape (num_return_sequences)
gsl::span<HypothesisScore> beams_; // Beam width sized array of hypotheses, sorted by highest scoring
int beams_used_; // Number of elements used in beams_
@ -60,13 +61,14 @@ struct BeamSearchScorer : IBeamScorer {
Tensor* output_sequences,
Tensor* output_sequence_scores) override;
void OutputScores(gsl::span<const float>& final_scores, Tensor* output_scores) override;
bool IsDone() const override { return not_done_count_ == 0; }
gsl::span<float> GetNextScores() override { return next_beam_scores_; }
gsl::span<int32_t> GetNextTokens() override { return next_beam_tokens_; }
gsl::span<int32_t> GetNextIndicesCPU() override { return next_beam_indices_; }
private:
size_t batch_size_;
size_t num_beams_;
size_t max_length_;

3
onnxruntime/contrib_ops/cpu/transformers/generation_shared.h
View file

@ -120,6 +120,9 @@ struct IBeamScorer {
Tensor* output_sequences,
Tensor* output_sequence_scores) = 0;
virtual void OutputScores(gsl::span<const float>& final_scores,
Tensor* output_scores) = 0;
virtual bool IsDone() const = 0; // GPU version will return false here, as it asynchronously queues up the event
virtual bool IsDoneLater() const { return false; } // GPU version waits for the asynchous result to complete here

63
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu
View file

@ -307,12 +307,13 @@ __device__ bool BeamHypotheses::CanImprove(float best_sum_logprobs, int current_
return beams_[beams_count_ - 1].score < current_score;
}
template <typename T>
__device__ void BeamHypotheses::Output(
int top_k,
int max_length,
int pad_token_id,
int32_t* sequences, // buffer of shape (num_return_sequences, max_length)
float* sequences_scores) // buffer of shape (num_return_sequences) or empty
T* sequences_scores) // buffer of shape (num_return_sequences) or empty
{
// Copy the top_k beams into the sequences
for (int index = 0; index < top_k; index++) {
@ -327,7 +328,7 @@ __device__ void BeamHypotheses::Output(
target[i] = pad_token_id;
if (sequences_scores)
sequences_scores[index] = item.score;
sequences_scores[index] = (T)item.score;
}
}
@ -501,13 +502,14 @@ void LaunchBeamSearchScorer_AppendNextTokenToSequences(BeamScorerState& state_cp
next_beam_tokens.data());
}
template <typename T>
__global__ void BeamSearchScorer_Finalize(BeamScorerState& state,
const int32_t* sequences_buffer,
int sequence_length,
BeamHypotheses* beam_hyps_,
const float* final_beam_scores,
int32_t* output,
float* sequence_scores) {
T* sequence_scores) {
int batch_index = blockIdx.x * blockDim.x + threadIdx.x;
if (batch_index >= state.batch_size_)
return;
@ -534,6 +536,7 @@ __global__ void BeamSearchScorer_Finalize(BeamScorerState& state,
sequence_scores ? sequence_scores + batch_index * state.num_return_sequences_ : nullptr);
}
template <typename T>
void LaunchBeamSearchScorer_Finalize(int batch_size,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
@ -541,7 +544,7 @@ void LaunchBeamSearchScorer_Finalize(int batch_size,
gsl::span<BeamHypotheses> beam_hyps,
gsl::span<const float> final_beam_scores,
gsl::span<int32_t> output,
gsl::span<float> sequence_scores,
gsl::span<T> sequence_scores,
cudaStream_t stream) {
BeamSearchScorer_Finalize<<<1, batch_size, 0, stream>>>(state,
sequences.data(),
@ -552,6 +555,58 @@ void LaunchBeamSearchScorer_Finalize(int batch_size,
sequence_scores.data());
}
template void LaunchBeamSearchScorer_Finalize<float>(
int batch_size,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps,
gsl::span<const float> final_beam_scores,
gsl::span<int32_t> output,
gsl::span<float> sequence_scores,
cudaStream_t stream);
template void LaunchBeamSearchScorer_Finalize<__half>(
int batch_size,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps,
gsl::span<const float> final_beam_scores,
gsl::span<int32_t> output,
gsl::span<__half> sequence_scores,
cudaStream_t stream);
template <typename T>
__global__ void FloatConvertAndCopyKernel(const float* src, T* dst, size_t total_elements) {
int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (index < total_elements) {
dst[index] = (T)src[index];
}
}
template <typename T>
void LaunchBeamSearchScoreCopy(gsl::span<const float> final_scores,
gsl::span<T> output_scores,
cudaStream_t stream) {
ORT_ENFORCE(final_scores.size() == output_scores.size());
constexpr unsigned ThreadPerBlock = 256;
unsigned num_blocks = (unsigned)((final_scores.size() + (ThreadPerBlock - 1))/ ThreadPerBlock);
typedef typename ToCudaType<float>::MappedType CudaT;
FloatConvertAndCopyKernel<<<num_blocks, ThreadPerBlock, 0, stream>>>(
final_scores.data(), (CudaT*)output_scores.data(), final_scores.size());
}
template void LaunchBeamSearchScoreCopy(gsl::span<const float> final_scores,
gsl::span<float> output_scores,
cudaStream_t stream);
template void LaunchBeamSearchScoreCopy(gsl::span<const float> final_scores,
gsl::span<MLFloat16> output_scores,
cudaStream_t stream);
__global__ void AddProbsKernel(float* log_probs,
float* cum_log_probs,
const int vocab_size,

19
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h
View file

@ -65,11 +65,12 @@ struct BeamHypotheses {
__device__ bool CanImprove(float best_sum_logprobs, int current_length) const;
// Output results
__device__ void Output(int top_k, // number of sequences to return
int max_length, // max sequence length
int pad_token_id, // pad token
int32_t* sequences, // buffer with pad token, shape (num_return_sequences, max_length)
float* sequences_scores); // buffer for sequence scores, with shape (num_return_sequences)
template <typename T>
__device__ void Output(int top_k, // number of sequences to return
int max_length, // max sequence length
int pad_token_id, // pad token
int32_t* sequences, // buffer with pad token, shape (num_return_sequences, max_length)
T* sequences_scores); // buffer for sequence scores, with shape (num_return_sequences)
};
struct BeamScorerState {
@ -110,6 +111,7 @@ void LaunchBeamSearchScorer_AppendNextTokenToSequences(BeamScorerState& state_cp
gsl::span<int32_t> next_beam_indices,
cudaStream_t stream);
template <typename T>
void LaunchBeamSearchScorer_Finalize(int batch_size,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
@ -117,9 +119,14 @@ void LaunchBeamSearchScorer_Finalize(int batch_size,
gsl::span<BeamHypotheses> beam_hyps_,
gsl::span<const float> final_beam_scores,
gsl::span<int32_t> output,
gsl::span<float> sequence_scores,
gsl::span<T> sequence_scores,
cudaStream_t stream);
template <typename T>
void LaunchBeamSearchScoreCopy(gsl::span<const float> final_scores,
gsl::span<T> output_scores,
cudaStream_t stream);
void LaunchNextTokenKernel(const int64_t* next_token_indices,
int32_t* next_indices,
int32_t* next_tokens,

55
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
View file

@ -620,6 +620,8 @@ struct CudaBeamSearchScorer : transformers::IBeamScorer {
Tensor* output_sequences,
Tensor* output_sequence_scores) override;
void OutputScores(gsl::span<const float>& final_scores, Tensor* output_scores) override;
bool IsDone() const override { return false; } // For CUDA we speculatively run the next step while we wait for the GPU to report status. We use 'IsDoneLater()' for this
bool IsDoneLater() const override;
@ -632,7 +634,6 @@ struct CudaBeamSearchScorer : transformers::IBeamScorer {
}
gsl::span<int32_t> GetNextIndicesGPU() override { return next_beam_indices_; }
private:
mutable cuda::AutoDestoryCudaEvent event_process_complete_;
IAllocatorUniquePtr<cuda::BeamScorerState> state_cpu_;
IAllocatorUniquePtr<cuda::BeamScorerState> state_gpu_;
@ -743,22 +744,58 @@ bool CudaBeamSearchScorer::IsDoneLater() const {
return state_cpu_->not_done_count_ == 0;
}
template <typename T>
void CudaOutputSequenceScores(CudaBeamSearchScorer* scorer,
transformers::ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) {
// Word IDs of each sequence, with shape (batch_size * num_return_sequences, max_sequence_length).
gsl::span<int32_t> output{output_sequences->MutableData<int32_t>(), static_cast<size_t>(output_sequences->Shape().Size())};
// Score of each sequence, with shape (batch_size * num_return_sequences).
using CudaT = typename ToCudaType<T>::MappedType;
gsl::span<CudaT> sequence_scores;
if (output_sequence_scores) {
sequence_scores = gsl::span<CudaT>{(CudaT*)output_sequence_scores->MutableData<T>(), static_cast<size_t>(output_sequence_scores->Shape().Size())};
}
cuda::LaunchBeamSearchScorer_Finalize(scorer->state_cpu_->batch_size_,
*scorer->state_gpu_,
sequences.GetCurrentDeviceSequences(),
sequences.GetSequenceLength(),
scorer->beam_hyps_,
final_beam_scores,
output,
sequence_scores,
scorer->stream_);
}
void CudaBeamSearchScorer::Finalize(transformers::ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) {
ORT_ENFORCE(output_sequences != nullptr);
// Word IDs of each sequence, with shape (batch_size * num_return_sequences, max_sequence_length).
gsl::span<int32_t> output{output_sequences->MutableData<int32_t>(), static_cast<size_t>(output_sequences->Shape().Size())};
// Score of each sequence, with shape (batch_size * num_return_sequences).
gsl::span<float> sequence_scores;
if (output_sequence_scores) {
sequence_scores = gsl::span<float>{output_sequence_scores->MutableData<float>(), static_cast<size_t>(output_sequence_scores->Shape().Size())};
if (output_sequence_scores == nullptr || output_sequence_scores->IsDataType<float>()) {
CudaOutputSequenceScores<float>(this, sequences, final_beam_scores, output_sequences, output_sequence_scores);
} else {
ORT_ENFORCE(output_sequence_scores->IsDataType<MLFloat16>());
CudaOutputSequenceScores<MLFloat16>(this, sequences, final_beam_scores, output_sequences, output_sequence_scores);
}
}
cuda::LaunchBeamSearchScorer_Finalize(state_cpu_->batch_size_, *state_gpu_, sequences.GetCurrentDeviceSequences(), sequences.GetSequenceLength(), beam_hyps_, final_beam_scores, output, sequence_scores, stream_);
void CudaBeamSearchScorer::OutputScores(gsl::span<const float>& final_scores, Tensor* output_scores) {
if (output_scores) {
if (output_scores->IsDataType<float>()) {
gsl::span<float> target(output_scores->MutableData<float>(), output_scores->Shape().Size());
cuda::LaunchBeamSearchScoreCopy(final_scores, target, stream_);
} else {
ORT_ENFORCE(output_scores->IsDataType<MLFloat16>());
gsl::span<MLFloat16> target(output_scores->MutableData<MLFloat16>(), output_scores->Shape().Size());
cuda::LaunchBeamSearchScoreCopy(final_scores, target, stream_);
}
}
}
std::unique_ptr<transformers::IBeamScorer> CreateBeamScorer(const transformers::IGenerationParameters& parameters,

31
onnxruntime/python/tools/transformers/models/whisper/whisper_chain.py
View file

@ -53,9 +53,9 @@ def chain_model(args):
beam_outputs = ["sequences"]
if args.output_sequence_scores:
beam_outputs.append("sequence_scores")
beam_outputs.append("sequence_scores_fp16" if args.precision == Precision.FLOAT16 else "sequence_scores")
if args.output_scores:
beam_outputs.append("scores")
beam_outputs.append("scores_fp16" if args.precision == Precision.FLOAT16 else "scores")
if args.use_whisper_beamsearch:
assert len(beam_inputs) == 12
@ -75,6 +75,7 @@ def chain_model(args):
beam_outputs.extend(["no_speech_probs_beam"])
input_features_cast_node, len_pen_cast_node, rep_pen_cast_node = None, None, None
output_scores_cast_node = output_sequence_scores_cast_node = None
if args.precision == Precision.FLOAT16:
input_features_cast_node = helper.make_node(
"Cast",
@ -97,6 +98,22 @@ def chain_model(args):
name="CastRepetitionPenaltyToFp16",
to=TensorProto.FLOAT16,
)
if args.output_sequence_scores:
output_sequence_scores_cast_node = helper.make_node(
"Cast",
inputs=["sequence_scores_fp16"],
outputs=["sequence_scores"],
name="CastOutputSequenceScoresToFp32",
to=TensorProto.FLOAT,
)
if args.output_scores:
output_scores_cast_node = helper.make_node(
"Cast",
inputs=["scores_fp16"],
outputs=["scores"],
name="CastScoresToFp32",
to=TensorProto.FLOAT,
)
operator_type = "WhisperBeamSearch" if args.use_whisper_beamsearch else "BeamSearch"
node = helper.make_node(operator_type, inputs=beam_inputs, outputs=beam_outputs, name="BeamSearch_zcode")
@ -214,10 +231,18 @@ def chain_model(args):
opset_import = [helper.make_opsetid(domain="com.microsoft", version=1), helper.make_opsetid(domain="", version=17)]
graph_nodes = (
[input_features_cast_node, len_pen_cast_node, rep_pen_cast_node, node]
[
input_features_cast_node,
len_pen_cast_node,
rep_pen_cast_node,
node,
output_sequence_scores_cast_node,
output_scores_cast_node,
]
if args.precision == Precision.FLOAT16
else [node]
)
graph_nodes = [node for node in graph_nodes if node is not None]
if args.output_no_speech_probs:
prob_cast_node = helper.make_node(
"Cast",