From 4bc025537a1546fd3eac211a0a50dd0844b28142 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 5 Feb 2025 01:18:02 -0800 Subject: [PATCH] dump --- .../cpu/transformers/generation_shared.h | 2 +- .../contrib_ops/cpu/utils/debug_macros.h | 2 +- .../cuda/transformers/generation_cuda_impl.cu | 114 +++++++++++++++++- .../cuda/transformers/generation_cuda_impl.h | 18 ++- .../transformers/generation_device_helper.cc | 11 +- 5 files changed, 137 insertions(+), 10 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h index 8145fbd4a4..635c485296 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h @@ -8,7 +8,7 @@ #include #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 { diff --git a/onnxruntime/contrib_ops/cpu/utils/debug_macros.h b/onnxruntime/contrib_ops/cpu/utils/debug_macros.h index d5cbaa0a3e..d6cea821ce 100644 --- a/onnxruntime/contrib_ops/cpu/utils/debug_macros.h +++ b/onnxruntime/contrib_ops/cpu/utils/debug_macros.h @@ -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 diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu index eb1943b59d..891a898521 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu @@ -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 #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 beam_hyps, __device__ void BeamHypotheses::Add(const int32_t* hypothesis, int hypothesis_length, float sum_logprobs) { float score = sum_logprobs / pow(static_cast(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(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(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(result)); + return result; } template @@ -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(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 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 sequences, int sequence_length, gsl::span beam_hyps, @@ -431,8 +510,18 @@ void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu, gsl::span next_tokens, gsl::span 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 next_beam_tokens, gsl::span 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 @@ -557,6 +657,8 @@ void LaunchBeamSearchScorer_Finalize(int batch_size, gsl::span output, gsl::span 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( diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h index 281cb6c725..bf0f999168 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h @@ -6,6 +6,7 @@ #include #include #include +#include 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 beam_hyps, float length_penalty, gsl::span beams, int num_beams, cudaStream_t stream); void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu, - BeamScorerState& state, + BeamScorerState* state, gsl::span sequences, int sequence_length, gsl::span beam_hyps_, diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc index 4e65336665..7b3376e949 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc @@ -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(), 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& next_scores, gsl::span& next_tokens, gsl::span& 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; }