aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge branch 'karpathy:master' into dataloader_win_fixes

Browse source
This commit is contained in:
rosslwheeler 2024-05-23 23:12:15 -07:00 • committed by GitHub
commit a241a00ad4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 473 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

3
README.md
View file

@ -403,6 +403,9 @@ Lastly, I will be a lot more sensitive to complexity in the root folder of the p
- [llm.rs](https://github.com/yijunyu/llm.rs) by @[Yijun Yu](https://github.com/yijunyu): a Rust rewrite with the aim to have same performance
- [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project
- Swift
- [llm.swift](https://github.com/otabuzzman/llm.swift) by @[otabuzzman](https://github.com/otabuzzman): a Swift port of this project
- Zig
- [llm.zig](https://github.com/Saimirbaci/llm.zig) by @[saimirbaci](https://github.com/Saimirbaci): a Zig port of this project

223
rand.h Normal file
View file

@ -0,0 +1,223 @@
/*
Mersenne Twisters implementation, numerically identical to torch.
Example usage:
mt19937_state state;
manual_seed(&state, 137);
printf("%u\n", randint32(&state));
printf("%u\n", randint32(&state));
printf("%u\n", randint32(&state));
printf("%u\n", randint32(&state));
printf("%u\n", randint32(&state));
float t8[8];
normal_(t8, 8, 0, 1, &state);
for (int i = 0; i < 8; i++) {
printf("%f\n", t8[i]);
}
printf("%u\n", randint32(&state));
float t16[16];
normal_(t16, 16, 0, 1, &state);
for (int i = 0; i < 16; i++) {
printf("%f\n", t16[i]);
}
printf("%u\n", randint32(&state));
PyTorch reference (producing identical results):
import torch
torch.manual_seed(137)
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
t = torch.zeros(8);
t.normal_()
for i in range(len(t)) :
print(t[i].item())
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
t = torch.zeros(16);
t.normal_()
for i in range(len(t)) :
print(t[i].item())
print(torch.randint(0, 0xFFFFFFFF, [1]).item())
Both output:
4053805790
2173880614
380293709
1237255315
2986595568
0.7947664260864258
1.4369317293167114
- 0.2292192131280899
0.47556325793266296
- 0.6334410905838013
- 0.5791953802108765
- 0.0925704762339592
- 0.8659197092056274
2186503452
- 1.2813878059387207
- 2.646395683288574
- 0.06569503247737885
0.2180829495191574
- 0.46536165475845337
- 0.33108410239219666
2.5485482215881348
0.10425379872322083
0.8460659980773926
0.9462448358535767
- 0.2913765013217926
0.34313806891441345
- 1.1186704635620117
- 0.18305328488349915
- 2.3153159618377686
0.3961987793445587
2756748748
*/
#ifndef RAND_H
#define RAND_H
#include <math.h>
#define MERSENNE_STATE_M 397u
#define MERSENNE_STATE_N 624u
#define LMASK 0x7ffffffful
#define UMASK 0x80000000ul
// Copyright(c) Makoto Matsumoto and Takuji Nishimura
// This implementation follows PyTorch so that we are numerically identical when running verification tests.
typedef struct {
unsigned long long seed_;
int left_;
unsigned int next_;
unsigned int state_[MERSENNE_STATE_N];
unsigned int MATRIX_A[2];
} mt19937_state;
void manual_seed(mt19937_state* state, unsigned int seed) {
state->MATRIX_A[0] = 0x0u;
state->MATRIX_A[1] = 0x9908b0df;
state->state_[0] = seed & 0xffffffff;
for (unsigned int j = 1; j < MERSENNE_STATE_N; j++) {
state->state_[j] = 1812433253 * (state->state_[j - 1] ^ (state->state_[j - 1] >> 30)) + j;
state->state_[j] &= 0xffffffff;
}
state->left_ = 1;
state->next_ = 0;
}
void next_state(mt19937_state* state) {
state->left_ = MERSENNE_STATE_N;
state->next_ = 0;
unsigned int y, j;
for (j = 0; j < MERSENNE_STATE_N - MERSENNE_STATE_M; j++) {
y = (state->state_[j] & UMASK) | (state->state_[j + 1] & LMASK);
state->state_[j] = state->state_[j + MERSENNE_STATE_M] ^ (y >> 1) ^ state->MATRIX_A[y & 0x1];
}
for (; j < MERSENNE_STATE_N - 1; j++) {
y = (state->state_[j] & UMASK) | (state->state_[j + 1] & LMASK);
state->state_[j] = state->state_[j + (MERSENNE_STATE_M - MERSENNE_STATE_N)] ^ (y >> 1) ^ state->MATRIX_A[y & 0x1];
}
y = (state->state_[MERSENNE_STATE_N - 1] & UMASK) | (state->state_[0] & LMASK);
state->state_[MERSENNE_STATE_N - 1] = state->state_[MERSENNE_STATE_M - 1] ^ (y >> 1) ^ state->MATRIX_A[y & 0x1];
}
unsigned int randint32(mt19937_state* state) {
if (!state) return 0;
if (state->MATRIX_A[0] != 0 || state->MATRIX_A[1] != 0x9908b0df) manual_seed(state, 5489); // auto-initialize
if (--state->left_ <= 0) {
next_state(state);
}
unsigned int y = state->state_[state->next_++];
y ^= y >> 11;
y ^= (y << 7) & 0x9d2c5680;
y ^= (y << 15) & 0xefc60000;
y ^= y >> 18;
return y;
}
inline unsigned long long randint64(mt19937_state* state) {
return (((unsigned long long)(randint32(state)) << 32) | randint32(state));
}
inline float randfloat32(mt19937_state* state) {
return (randint32(state) & ((1ull << 24) - 1)) * (1.0f / (1ull << 24));
}
inline double randfloat64(mt19937_state* state) {
return (randint64(state) & ((1ull << 53) - 1)) * (1.0 / (1ull << 53));
}
void uniform_(float* data, unsigned int numel, float from, float to, mt19937_state* state) {
for (unsigned int t = 0; t < numel; t++) {
data[t] = randfloat32(state) * (to - from) + from;
}
}
// Box<6F>Muller transform
void normal_fill_16(float* data, float mean, float std, mt19937_state* state) {
#define EPSILONE 1e-12
for (unsigned int t = 0; t < 8; t++) {
float u1 = 1 - data[t];
float u2 = data[t + 8];
float radius = sqrtf(-2 * logf(u1 + EPSILONE));
float theta = 2.0 * M_PI * u2;
data[t] = (radius * cosf(theta) * std + mean);
data[t + 8] = (radius * sinf(theta) * std + mean);
}
}
void normal_fill(float* data, unsigned int numel, float mean, float std, mt19937_state* state) {
for (unsigned int t = 0; t < numel; t++) {
data[t] = randfloat32(state);
}
for (unsigned int i = 0; i < numel - 15; i += 16) {
normal_fill_16(data + i, mean, std, state);
}
if (numel % 16 != 0) {
// recompute the last 16 values
data = data + numel - 16;
for (unsigned int i = 0; i < 16; i++) {
data[i] = randfloat32(state);
}
normal_fill_16(data, mean, std, state);
}
}
void normal_(float* data, unsigned int numel, float mean, float std, mt19937_state* state) {
#define EPSILONE 1e-12
if (numel >= 16) {
normal_fill(data, numel, mean, std, state);
}
else {
double next_double_normal_sample;
int has_next_double_normal_sample = 0;
for (unsigned int t = 0; t < numel; t++) {
if (has_next_double_normal_sample) {
data[t] = (float)(next_double_normal_sample * std + mean);
has_next_double_normal_sample = 0;
continue;
}
// for numel < 16 we draw a double (float64)
float u1 = randfloat64(state);
float u2 = randfloat64(state);
float radius = sqrtf(-2 * logf(1 - u2 + EPSILONE));
float theta = 2.0 * M_PI * u1;
next_double_normal_sample = radius * sinf(theta);
has_next_double_normal_sample = 1;
data[t] = (radius * cosf(theta) * std + mean);
}
}
}
#endif

18
test_gpt2.cu
View file

@ -274,7 +274,7 @@ int main(int argc, char *argv[]) {
allok = allok & check_tensor(tensors1[15], tensors2[15], C, "lnfb", 2e-2f);
}
gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.01f, 1.f, step+1, &multi_gpu_config);
gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+1, &multi_gpu_config);
// print the timing information at the end
printf("step %d: loss %f (took %f ms)\n", step+1, model.mean_loss, time_elapsed_s * 1000);
@ -285,14 +285,14 @@ int main(int argc, char *argv[]) {
float expected_losses[10] = {
5.2700,
4.0607,
3.3166,
2.7115,
2.1702,
1.6349,
1.1419,
0.7038,
0.3769,
0.1743
3.3202,
2.7176,
2.1811,
1.6538,
1.1680,
0.7367,
0.4008,
0.1874
};
// compare

200
train_gpt2.cu
View file

@ -57,7 +57,9 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200),
// defines: dataloader_init, dataloader_reset, dataloader_next_batch, dataloader_free
// defines: evalloader_init, evalloader_reset, evalloader_next_batch, evalloader_free
#include "dataloader.h"
// defines: manual_seed, normal_
// numerically identical to PyTorch's torch.manual_seed and torch.normal
#include "rand.h"
// ----------------------------------------------------------------------------
// CUDA precision settings
@ -2035,6 +2037,110 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
model->recompute = 1; // default to recompute gelu during backward
}
void gpt2_build_from_random(GPT2 *model, int depth) {
// init random (training from scratch)
// parameterize the size of gpt2 based only on the depth of the model (num_layers)
model->config.num_layers = depth;
// follows GPT-2 sizes
int channels, num_heads;
if (depth == 12) { channels = 768; num_heads = 12; } // gpt2 (124M)
else if (depth == 24) { channels = 1024; num_heads = 16; } // gpt2-medium (350M)
else if (depth == 36) { channels = 1280; num_heads = 20; } // gpt2-large (774M)
else if (depth == 48) { channels = 1600; num_heads = 25; } // gpt2-xl (1558M)
else { fprintf(stderr, "Unsupported depth for now\n"); exit(EXIT_FAILURE); }
model->config.channels = channels;
model->config.num_heads = num_heads;
model->config.max_seq_len = 1024;
model->config.vocab_size = 50257;
model->config.padded_vocab_size = 50304; // padded to 128
// fill in all the parameter tensor dimensions and types
fill_in_parameter_sizes(model->param_elements, model->param_sizeof, model->config);
model->num_parameters = 0;
model->num_parameters_bytes = 0;
for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
model->num_parameters += model->param_elements[i];
model->num_parameters_bytes += model->param_elements[i] * model->param_sizeof[i];
}
// create memory for model parameters on the device
model->params_memory = malloc_and_point_parameters(&model->params, model->param_elements, model->param_sizeof);
// allocate and random init the memory for all the parameters with GPT-2 schema
// weights ~N(0, 0.02), biases 0, c_proj weights ~N(0, 0.02/(2*L)**0.5)
// NOTE: assuming all parameters are of the type floatX, could be relaxed later
mt19937_state init_rng;
manual_seed(&init_rng, 42);
floatX* params_memory_cpu = (floatX*)mallocCheck(model->num_parameters_bytes);
memset(params_memory_cpu, 0, model->num_parameters_bytes);
// fill in all the weights with random values
float residual_scale = 1.0f / sqrtf(2.0f * model->config.num_layers);
// we have to init all these tensors exactly in the order that PyTorch initializes them
// so that we can match them up and get correctness and exactly the same initial conditions
size_t L = model->config.num_layers;
size_t offset = 0;
for (int l = 0; l < L; l++) {
offset = 0;
for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
// the layernorm parameters are all initialized to 1
if (l == 0 && (i == 2 || i == 8 || i == 14)) { // only at l = 0 to init these just once
for (size_t j = 0; j < model->param_elements[i]; j++) {
params_memory_cpu[offset + j] = 1.0f;
}
}
// weights tensors are handled here
if ((l == 0 && (i == 0 || i == 1)) // only at l = 0, init the wte and wpe tensors
|| i == 4 || i == 6 || i == 10 || i == 12) {
int n = model->param_elements[i];
size_t layer_offset = 0;
if (i == 0) {
// for wte tensor (padded vocab) override to init V instead of Vp rows
n = model->config.vocab_size * model->config.channels;
}
if (i == 4 || i == 6 || i == 10 || i == 12) {
// weight tensors, we are only initializing layer l
assert(n % L == 0);
n = n / L;
layer_offset = l * n;
}
// in GPT-2, the projections back into the residual stream are additionally
// scaled by 1/sqrt(2*L) for training stability
float scale = (i == 6 || i == 12) ? 0.02f * residual_scale : 0.02f;
// okay let's draw the random numbers and write them
float *fp32_buffer = (float*)mallocCheck(n * sizeof(float));
normal_(fp32_buffer, n, 0.0f, scale, &init_rng);
for (size_t j = 0; j < n; j++) {
params_memory_cpu[offset + layer_offset + j] = (floatX)fp32_buffer[j];
}
free(fp32_buffer);
}
offset += model->param_elements[i];
}
}
// copy them to GPU
cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, model->num_parameters_bytes, cudaMemcpyHostToDevice));
free(params_memory_cpu);
// other inits and defaults
model->acts_memory = NULL;
model->grads_memory = NULL;
model->m_memory = NULL;
model->v_memory = NULL;
model->master_weights = NULL;
model->grads_acts_memory = NULL;
model->inputs = NULL;
model->targets = NULL;
model->cpu_losses = NULL;
model->cpu_losses_fp32 = NULL;
model->batch_size = 0;
model->seq_len = 0;
model->mean_loss = -1.0f; // -1.0f designates no loss
model->rng_state = 13371337;
model->use_master_weights = 1; // keep master weights copy in float for optim update?
model->recompute = 1; // default to recompute gelu during backward
}
void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, int grad_accum_steps=1) {
NVTX_RANGE_FN();
// targets are optional and could be NULL
@ -2524,14 +2630,17 @@ typedef struct {
} Logger;
void logger_init(Logger *logger, const char *filename) {
logger->flush_every = 20;
logger->flush_every = 10;
logger->logfile = NULL;
if (filename != NULL) { logger->logfile = fopenCheck(filename, "w"); }
// only rank 0 process will log
if (filename != NULL && multi_gpu_config.process_rank == 0) {
logger->logfile = fopenCheck(filename, "w");
}
}
void logger_log_eval(Logger *logger, int step, float val_loss) {
void logger_log_eval(Logger *logger, int step, float val) {
if (logger->logfile != NULL) {
fprintf(logger->logfile, "s:%d eval:%.4f\n", step, val_loss);
fprintf(logger->logfile, "s:%d eval:%.4f\n", step, val);
}
}
@ -2544,7 +2653,7 @@ void logger_log_val(Logger *logger, int step, float val_loss) {
void logger_log_train(Logger *logger, int step, float train_loss) {
if (logger->logfile != NULL) {
fprintf(logger->logfile, "s:%d trl:%.4f\n", step, train_loss);
if (step % 10 == 0) { fflush(logger->logfile); }
if (step % logger->flush_every == 0) { fflush(logger->logfile); }
}
}
@ -2566,6 +2675,9 @@ void error_usage() {
fprintf(stderr, " -t <int> sequence length T (default = 1024)\n");
fprintf(stderr, " -d <int> total desired batch size (default = B * T * num_processes, i.e. no grad accumulation\n");
fprintf(stderr, " -l <float> learning rate (default = 3e-4f)\n");
fprintf(stderr, " -u <int> learning rate warmup iterations (default = 0, no warmup)\n");
fprintf(stderr, " -q <float> learning rate decay: final fraction, at end of training (default = 1.0 (no decay))\n");
fprintf(stderr, " -c <float> weight decay (default = 0.0f)\n");
fprintf(stderr, " -x <int> max_steps of optimization to run (-1 (default) = disable, run 1 epoch)\n");
fprintf(stderr, " -v <int> val_loss_every, how often we evaluate val loss (default = 20)\n");
fprintf(stderr, " -m <int> val_max_batches, up to how many val batches to estimate val loss? (default = 20)\n");
@ -2576,6 +2688,7 @@ void error_usage() {
fprintf(stderr, " -w <int> keep f32 copy of weights for the optimizer? (default: 1)\n");
fprintf(stderr, " -z <int> zero_stage, Zero Optimization Stage, 0,1,2,3 (default = 0)\n");
fprintf(stderr, " -r <int> recompute: saves memory at cost of speed. (default = 1), 0 = none. 1 = recompute gelu\n");
fprintf(stderr, " -h <int> hellaswag eval run? (default = 0)\n");
exit(EXIT_FAILURE);
}
@ -2593,6 +2706,9 @@ int main(int argc, char *argv[]) {
int T = 1024; // sequence length max
int total_batch_size = -1; // will be calculated down below later, if not provided
float learning_rate = 3e-4f;
int warmup_iterations = 0;
float final_learning_rate_frac = 1.0f; // final fraction of learning rate, at end of training
float weight_decay = 0.0f;
int val_loss_every = 20; // every how many steps do we eval validation loss?
int val_max_batches = 20; // how many batches max do we eval for validation loss?
int sample_every = 20; // every how many steps to do inference?
@ -2604,6 +2720,7 @@ int main(int argc, char *argv[]) {
int recompute = 1; // recompute during backward setting, 0 = none, 1 = recompute gelu
int zero_stage = 0; // Zero Optimization Stage for Multi-GPU training
float grad_clip = 1.0f;
int hellaswag_eval = 0;
for (int i = 1; i < argc; i+=2) {
if (i + 1 >= argc) { error_usage(); } // must have arg after flag
if (argv[i][0] != '-') { error_usage(); } // must start with dash
@ -2617,6 +2734,9 @@ int main(int argc, char *argv[]) {
else if (argv[i][1] == 't') { T = atoi(argv[i+1]); }
else if (argv[i][1] == 'd') { total_batch_size = atoi(argv[i+1]); }
else if (argv[i][1] == 'l') { learning_rate = atof(argv[i+1]); }
else if (argv[i][1] == 'u') { warmup_iterations = atoi(argv[i+1]); }
else if (argv[i][1] == 'q') { final_learning_rate_frac = atof(argv[i+1]); }
else if (argv[i][1] == 'c') { weight_decay = atof(argv[i+1]); }
else if (argv[i][1] == 'x') { max_steps = atoi(argv[i+1]); }
else if (argv[i][1] == 'v') { val_loss_every = atoi(argv[i+1]); }
else if (argv[i][1] == 'm') { val_max_batches = atoi(argv[i+1]); }
@ -2628,8 +2748,11 @@ int main(int argc, char *argv[]) {
else if (argv[i][1] == 'c') { grad_clip = atof(argv[i+1]); }
else if (argv[i][1] == 'z') { zero_stage = atoi(argv[i+1]); }
else if (argv[i][1] == 'r') { recompute = atoi(argv[i+1]); }
else if (argv[i][1] == 'h') { hellaswag_eval = atoi(argv[i+1]); }
else { error_usage(); }
}
// should do a bit more error checking here
assert(warmup_iterations >= 0);
// calculate a sensible default for total batch size by assuming no gradient accumulation
if (total_batch_size == -1) { total_batch_size = B * T * multi_gpu_config.num_processes; }
// if we're only overfitting a single batch for debugging, let's overfit the first batch
@ -2644,7 +2767,10 @@ int main(int argc, char *argv[]) {
printf0("| micro batch size B | %-50d |\n", B);
printf0("| sequence length T | %-50d |\n", T);
printf0("| total batch size | %-50d |\n", total_batch_size);
printf0("| learning rate | %-50e |\n", learning_rate);
printf0("| learning rate (LR) | %-50e |\n", learning_rate);
printf0("| warmup iterations | %-50d |\n", warmup_iterations);
printf0("| final LR fraction | %-50e |\n", final_learning_rate_frac);
printf0("| weight decay | %-50e |\n", weight_decay);
printf0("| grad_clip | %-50e |\n", grad_clip);
printf0("| max_steps | %-50d |\n", max_steps);
printf0("| val_loss_every | %-50d |\n", val_loss_every);
@ -2666,9 +2792,23 @@ int main(int argc, char *argv[]) {
printf0("| precision | %-50s |\n", precision_str);
printf0("+-----------------------+----------------------------------------------------+\n");
// build the GPT-2 model from a checkpoint
// build the GPT-2 model
GPT2 model;
gpt2_build_from_checkpoint(&model, load_filename);
// if load_filename is of the form "dX" where X is an integer (e.g. d12), then we build
// a random model with the depth of the model specified by X (e.g. 12). otherwise interpret
// this variable as a checkpoint filename, and load that checkpoint
assert(strlen(load_filename) >= 2);
if (load_filename[0] == 'd') {
int depth = atoi(load_filename + 1);
if (depth > 1 && depth <= 1000) { // we're not going to train models this big right? heh
gpt2_build_from_random(&model, depth);
} else {
exit(EXIT_FAILURE);
}
} else {
gpt2_build_from_checkpoint(&model, load_filename);
}
model.use_master_weights = use_master_weights;
model.recompute = recompute;
printf0("| load_filename | %-50s |\n", load_filename);
@ -2695,10 +2835,11 @@ int main(int argc, char *argv[]) {
EvalLoader eval_loader;
const char* hellaswag_path = "dev/data/hellaswag/hellaswag_val.bin";
const char hellaswag_available = access(hellaswag_path, F_OK) == 0;
if (hellaswag_available) {
const char run_hellaswag = hellaswag_eval && hellaswag_available;
if (run_hellaswag) {
evalloader_init(&eval_loader, hellaswag_path, B, T, multi_gpu_config.process_rank, multi_gpu_config.num_processes);
}
printf0("| hellaswag available | %-50s |\n", hellaswag_available ? "yes" : "no");
printf0("| run hellaswag | %-50s |\n", run_hellaswag ? "yes" : "no");
printf0("+-----------------------+----------------------------------------------------+\n");
// pretty print in a table the multi-gpu configuration as well
@ -2710,7 +2851,7 @@ int main(int argc, char *argv[]) {
// prints outside of pretty table to here and below
if (!hellaswag_available) {
printf0("HellaSwag eval not found at %s, skipping its evaluation\n", hellaswag_path);
printf0("You can run `python dev/data/hellaswag.py` to export and use it.\n");
printf0("You can run `python dev/data/hellaswag.py` to export and use it with `-h 1`.\n");
}
// more prints related to allocations from gpt2_build_from_checkpoint down here to not mess up our table above
printf0("num_parameters: %zu => bytes: %zu\n", model.num_parameters, model.num_parameters_bytes);
@ -2724,9 +2865,11 @@ int main(int argc, char *argv[]) {
B, T, multi_gpu_config.num_processes, total_batch_size);
printf0("=> setting grad_accum_steps=%d\n", grad_accum_steps);
// set up the Logger & Tokenizer
// set up the Logger
Logger logger;
logger_init(&logger, output_log_file);
// set up the Tokenizer
Tokenizer tokenizer;
tokenizer_init(&tokenizer, "gpt2_tokenizer.bin");
@ -2765,8 +2908,8 @@ int main(int argc, char *argv[]) {
}
// once in a while estimate HellaSwag accuracy
if (hellaswag_available &&
(step % val_loss_every == 0 || last_step)) {
if (run_hellaswag &&
((step > 0 && step % val_loss_every == 0) || last_step)) {
NvtxRange evaluation_range("evaluation");
float eval_acc_norm = 0.0f;
evalloader_reset(&eval_loader);
@ -2780,7 +2923,7 @@ int main(int argc, char *argv[]) {
// careful because not all ranks may have the exact same allocation of number of examples
eval_acc_norm = multi_gpu_cpu_float_sum(eval_acc_norm);
printf0("HellaSwag: %d/%d = %f\n", (int)eval_acc_norm, eval_loader.num_examples, eval_acc_norm / eval_loader.num_examples);
logger_log_eval(&logger, step, eval_acc_norm);
logger_log_eval(&logger, step, eval_acc_norm / eval_loader.num_examples);
}
// once in a while do model inference to print generated text
@ -2857,7 +3000,20 @@ int main(int argc, char *argv[]) {
model.mean_loss = lossf;
// update the parameters
gpt2_multi_gpu_accumulate(&model, &multi_gpu_config);
float grad_norm = gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, grad_clip, step+1, &multi_gpu_config);
// learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac
float step_learning_rate = learning_rate;
if (step < warmup_iterations) {
step_learning_rate = learning_rate * ((float)(step + 1)) / warmup_iterations;
} else {
float decay_ratio = ((float)(step - warmup_iterations)) / (train_num_batches - warmup_iterations);
assert(0.0f <= decay_ratio && decay_ratio <= 1.0f);
float coeff = 0.5f * (1.0f + cosf(M_PI * decay_ratio)); // coeff starts at 1 and goes to 0
assert(0.0f <= coeff && coeff <= 1.0f);
float min_lr = learning_rate * final_learning_rate_frac;
step_learning_rate = min_lr + coeff * (learning_rate - min_lr);
}
// update the model parameters
float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, grad_clip, step+1, &multi_gpu_config);
gpt2_multi_gpu_gather(&model, &multi_gpu_config);
// zero out the gradients for the next iteration
gpt2_zero_grad(&model);
@ -2870,17 +3026,17 @@ int main(int argc, char *argv[]) {
float time_elapsed_ms;
cudaCheck(cudaEventElapsedTime(&time_elapsed_ms, start, end));
size_t tokens_processed = (size_t)multi_gpu_config.num_processes * B * T * grad_accum_steps;
float tokens_per_second = tokens_processed / time_elapsed_ms * 1000.0;
float tokens_per_second = tokens_processed / time_elapsed_ms * 1000.0f;
float bias_corrected_ema_tokens_per_second = tokens_per_second; // by default set to non-ema version
if (step > 0) { // consider the first batch to be a warmup (e.g. cuBLAS/cuDNN initialisation)
total_sum_iteration_time_s += time_elapsed_ms / 1000.0;
total_sum_iteration_time_s += time_elapsed_ms / 1000.0f;
// smooth out the tok/s with an exponential moving average, and bias correct just like in AdamW
ema_tokens_per_second = 0.95f * ema_tokens_per_second + 0.05f * tokens_per_second;
bias_corrected_ema_tokens_per_second = ema_tokens_per_second / (1.0f - powf(0.95f, step));
}
float accumulated_loss = multi_gpu_config.num_processes == 1 ? model.mean_loss : model.accumulated_mean_loss;
printf0("step %4d/%d: train loss %f norm %.4f (%.2f ms, %.0f tok/s)\n",
step + 1, train_num_batches, accumulated_loss, grad_norm,
printf0("step %4d/%d: train loss %f norm %.4f lr %.2e (%.2f ms, %.0f tok/s)\n",
step + 1, train_num_batches, accumulated_loss, grad_norm, step_learning_rate,
time_elapsed_ms, bias_corrected_ema_tokens_per_second);
logger_log_train(&logger, step, model.mean_loss);
@ -2893,7 +3049,7 @@ int main(int argc, char *argv[]) {
// free and destroy everything
cudaCheck(cudaEventDestroy(end));
cudaCheck(cudaEventDestroy(start));
if (hellaswag_available) { evalloader_free(&eval_loader); }
if (run_hellaswag) { evalloader_free(&eval_loader); }
dataloader_free(&train_loader);
dataloader_free(&val_loader);
tokenizer_free(&tokenizer);

76
train_gpt2.py
View file

@ -47,6 +47,7 @@ class CausalSelfAttention(nn.Module):
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
# output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1
# regularization
self.n_head = config.n_head
self.n_embd = config.n_embd
@ -84,6 +85,7 @@ class MLP(nn.Module):
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
self.gelu = NewGELU()
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1
def forward(self, x):
x = self.c_fc(x)
@ -126,8 +128,27 @@ class GPT(nn.Module):
ln_f = nn.LayerNorm(config.n_embd),
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
self.lm_head.LLMC_SKIP_INIT = 1 # don't init this one, we will tie weights
self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
# init all weights, use a torch rng object to be very careful
self.init_rng = torch.Generator()
self.init_rng.manual_seed(42)
self.apply(self._init_weights)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
# apply special scaled init to the residual projections, per GPT-2 paper
std = 0.02 if not hasattr(module, 'LLMC_RESIDUAL_SCALE_FLAG') else 0.02/math.sqrt(2 * self.config.n_layer)
# we want to skip initializing lm_head, which shares parameters with wte
# and wte was already initialized down below during the Embedding init
if not hasattr(module, 'LLMC_SKIP_INIT'):
torch.nn.init.normal_(module.weight, mean=0.0, std=std, generator=self.init_rng)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02, generator=self.init_rng)
def forward(self, idx, targets=None, return_logits=True):
device = idx.device
b, t = idx.size()
@ -395,7 +416,7 @@ if __name__ == "__main__":
# python train_gpt2.py --inference_only 1 --write_tensors 0 --sequence_length 1024
parser = argparse.ArgumentParser()
parser.add_argument("--input_bin", type=str, default="dev/data/tinyshakespeare/tiny_shakespeare_val.bin", help="input .bin to train on")
parser.add_argument("--model", type=str, default="gpt2", help="gpt2|gpt2-medium|gpt2-large|gpt2-xl")
parser.add_argument("--model", type=str, default="gpt2", help="gpt2|gpt2-medium|gpt2-large|gpt2-xl|d12|d24|d36|d48")
parser.add_argument("--write_tensors", type=int, default=1, help="write tensors to disk")
parser.add_argument("--inference_only", type=int, default=0, help="only run inference")
parser.add_argument("--dtype", type=str, default="float32", help="float32|float16|bfloat16")
@ -408,12 +429,12 @@ if __name__ == "__main__":
parser.add_argument("--sequence_length", type=int, default=64, help="sequence length")
parser.add_argument("--total_batch_size", type=int, default=256, help="total desired batch size, in units of #tokens")
parser.add_argument("--grad_clip", type=float, default=1.0, help="maximum gradient magnitude")
parser.add_argument("--overfit_single_batch", type=int, default=1, help="overfit just one batch of data")
args = parser.parse_args()
B, T = args.batch_size, args.sequence_length
assert 1 <= T <= 1024
assert args.dtype in {"float32", "float16", "bfloat16"}
assert args.model in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"}
model_to_size = {"gpt2": "124M", "gpt2-medium": "355M", "gpt2-large": "774M", "gpt2-xl": "1558M"}
assert args.model in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"}
# set up DDP (distributed data parallel). torchrun sets this env variable
ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
@ -427,8 +448,10 @@ if __name__ == "__main__":
device = f'cuda:{ddp_local_rank}'
torch.cuda.set_device(device)
master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
seed_offset = ddp_rank # each process gets a different seed
seed_offset = 0 # each process gets the exact same seed
else:
ddp_rank = 0
ddp_local_rank = 0
ddp_world_size = 1
master_process = True
seed_offset = 0
@ -449,8 +472,9 @@ if __name__ == "__main__":
tokens_per_fwdbwd = B * T * ddp_world_size
assert args.total_batch_size % tokens_per_fwdbwd == 0
grad_accum_steps = args.total_batch_size // tokens_per_fwdbwd
print(f"total desired batch size: {args.total_batch_size}")
print(f"=> calculated gradient accumulation steps: {grad_accum_steps}")
if master_process:
print(f"total desired batch size: {args.total_batch_size}")
print(f"=> calculated gradient accumulation steps: {grad_accum_steps}")
# set up a context manager following the desired dtype and device
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype]
@ -480,8 +504,19 @@ if __name__ == "__main__":
if master_process and args.write_tensors: # tokenizer is technically not tensors but ok
write_tokenizer(enc, "gpt2_tokenizer.bin")
# load the GPT-2 model weights
model = GPT.from_pretrained(args.model)
# init the model, either from scratch or from OpenAI pretrained checkpoint
if args.model[0] == "d":
# from scratch (random weights)
model_config = {
"d12": GPTConfig(block_size=1024, vocab_size=50257, n_layer=12, n_head=12, n_embd=768),
"d24": GPTConfig(block_size=1024, vocab_size=50257, n_layer=24, n_head=16, n_embd=1024),
"d36": GPTConfig(block_size=1024, vocab_size=50257, n_layer=36, n_head=20, n_embd=1280),
"d48": GPTConfig(block_size=1024, vocab_size=50257, n_layer=48, n_head=25, n_embd=1600),
}[args.model]
model = GPT(model_config)
else:
# load the GPT-2 model weights
model = GPT.from_pretrained(args.model)
model.train()
model.to(device)
if args.compile:
@ -526,18 +561,20 @@ if __name__ == "__main__":
def get_batch():
assert B*T+1 <= len(tokens), "not enough tokens"
# for 338,025 tokens. E.g. with B=8 T=1024, this will yield 41 batches before looping
i = 0
i = B*T*ddp_rank
while True:
x = tokens[i:i+B*T].view(B, T)
y = tokens[i+1:i+B*T+1].view(B, T)
yield x, y
i += B*T
i += B*T*ddp_world_size
if i + B*T + 1 >= len(tokens):
i = 0 # in prod we'd want to randomize the start point a bit
print("We do not expect to reach here in PyTorch right now")
import sys; sys.exit()
# fetch one batch of data, which we will overfit to
# fetch one batch of data
data_iter = iter(get_batch())
x, y = next(data_iter) # we'll overfit this batch below
x, y = next(data_iter)
x = x.to(device)
y = y.to(device)
@ -549,7 +586,9 @@ if __name__ == "__main__":
logits, loss = model(x, y)
loss.backward()
# save model params, in both float32 and bfloat16
model_size_str = model_to_size[args.model] # e.g. "124M"
model_to_size = {"gpt2": "124M", "gpt2-medium": "355M", "gpt2-large": "774M", "gpt2-xl": "1558M"}
model_to_size.update({f"d{d}": f"d{d}" for d in [12, 24, 36, 48]})
model_size_str = model_to_size[args.model] # e.g. "124M", or "d12"
write_model(model, f"gpt2_{model_size_str}.bin", dtype="float32")
write_model(model, f"gpt2_{model_size_str}_bf16.bin", dtype="bfloat16")
# save x, y, logits, loss, and parameter gradients, for debugging C
@ -566,7 +605,7 @@ if __name__ == "__main__":
# init the optimizer
adam_use_fused = device == "cuda" # only works on CUDA (?)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, fused=adam_use_fused)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, fused=adam_use_fused)
if device == "cuda":
torch.cuda.reset_peak_memory_stats()
@ -587,12 +626,17 @@ if __name__ == "__main__":
# instead of a SUM we want MEAN, so we scale the loss here
loss = loss / grad_accum_steps
lossf += loss.item() # keep track of the mean loss
# advance the dataset for the next batch
if not args.overfit_single_batch:
x, y = next(data_iter)
x = x.to(device)
y = y.to(device)
# backward pass
if ddp:
# we want only the last micro-step to sync grads in a DDP model
# the official way to do this is with model.no_sync(), but that is a
# context manager that bloats the code, so we just toggle this variable
model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1)
# backward pass
if not args.inference_only:
loss.backward()
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
@ -608,7 +652,7 @@ if __name__ == "__main__":
t1 = time.time()
# the 0th iteration is often an outlier (much slower) => skip logging it
tokens_per_second = grad_accum_steps * ddp_world_size * B * T / (t1-t0)
print0(f"iteration {step+1}, loss: {lossf:.4f}, time: {(t1-t0)*1000:.3f}ms, tok/s: {tokens_per_second:.2f}, norm: {norm:.3f}")
print0(f"step {step+1:4d}/{args.num_iterations}: train loss {lossf:.6f} norm {norm:.4f} lr 1.00e-04 ({(t1-t0)*1000:.3f} ms, {tokens_per_second:.0f} tok/s)")
if step > 0 and step > args.num_iterations - 20:
timings.append(t1-t0)