From 0636bc4ee31a86c84f102aa0cf6b3cbebb4b411a Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 21 Apr 2024 16:45:46 +0300 Subject: [PATCH] removed the union and moved the logit->probability conversion into the sampling function --- test_gpt2.cu | 2 +- train_gpt2.cu | 42 +++++++++++++++++++++--------------------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/test_gpt2.cu b/test_gpt2.cu index 838646a..93c2e1b 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -95,7 +95,7 @@ int main(int argc, char *argv[]) { // at this point, target should be equal to expected_logits, let's compare // copy logits to CPU so we can compare them float* logits_cpu = (float*)mallocCheck(B * T * V * sizeof(float)); - cudaMemcpy(logits_cpu, model.acts.output.logits, B * T * V * sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(logits_cpu, model.acts.output, B * T * V * sizeof(float), cudaMemcpyDeviceToHost); int logits_ok = 1; for (int i=0; iencoded, &acts->ln1, &acts->ln1_mean, &acts->ln1_rstd, &acts->qkv, &acts->atty, &acts->preatt, &acts->att, &acts->attproj, &acts->residual2, &acts->ln2, &acts->ln2_mean, &acts->ln2_rstd, &acts->fch, &acts->fch_gelu, &acts->fcproj, &acts->residual3, &acts->lnf, - &acts->lnf_mean, &acts->lnf_rstd, &acts->output.logits, &acts->losses, + &acts->lnf_mean, &acts->lnf_rstd, &acts->output, &acts->losses, &acts->qkvr, &acts->v_accum }; float* acts_memory_iterator = acts_memory; @@ -1425,13 +1420,13 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, int B, int T) { residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 layernorm_forward(acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C); - matmul_forward_cublas(acts.output.logits, acts.lnf, params.wte, NULL, B, T, C, V); + matmul_forward_cublas(acts.output, acts.lnf, params.wte, NULL, B, T, C, V); // also forward the cross-entropy loss function if we have the targets if (targets != NULL) { // fused classifier: does the forward pass and first part of the backward pass // we're passing dlosses = NULL, which will default them to 1.0f/(B*T), i.e. uniform loss - fused_classifier3(acts.output.dlogits, acts.losses, NULL, model->targets, B, T, V, V); + fused_classifier3(acts.output, acts.losses, NULL, model->targets, B, T, V, V); // for convenience also evaluate the mean loss (TODO re-think this compute+sync point) // move the (B,T) losses to CPU cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(float), cudaMemcpyDeviceToHost)); @@ -1442,7 +1437,6 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, int B, int T) { } else { // if we don't have targets, we don't have loss - softmax_forward(acts.output.probs, acts.output.logits, B*T, V); model->mean_loss = -1.0f; } } @@ -1512,7 +1506,7 @@ void gpt2_backward(GPT2 *model) { // technically that is a small, inline backward() pass of calculating // total, final loss as the mean over all losses over all (B,T) positions in the batch // next: backward the classifier matmul - matmul_backward(grads_acts.lnf, grads.wte, NULL, acts.output.dlogits, acts.lnf, params.wte, B, T, C, V); + matmul_backward(grads_acts.lnf, grads.wte, NULL, acts.output, acts.lnf, params.wte, B, T, C, V); // backward the final layernorm float* residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 float* dresidual = grads_acts.residual3; // the main buffer holding the gradient in the backward pass @@ -1706,12 +1700,18 @@ float random_f32(unsigned long long *state) { // random float32 in [0,1) return (random_u32(state) >> 8) / 16777216.0f; } -int sample_mult(float* probabilities, int n, float coin) { - // sample index from probabilities (they must sum to 1!) +int sample_softmax(const float* logits, int n, float coin) { + // sample index from logits (converted to probabilities using softmax) // coin is a random number in [0, 1), usually from random_f32() + double norm = 0; + for (int i = 0; i < n; i++) { + norm += exp(logits[i]); + } + // instead of dividing all exp(logits), we can just multiply coin. + coin *= norm; float cdf = 0.0f; for (int i = 0; i < n; i++) { - cdf += probabilities[i]; + cdf += exp(logits[i]); if (coin < cdf) { return i; } @@ -1942,7 +1942,7 @@ int main(int argc, char *argv[]) { // some memory for generating samples from the model unsigned long long rng_state = 1337; int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); - float* cpu_probs = (float*)mallocCheck(model.config.vocab_size * sizeof(float)); + float* cpu_logits = (float*)mallocCheck(model.config.vocab_size * sizeof(float)); // train struct timespec start, end; @@ -1982,11 +1982,11 @@ int main(int argc, char *argv[]) { // we're in principle running B "inference streams" in parallel here // only using position 0 because it's a bit faster (copy less probs from GPU -> CPU) // get the V-dimensional vector probs[0, t-1, :] - float* probs = model.acts.output.probs + (t-1) * model.config.vocab_size; + float* logits = model.acts.output + (t-1) * model.config.vocab_size; // move probs back to CPU and sample - cudaCheck(cudaMemcpy(cpu_probs, probs, model.config.vocab_size * sizeof(float), cudaMemcpyDeviceToHost)); + cudaCheck(cudaMemcpy(cpu_logits, logits, model.config.vocab_size * sizeof(float), cudaMemcpyDeviceToHost)); float coin = random_f32(&rng_state); - int next_token = sample_mult(cpu_probs, model.config.vocab_size, coin); + int next_token = sample_softmax(cpu_logits, model.config.vocab_size, coin); gen_tokens[t] = next_token; // print the generated token, either using the Tokenizer or a fallback if (tokenizer.init_ok) { @@ -2029,7 +2029,7 @@ int main(int argc, char *argv[]) { dataloader_free(&val_loader); tokenizer_free(&tokenizer); gpt2_free(&model); - free(cpu_probs); + free(cpu_logits); free(gen_tokens); cudaCheck(cudaFree(cublaslt_workspace)); cublasCheck(cublasDestroy(cublas_handle));