mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-20 23:05:08 -04:00
74 lines
2.9 KiB
Text
74 lines
2.9 KiB
Text
/*
|
|
This code is a convenience tool for profiling the CUDA kernels in the training
|
|
loop of train_gpt2.cu. Compile:
|
|
|
|
make profile_gpt2cu NO_MULTI_GPU=1
|
|
|
|
And then e.g. use ncu from NVIDIA. The CLI docs for example:
|
|
https://docs.nvidia.com/nsight-compute/NsightComputeCli/
|
|
|
|
TLDR run like:
|
|
|
|
sudo ncu --set full --import-source yes -o profile -f ./profile_gpt2cu
|
|
|
|
This:
|
|
- `--set full` means we'll collect A LOT of metrics. take out for less
|
|
- `--import-source yes` means we'll get the source code in the profile
|
|
- `-o profile` writes the results into file profile.ncu-rep
|
|
- `-f` forces overwrite of the profile.ncu-rep file
|
|
- `./profile_gpt2cu` is the executable we want to profile
|
|
|
|
This writes results into profile.ncu-rep output file.
|
|
You can open this up in NVIDIA Nsight Compute UI.
|
|
For example, I have NVIDIA Nsight Compute installed on my Mac, and I rsync
|
|
the profile.ncu-rep from a cloud box to local to pretty view.
|
|
*/
|
|
|
|
#define TESTING
|
|
#include "train_gpt2.cu"
|
|
|
|
int main(int argc, char *argv[]) {
|
|
char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi"
|
|
int num_processes = -1; // doesn't matter when using MPI
|
|
int process_rank = -1; // doesn't matter when using MPI
|
|
int gpus_per_node = -1; // doesn't matter when using MPI
|
|
char server_ip[256] = ""; // doesn't matter when using MPI
|
|
char fs_path[256] = ""; // doesn't matter when using MPI
|
|
multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method);
|
|
common_start(true, true);
|
|
|
|
// build the GPT-2 model from a checkpoint
|
|
GPT2 model;
|
|
gpt2_init_common(&model);
|
|
gpt2_build_from_checkpoint(&model, "gpt2_124M_bf16.bin");
|
|
|
|
int B = 24; // if program OOMs decrease this number, e.g. all the way down to 4 or etc
|
|
int T = 1024; // if even that OOMs move on to this one. keep them nice and powers of 2
|
|
printf("batch size: %d\n", B);
|
|
printf("sequence length: %d\n", T);
|
|
|
|
int* x = (int*)mallocCheck(B * T * sizeof(int));
|
|
int* y = (int*)mallocCheck(B * T * sizeof(int));
|
|
for(int i = 0; i < B * T; ++i) {
|
|
x[i] = i % model.config.vocab_size;
|
|
y[i] = i % model.config.vocab_size;
|
|
}
|
|
|
|
// override number of layers to 1 because all layers repeat the same kernels, only profile once
|
|
model.config.num_layers = 1;
|
|
set_zero_configs(&multi_gpu_config, 0, model.num_parameters);
|
|
|
|
gpt2_allocate_state(&model, B, T);
|
|
// do a training step
|
|
gpt2_forward(&model, x, B, T);
|
|
gpt2_backward_and_reduce(&model, x, y, 1, 0);
|
|
float grad_norm = gpt2_calculate_grad_norm(&model, &multi_gpu_config);
|
|
float grad_scale = (grad_norm > 1.0f) ? 1.0f / grad_norm : 1.0f;
|
|
gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, grad_scale, 1, &multi_gpu_config);
|
|
cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings
|
|
|
|
// free
|
|
gpt2_free(&model);
|
|
common_free(model);
|
|
return 0;
|
|
}
|