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

dump

Browse source
This commit is contained in:
Tianlei Wu 2025-02-05 01:18:02 -08:00
parent c89a798b73
commit 4bc025537a
5 changed files with 137 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -8,7 +8,7 @@
#include <gsl/gsl>
#include "core/framework/allocator.h"
#include "core/framework/ort_value.h"
#include "contrib_ops/cpu/utils/debug_macros.h"
#include "contrib_ops/cpu/utils/console_dumper.h"
namespace onnxruntime {

2
onnxruntime/contrib_ops/cpu/utils/debug_macros.h
View file

@ -1,7 +1,7 @@
#pragma once
#include "core/common/make_string.h"
// #define DEBUG_GENERATION 1 // uncomment it for debugging generation (like beam search etc)
#define DEBUG_GENERATION 1 // uncomment it for debugging generation (like beam search etc)
#ifdef DEBUG_GENERATION
#define DUMP_TENSOR_LEVEL 2

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

@ -5,7 +5,7 @@
// cub.cuh includes device/dispatch_radix_sort.cuh which has assignment in conditional expressions
#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4706)
#pragma warning(disable : 4706)
#endif
#include <cub/cub.cuh>
#if defined(_MSC_VER)
@ -267,6 +267,37 @@ __global__ void InitializeBeamHypotheses(BeamHypotheses* beam_hyps, int beam_hyp
beam_hyp.done_ = false;
}
// Function to dump the data in the struct
__device__ void dump_hypothesis_score(const struct HypothesisScore* hs) {
printf("HypothesisScore Dump:\n");
printf(" hypothesis_length: %d\n", hs->hypothesis_length);
printf(" score: %f\n", hs->score);
printf(" hypothesis: ");
if (hs->hypothesis_length > 0 && hs->hypothesis != NULL) {
for (int i = 0; i < hs->hypothesis_length; ++i) {
printf("%d ", hs->hypothesis[i]);
}
} else {
printf("(empty)");
}
printf("\n");
}
__device__ void dump_beam_hypotheses(const struct BeamHypotheses* bh) {
printf("BeamHypotheses Dump:\n");
printf(" beams_count_: %d\n", bh->beams_count_);
printf(" beams_used_: %d\n", bh->beams_used_);
printf(" length_penalty_: %f\n", bh->length_penalty_);
printf(" done_: %s\n", bh->done_ ? "true" : "false");
printf(" beams_: %d\n", bh->beams_used_);
for (int i = 0; i < bh->beams_used_; ++i) {
printf(" Beam %d:\n", i + 1);
dump_hypothesis_score(&bh->beams_[i]);
}
}
// For counts that are typically far less than 256, this will round up the count to the next multiple of 32
// If this winds up being >256 then it uses a block size of 256 and calculates the appropriate grid_size
struct GridBlock32 {
@ -298,6 +329,8 @@ void LaunchInitializeBeamHypotheses(gsl::span<BeamHypotheses> beam_hyps,
__device__ void BeamHypotheses::Add(const int32_t* hypothesis, int hypothesis_length, float sum_logprobs) {
float score = sum_logprobs / pow(static_cast<float>(hypothesis_length), length_penalty_);
printf("\n BeamHypotheses::Add (score=%f hypothesis_length=%d sum_logprobs=%f) \n", score, hypothesis_length, sum_logprobs);
size_t index = beams_used_;
// If the array is full, don't add unless it's better than the worst element
if (index == beams_count_) {
@ -311,11 +344,17 @@ __device__ void BeamHypotheses::Add(const int32_t* hypothesis, int hypothesis_le
beams_[index] = beams_[index - 1];
beams_[index] = HypothesisScore{hypothesis, hypothesis_length, score};
printf("\n BeamHypotheses::Add (index=%d) \n", static_cast<int>(index));
// dump_hypothesis_score(&beams_[index]);
}
__device__ bool BeamHypotheses::CanImprove(float best_sum_logprobs, int current_length) const {
float current_score = best_sum_logprobs / pow(static_cast<float>(current_length), length_penalty_);
return beams_[beams_count_ - 1].score < current_score;
bool result = beams_[beams_count_ - 1].score < current_score;
printf("\n BeamHypotheses::CanImprove (current_score=%f beams_[%d].score=%f can_improve=%d) \n",
current_score, beams_count_ - 1, beams_[beams_count_ - 1].score, static_cast<int>(result));
return result;
}
template <typename T>
@ -344,7 +383,7 @@ __device__ void BeamHypotheses::Output(
}
__global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
BeamScorerState& state,
BeamScorerState* state_gpu,
const int32_t* sequences_buffer,
int sequence_length,
BeamHypotheses* beam_hyps_,
@ -358,11 +397,15 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
// Sequences shape is (batch_size * num_beams, total_sequence_length)
// It contains word ID of whole sequence generated so far.
// It is different from subgraph input_ids, which only need one word when past state is not empty.
BeamScorerState& state = *state_gpu;
printf("\n >>> BeamSearchScorer_Process \n");
int batch = threadIdx.x;
int batch_start = batch * state.num_beams_;
cuda::BeamHypotheses& beam_hyp = beam_hyps_[batch];
if (!beam_hyp.done_) {
// Next tokens for this sentence.
size_t beam_idx = 0;
@ -373,6 +416,10 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
int32_t next_index = next_indices[batch * top_k + j];
int batch_beam_idx = batch_start + next_index;
printf("\nbatch=%d batch_beam_idx=%d j=%d next_token=%d eos_token_id=%d next_score=%f next_index=%d\n",
batch, batch_beam_idx, static_cast<int>(j), next_token, state.eos_token_id_, next_score, next_index);
// Add to generated hypotheses if end of sentence.
if ((state.eos_token_id_ >= 0) && (next_token == state.eos_token_id_)) {
bool is_beam_token_worse_than_top_num_beams = (j >= state.num_beams_);
@ -401,11 +448,15 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
}
// Check if we are done so that we can save a pad step if all(done)
printf("\n beam_hyp.beams_used_ == state.num_beams_ is %d\n", beam_hyp.beams_used_ == state.num_beams_ ? 1 : 0);
if (beam_hyp.beams_used_ == state.num_beams_) {
if (state.early_stopping_ || !beam_hyp.CanImprove(*std::max_element(next_scores + batch_start, next_scores + batch_start + top_k), sequence_length)) {
beam_hyp.done_ = true;
if (atomicAdd(&state.not_done_count_, -1) == 1)
state_cpu.not_done_count_ = 0; // Update the CPU side
printf("\n --- BeamSearchScorer_Process updated cpu state for batch %d\n", threadIdx.x);
}
}
} else {
@ -416,10 +467,38 @@ __global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
next_beam_indices_[batch_start + beam_idx] = 0;
}
}
printf("\n <<< BeamSearchScorer_Process \n");
}
__global__ void DumpBeamScorerState(BeamScorerState* state) {
state->Print(false);
}
void DumpBeamScorerStates(BeamScorerState& state_cpu, BeamScorerState* state, cudaStream_t stream){
state_cpu.Print(true);
cudaDeviceSynchronize();
DumpBeamScorerState<<<1, 1, 0, stream>>>(state);
cudaDeviceSynchronize();
}
__global__ void DumpBeamSearchScorer(const BeamHypotheses& beam_hyp) {
dump_beam_hypotheses(&beam_hyp);
}
void DumpBeamHypotheses(gsl::span<BeamHypotheses> beam_hyps, cudaStream_t stream) {
printf("\n BeamHypotheses of size %zu: \n", beam_hyps.size());
for (size_t i = 0; i < beam_hyps.size(); i++) {
printf("\n [%zu]:\n", i);
cudaDeviceSynchronize();
DumpBeamSearchScorer<<<1, 1, 0, stream>>>(beam_hyps[i]);
cudaDeviceSynchronize();
}
}
void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
BeamScorerState& state,
BeamScorerState* state_gpu,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps,
@ -431,8 +510,18 @@ void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
gsl::span<const int32_t> next_tokens,
gsl::span<const int32_t> next_indices,
cudaStream_t stream) {
// size_t printfBufferSize = 10 * 1024 * 1024;
// cudaDeviceSetLimit(cudaLimitPrintfFifoSize, printfBufferSize);
// cudaDeviceSynchronize();
printf("\n >>> LaunchBeamSearchScorer_Process \n");
// cudaDeviceSynchronize();
DumpBeamHypotheses(beam_hyps, stream);
cudaDeviceSynchronize();
DumpBeamScorerStates(state_cpu, state_gpu, stream);
cudaDeviceSynchronize();
BeamSearchScorer_Process<<<1, state_cpu.batch_size_, 0, stream>>>(state_cpu,
state,
state_gpu,
sequences.data(),
sequence_length,
beam_hyps.data(),
@ -443,6 +532,14 @@ void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
next_scores.data(),
next_tokens.data(),
next_indices.data());
cudaDeviceSynchronize();
DumpBeamHypotheses(beam_hyps, stream);
cudaDeviceSynchronize();
DumpBeamScorerStates(state_cpu, state_gpu, stream);
cudaDeviceSynchronize();
printf("\n <<< LaunchBeamSearchScorer_Process \n");
// cudaDeviceSynchronize();
}
__global__ void BeamSearchScorer_AppendNextTokenToSequences1(BeamScorerState& state,
@ -478,6 +575,8 @@ void LaunchBeamSearchScorer_AppendNextTokenToSequences(BeamScorerState& state_cp
gsl::span<int32_t> next_beam_tokens,
gsl::span<int32_t> next_beam_indices,
cudaStream_t stream) {
printf("\n >>> LaunchBeamSearchScorer_AppendNextTokenToSequences \n");
const int max_threads = 512;
int batch_beam_size = state_cpu.batch_size_ * state_cpu.num_beams_;
dim3 block_size;
@ -511,6 +610,7 @@ void LaunchBeamSearchScorer_AppendNextTokenToSequences(BeamScorerState& state_cp
next_sequences.data(),
sequence_length,
next_beam_tokens.data());
printf("\n <<< LaunchBeamSearchScorer_AppendNextTokenToSequences \n");
}
template <typename T>
@ -557,6 +657,8 @@ void LaunchBeamSearchScorer_Finalize(int batch_size,
gsl::span<int32_t> output,
gsl::span<T> sequence_scores,
cudaStream_t stream) {
printf("\n >>> LaunchBeamSearchScorer_Finalize \n");
cudaDeviceSynchronize();
BeamSearchScorer_Finalize<<<1, batch_size, 0, stream>>>(state,
sequences.data(),
sequence_length,
@ -564,6 +666,8 @@ void LaunchBeamSearchScorer_Finalize(int batch_size,
final_beam_scores.data(),
output.data(),
sequence_scores.data());
cudaDeviceSynchronize();
printf("\n <<< LaunchBeamSearchScorer_Finalize \n");
}
template void LaunchBeamSearchScorer_Finalize<float>(

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

@ -6,6 +6,7 @@
#include <stdint.h>
#include <cuda_fp16.h>
#include <curand_kernel.h>
#include <cstdio>
namespace onnxruntime {
namespace contrib {
@ -82,14 +83,27 @@ struct BeamScorerState {
int eos_token_id_;
bool early_stopping_;
int not_done_count_; // When zero, every batch entry is done (starts at batch_size_)
int hypothesis_buffer_used_; // Offset of available buffer, or length of used buffer.
// Function to dump the struct data to stdout
__host__ __device__ void Print(bool is_cpu) const {
printf("BeamScorerState (cpu=%d) Dump:\n", is_cpu ? 1 : 0);
printf(" batch_size_: %d\n", batch_size_);
printf(" num_beams_: %d\n", num_beams_);
printf(" max_length_: %d\n", max_length_);
printf(" num_return_sequences_: %d\n", num_return_sequences_);
printf(" pad_token_id_: %d\n", pad_token_id_);
printf(" eos_token_id_: %d\n", eos_token_id_);
printf(" early_stopping_: %s\n", early_stopping_ ? "true" : "false");
printf(" not_done_count_: %d\n", not_done_count_);
printf(" hypothesis_buffer_used_: %d\n", hypothesis_buffer_used_);
}
};
void LaunchInitializeBeamHypotheses(gsl::span<BeamHypotheses> beam_hyps, float length_penalty, gsl::span<HypothesisScore> beams, int num_beams, cudaStream_t stream);
void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
BeamScorerState& state,
BeamScorerState* state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps_,

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

@ -588,6 +588,7 @@ Status ProcessLogits(const OrtValue& logits, //
cuda::LaunchNextTokenKernel(next_token_indices, beam_state->next_indices.data(), beam_state->next_tokens.data(),
batch_size, top_k, vocab_size, cuda_stream);
// BUG?: Copy topk_scores to next_scores
#ifdef DEBUG_GENERATION
dumper->Print("next_scores before scorer", topk_scores->Data<float>(), batch_size, top_k);
dumper->Print("next_tokens before scorer", beam_state->next_tokens.data(), batch_size, top_k);
@ -722,8 +723,11 @@ void CudaBeamSearchScorer::Process(transformers::ISequences& sequences,
gsl::span<const float>& next_scores,
gsl::span<const int32_t>& next_tokens,
gsl::span<const int32_t>& next_indices) {
printf("\n---Process ---\n");
state_cpu_->Print(true);
cuda::LaunchBeamSearchScorer_Process(*state_cpu_,
*state_gpu_,
state_gpu_.get(),
sequences.GetCurrentDeviceSequences(),
sequences.GetSequenceLength(),
beam_hyps_,
@ -735,6 +739,7 @@ void CudaBeamSearchScorer::Process(transformers::ISequences& sequences,
next_tokens,
next_indices,
stream_);
CUDA_CALL_THROW(cudaEventRecord(event_process_complete_.Get(), stream_));
cuda::LaunchBeamSearchScorer_AppendNextTokenToSequences(*state_cpu_,
@ -749,6 +754,10 @@ void CudaBeamSearchScorer::Process(transformers::ISequences& sequences,
bool CudaBeamSearchScorer::IsDoneLater() const {
CUDA_CALL_THROW(cudaEventSynchronize(event_process_complete_.Get()));
printf("\n---IsDoneLater ---\n");
state_cpu_->Print(true);
return state_cpu_->not_done_count_ == 0;
}