420 lines
16 KiB
Plaintext
420 lines
16 KiB
Plaintext
//-----------------------------------------------------------------------------------------------
|
|
// Copyright (c) 2024 Andrej Karpathy
|
|
// Licensed under the MIT License. See the LICENSE file for details.
|
|
//
|
|
// Modifications Copyright (c) 2025 Hanwen Liu, Hrishi Shah, Kelin Zeng, Charles Pei, and Vijay Daita, ALL RIGHTS RESERVED.
|
|
//-----------------------------------------------------------------------------------------------
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <math.h>
|
|
#include <time.h>
|
|
#include <assert.h>
|
|
#include <float.h>
|
|
#include <string.h>
|
|
#include <unistd.h>
|
|
|
|
// GPU / CUDA related
|
|
#include <cublas_v2.h>
|
|
#include <cuda_runtime.h>
|
|
#include "utils/utils.h"
|
|
#include "utils/tokenizer.h"
|
|
#include "utils/cuda_utils.cuh"
|
|
|
|
// kernels
|
|
#include "kernels/attention.cuh"
|
|
#include "kernels/encoder.cuh"
|
|
#include "kernels/gelu.cuh"
|
|
#include "kernels/layernorm.cuh"
|
|
#include "kernels/matmul.cuh"
|
|
#include "kernels/residual.cuh"
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// GPT-2 model definition
|
|
|
|
typedef struct {
|
|
int max_seq_len;
|
|
int vocab_size;
|
|
int padded_vocab_size; // padded to multiple of 128
|
|
int num_layers;
|
|
int num_heads;
|
|
int channels;
|
|
} GPT2Config;
|
|
|
|
// the parameters of the model
|
|
#define NUM_PARAMETER_TENSORS 16
|
|
typedef struct {
|
|
float* wte; // (Vp, C)
|
|
float* wpe; // (maxT, C)
|
|
float* ln1w; // (L, C)
|
|
float* ln1b; // (L, C)
|
|
float* qkvw; // (L, 3*C, C)
|
|
float* qkvb; // (L, 3*C)
|
|
float* attprojw; // (L, C, C)
|
|
float* attprojb; // (L, C)
|
|
float* ln2w; // (L, C)
|
|
float* ln2b; // (L, C)
|
|
float* fcw; // (L, 4*C, C)
|
|
float* fcb; // (L, 4*C)
|
|
float* fcprojw; // (L, C, 4*C)
|
|
float* fcprojb; // (L, C)
|
|
float* lnfw; // (C)
|
|
float* lnfb; // (C)
|
|
} ParameterTensors;
|
|
|
|
void fill_in_parameter_sizes(size_t* param_sizes, GPT2Config config) {
|
|
int Vp = config.padded_vocab_size;
|
|
int C = config.channels;
|
|
int maxT = config.max_seq_len;
|
|
int L = config.num_layers;
|
|
param_sizes[0] = Vp * C; // wte
|
|
param_sizes[1] = maxT * C; // wpe
|
|
param_sizes[2] = L * C; // ln1w
|
|
param_sizes[3] = L * C; // ln1b
|
|
param_sizes[4] = L * (3 * C) * C; // qkvw
|
|
param_sizes[5] = L * (3 * C); // qkvb
|
|
param_sizes[6] = L * C * C; // attprojw
|
|
param_sizes[7] = L * C; // attprojb
|
|
param_sizes[8] = L * C; // ln2w
|
|
param_sizes[9] = L * C; // ln2b
|
|
param_sizes[10] = L * (4 * C) * C; // fcw
|
|
param_sizes[11] = L * (4 * C); // fcb
|
|
param_sizes[12] = L * C * (4 * C); // fcprojw
|
|
param_sizes[13] = L * C; // fcprojb
|
|
param_sizes[14] = C; // lnfw
|
|
param_sizes[15] = C; // lnfb
|
|
}
|
|
|
|
// allocate memory for the parameters and point the individual tensors to the right places
|
|
float* malloc_and_point_parameters(ParameterTensors* params, size_t* param_sizes, int on_device) {
|
|
// on_device: 0 = CPU, 1 = GPU
|
|
// calculate the number of parameters
|
|
size_t num_parameters = 0;
|
|
for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) {
|
|
num_parameters += param_sizes[i];
|
|
}
|
|
// malloc all parameters all at once on the device
|
|
float* params_memory;
|
|
if (on_device) {
|
|
cudaCheck(cudaMalloc((void**)¶ms_memory, num_parameters * sizeof(float)));
|
|
} else {
|
|
params_memory = (float*)mallocCheck(num_parameters * sizeof(float));
|
|
}
|
|
// assign all the tensors their place in the array
|
|
float** ptrs[] = {
|
|
¶ms->wte, ¶ms->wpe, ¶ms->ln1w, ¶ms->ln1b, ¶ms->qkvw, ¶ms->qkvb,
|
|
¶ms->attprojw, ¶ms->attprojb, ¶ms->ln2w, ¶ms->ln2b, ¶ms->fcw, ¶ms->fcb,
|
|
¶ms->fcprojw, ¶ms->fcprojb, ¶ms->lnfw, ¶ms->lnfb
|
|
};
|
|
float* params_memory_iterator = params_memory;
|
|
for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) {
|
|
*(ptrs[i]) = params_memory_iterator;
|
|
params_memory_iterator += param_sizes[i];
|
|
}
|
|
return params_memory;
|
|
}
|
|
|
|
#define NUM_ACTIVATION_TENSORS 21
|
|
typedef struct {
|
|
float* encoded; // (B, T, C)
|
|
float* ln1; // (L, B, T, C)
|
|
float* ln1_mean; // (L, B, T)
|
|
float* ln1_rstd; // (L, B, T)
|
|
float* atty; // (L, B, T, C)
|
|
float* att; // (L, B, NH, T, T)
|
|
float* attproj; // (L, B, T, C)
|
|
float* residual2; // (L, B, T, C)
|
|
float* ln2; // (L, B, T, C)
|
|
float* ln2_mean; // (L, B, T)
|
|
float* ln2_rstd; // (L, B, T)
|
|
float* fch; // (L, B, T, 4*C)
|
|
float* fch_gelu; // (L, B, T, 4*C)
|
|
float* fcproj; // (L, B, T, C)
|
|
float* residual3; // (L, B, T, C)
|
|
float* lnf; // (B, T, C)
|
|
float* lnf_mean; // (B, T)
|
|
float* lnf_rstd; // (B, T)
|
|
float* losses; // (B, T)
|
|
float* qkvr; // (L, B, T, 3*C)
|
|
// This buffer will store the logits
|
|
// during the processing of transformer blocks, we will also use this as a
|
|
// general scratchpad buffer. Allocation is made large enough to hold (B, T, 3C),
|
|
// (B, NH, T, T), and (B, T, V) shaped tensors.
|
|
float* output;
|
|
} ActivationTensors;
|
|
|
|
void fill_in_activation_sizes(size_t* act_sizes, int B, int T, GPT2Config config) {
|
|
size_t Vp = config.padded_vocab_size;
|
|
size_t L = config.num_layers;
|
|
size_t NH = config.num_heads;
|
|
size_t C = config.channels;
|
|
act_sizes[0] = B * T * C; // encoded
|
|
act_sizes[1] = L * B * T * C; // ln1
|
|
//to align making 1 to 4
|
|
act_sizes[2] = 4; // ln1_mean, reserved for future use
|
|
act_sizes[3] = 4; // ln1_rstd, reserved for future use
|
|
act_sizes[4] = L * B * T * C; // atty
|
|
act_sizes[5] = L * B * NH * T * T; // att
|
|
act_sizes[6] = L * B * T * C; // attproj
|
|
act_sizes[7] = L * B * T * C; // residual2
|
|
act_sizes[8] = L * B * T * C; // ln2
|
|
act_sizes[9] = 4; // ln2_mean, reserved for future use
|
|
act_sizes[10] = 4; // ln2_rstd, reserved for future use
|
|
act_sizes[11] = L * B * T * 4*C; // fch
|
|
act_sizes[12] = L * B * T * 4*C; // fch_gelu
|
|
act_sizes[13] = L * B * T * C; // fcproj
|
|
act_sizes[14] = L * B * T * C; // residual3
|
|
act_sizes[15] = B * T * C; // lnf
|
|
act_sizes[16] = 4; // lnf_mean, reserved for future use
|
|
act_sizes[17] = 4; // lnf_rstd, reserved for future use
|
|
act_sizes[18] = B * T; // losses
|
|
act_sizes[19] = L * B * T * 3*C; // qkvr
|
|
act_sizes[20] = B * T * max(3*C, max(NH*T, Vp)); // output / scratch
|
|
}
|
|
|
|
// Reserved for future use
|
|
#define NUM_BACKWARD_TENSORS 3
|
|
typedef struct {
|
|
float* bt4c; // (B, T, 4*C)
|
|
float* preatt; // (B, NH, T, T)
|
|
float* residual3; // (B, T, C)
|
|
} GradActTensors;
|
|
|
|
// op_13: 256-byte aligned activations.
|
|
// Each activation tensor's stride is rounded up to a multiple of 64 floats (256 bytes).
|
|
#define align_to_256_bytes(x) \
|
|
(x + (256/sizeof(float) - (x % (256/sizeof(float)))) % (256/sizeof(float)))
|
|
|
|
float* malloc_and_point(float** targets[], const size_t* act_sizes, int n) {
|
|
size_t num_activations = 0;
|
|
for (size_t i = 0; i < n; i++) {
|
|
num_activations += align_to_256_bytes(act_sizes[i]);
|
|
}
|
|
float* acts_memory;
|
|
cudaCheck(cudaMalloc((void**)&acts_memory, num_activations * sizeof(float)));
|
|
cudaCheck(cudaMemset(acts_memory, 0, num_activations * sizeof(float)));
|
|
float* acts_memory_iterator = acts_memory;
|
|
for (size_t i = 0; i < n; i++) {
|
|
*(targets[i]) = acts_memory_iterator;
|
|
acts_memory_iterator += align_to_256_bytes(act_sizes[i]);
|
|
}
|
|
return acts_memory;
|
|
}
|
|
|
|
float* malloc_and_point_activations(ActivationTensors* acts, const size_t* act_sizes) {
|
|
float** ptrs[] = {
|
|
&acts->encoded, &acts->ln1, &acts->ln1_mean, &acts->ln1_rstd, &acts->atty,
|
|
&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->losses, &acts->qkvr, &acts->output
|
|
};
|
|
return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS);
|
|
}
|
|
|
|
typedef struct {
|
|
GPT2Config config;
|
|
// the weights of the model, and their sizes
|
|
ParameterTensors params;
|
|
size_t param_sizes[NUM_PARAMETER_TENSORS];
|
|
float* params_memory;
|
|
size_t num_parameters;
|
|
ParameterTensors grads;
|
|
float* grads_memory;
|
|
float* m_memory;
|
|
float* v_memory;
|
|
// the activations of the model, and their sizes
|
|
ActivationTensors acts;
|
|
size_t act_sizes[NUM_ACTIVATION_TENSORS];
|
|
float* acts_memory;
|
|
size_t num_activations;
|
|
GradActTensors grads_acts;
|
|
size_t num_grad_acts;
|
|
float* grads_acts_memory;
|
|
// other run state configuration
|
|
int batch_size; // the batch size (B) of current forward pass
|
|
int seq_len; // the sequence length (T) of current forward pass
|
|
int* inputs; // the input tokens for the current forward pass
|
|
int* targets; // the target tokens for the current forward pass
|
|
float mean_loss;
|
|
float* cpu_losses;
|
|
} GPT2;
|
|
|
|
void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
|
|
|
|
// read in model from a checkpoint file
|
|
FILE *model_file = fopenCheck(checkpoint_path, "rb");
|
|
int model_header[256];
|
|
freadCheck(model_header, sizeof(int), 256, model_file);
|
|
if (model_header[0] != 20240326) { fprintf(stderr, "GPT2 - Bad magic model file\n"); exit(EXIT_FAILURE); }
|
|
if (model_header[1] != 3) {
|
|
// was bumped from 1 -> 3 to incorporate the padded vocab size
|
|
fprintf(stderr, "GPT2 - Bad version in model file\n");
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
|
|
// read in hyperparameters
|
|
model->config.max_seq_len = model_header[2];
|
|
model->config.vocab_size = model_header[3];
|
|
model->config.num_layers = model_header[4];
|
|
model->config.num_heads = model_header[5];
|
|
model->config.channels = model_header[6];
|
|
model->config.padded_vocab_size = model_header[7];
|
|
|
|
// allocate space for all the parameters and read them in
|
|
fill_in_parameter_sizes(model->param_sizes, model->config);
|
|
|
|
// count the number of parameters
|
|
size_t num_parameters = 0;
|
|
for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) {
|
|
num_parameters += model->param_sizes[i];
|
|
}
|
|
model->num_parameters = num_parameters;
|
|
|
|
// create memory for model parameters on the device
|
|
model->params_memory = malloc_and_point_parameters(&model->params, model->param_sizes, 1);
|
|
|
|
// read in all the parameters from file and copy them to device
|
|
float* params_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float));
|
|
freadCheck(params_memory_cpu, sizeof(float), num_parameters, model_file);
|
|
cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, num_parameters * sizeof(float), cudaMemcpyHostToDevice));
|
|
free(params_memory_cpu);
|
|
fcloseCheck(model_file);
|
|
|
|
// other inits
|
|
model->acts_memory = NULL;
|
|
model->grads_memory = NULL;
|
|
model->m_memory = NULL;
|
|
model->v_memory = NULL;
|
|
model->grads_acts_memory = NULL;
|
|
model->inputs = NULL;
|
|
model->targets = NULL;
|
|
model->cpu_losses = NULL;
|
|
model->batch_size = 0;
|
|
model->seq_len = 0;
|
|
model->mean_loss = -1.0f; // -1.0f will designate no loss
|
|
}
|
|
|
|
void gpt2_forward(GPT2 *model, int* inputs, int B, int T) {
|
|
|
|
// ensure the model was initialized or error out
|
|
if (model->params_memory == NULL) {
|
|
printf("GPT2 - Error: model was not initialized properly.\n");
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
|
|
// convenience parameters
|
|
int V = model->config.vocab_size;
|
|
int Vp = model->config.padded_vocab_size;
|
|
int L = model->config.num_layers;
|
|
int NH = model->config.num_heads;
|
|
int C = model->config.channels;
|
|
|
|
// validate inputs, all indices must be in the range [0, V)
|
|
for(int i = 0; i < B * T; i++) {
|
|
assert(0 <= inputs[i] && inputs[i] < V);
|
|
}
|
|
|
|
// allocate space for all the activations if needed (done here, lazily)
|
|
if(model->acts_memory == NULL || B != model->batch_size || T != model->seq_len) {
|
|
// record the current B,T as well
|
|
model->batch_size = B;
|
|
model->seq_len = T;
|
|
// and now allocate the space
|
|
fill_in_activation_sizes(model->act_sizes, B, T, model->config);
|
|
size_t num_activations = 0;
|
|
for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) {
|
|
num_activations += model->act_sizes[i];
|
|
}
|
|
model->num_activations = num_activations;
|
|
if (model->acts_memory != NULL) {
|
|
printf("GPT2 - Reallocating activations for B=%d T=%d\n", B, T);
|
|
cudaFree(model->acts_memory);
|
|
cudaFree(model->inputs);
|
|
}
|
|
model->acts_memory = malloc_and_point_activations(&model->acts, model->act_sizes);
|
|
printf("GPT2 - Allocated %zu MiB for activations\n", (num_activations * sizeof(float)) >> 20); // >> 20 is /(1024*1024)
|
|
cudaCheck(cudaMalloc((void**)&model->inputs, B * T * sizeof(int)));
|
|
cudaCheck(cudaMallocHost((void**)&model->cpu_losses, B * T * sizeof(float)));
|
|
}
|
|
|
|
// copy inputs to the model
|
|
cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice));
|
|
|
|
// forward pass
|
|
ParameterTensors params = model->params; // for brevity
|
|
ActivationTensors acts = model->acts;
|
|
float* residual;
|
|
encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0]
|
|
|
|
for (int l = 0; l < L; l++) {
|
|
|
|
residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C;
|
|
|
|
// get the pointers of the weights for this layer
|
|
float* l_ln1w = params.ln1w + l * C;
|
|
float* l_ln1b = params.ln1b + l * C;
|
|
float* l_qkvw = params.qkvw + l * 3*C * C;
|
|
float* l_qkvb = params.qkvb + l * 3*C;
|
|
float* l_attprojw = params.attprojw + l * C * C;
|
|
float* l_attprojb = params.attprojb + l * C;
|
|
float* l_ln2w = params.ln2w + l * C;
|
|
float* l_ln2b = params.ln2b + l * C;
|
|
float* l_fcw = params.fcw + l * 4*C * C;
|
|
float* l_fcb = params.fcb + l * 4*C;
|
|
float* l_fcprojw = params.fcprojw + l * C * 4*C;
|
|
float* l_fcprojb = params.fcprojb + l * C;
|
|
|
|
// get the pointers of the activations for this layer
|
|
float* l_ln1 = acts.ln1 + l * B * T * C;
|
|
float* l_ln1_mean = acts.ln1_mean;
|
|
float* l_ln1_rstd = acts.ln1_rstd;
|
|
float* l_qkvr = acts.qkvr + l * B * T * 3*C;
|
|
float* l_atty = acts.atty + l * B * T * C;
|
|
float* l_att = acts.att + l * B * NH * T * T;
|
|
float* l_attproj = acts.attproj + l * B * T * C;
|
|
float* l_residual2 = acts.residual2 + l * B * T * C;
|
|
float* l_ln2 = acts.ln2 + l * B * T * C;
|
|
float* l_ln2_mean = acts.ln2_mean;
|
|
float* l_ln2_rstd = acts.ln2_rstd;
|
|
float* l_fch = acts.fch + l * B * T * 4*C;
|
|
float* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C;
|
|
float* l_fcproj = acts.fcproj + l * B * T * C;
|
|
float* l_residual3 = acts.residual3 + l * B * T * C;
|
|
// these are only needed as scratchpads for the forward pass
|
|
float* scratch = acts.output;
|
|
|
|
// now do the forward pass
|
|
layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C);
|
|
matmul_forward(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C);
|
|
attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH);
|
|
matmul_forward(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C);
|
|
residual_forward(l_residual2, residual, l_attproj, B*T*C);
|
|
layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C);
|
|
//op16 matmul + bias + GeLU fusion
|
|
// matmul_forward(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C);
|
|
// gelu_forward(l_fch_gelu, l_fch, B*T*4*C);
|
|
matmul_forward_fused_gelu(l_fch_gelu, l_ln2, l_fcw, l_fcb, B, T, C, 4*C);
|
|
matmul_forward(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C);
|
|
residual_forward(l_residual3, l_residual2, l_fcproj, B*T*C);
|
|
}
|
|
|
|
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(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp);
|
|
|
|
model->mean_loss = -1.0f;
|
|
}
|
|
|
|
void gpt2_free(GPT2 *model) {
|
|
cudaCheck(cudaFree(model->params_memory));
|
|
cudaCheck(cudaFree(model->grads_memory));
|
|
cudaCheck(cudaFree(model->m_memory));
|
|
cudaCheck(cudaFree(model->v_memory));
|
|
cudaCheck(cudaFree(model->acts_memory));
|
|
cudaCheck(cudaFree(model->grads_acts_memory));
|
|
cudaCheck(cudaFree(model->inputs));
|
|
cudaCheck(cudaFree(model->targets));
|
|
cudaFreeHost(model->cpu_losses);
|
|
}
|
|
|
|
#define GPT2_EOT 50256 |