mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-27 20:25:09 -04:00
Merge remote-tracking branch 'karpathy/master' into cudnn_try2
This commit is contained in:
ademeure
commit
c4ecc04dc3
8 changed files with 236 additions and 66 deletions
|
|
@ -290,6 +290,10 @@ Lastly, I will be a lot more sensitive to complexity in the root folder of the p
|
|||
|
||||
## notable forks
|
||||
|
||||
- CUDA C++
|
||||
- [llm.cpp](https://github.com/gevtushenko/llm.c) by @[gevtushenko](https://github.com/gevtushenko): a port of this project using the [CUDA C++ Core Libraries](https://github.com/NVIDIA/cccl)
|
||||
- A presentation this fork was covered in [this lecture](https://www.youtube.com/watch?v=WiB_3Csfj_Q) in the [CUDA MODE Discord Server](https://discord.gg/cudamode)
|
||||
|
||||
- Mojo
|
||||
- [llm.🔥](https://github.com/dorjeduck/llm.mojo) by @[dorjeduck](https://github.com/dorjeduck): a Mojo port of this project
|
||||
|
||||
|
|
|
|||
|
|
@ -31,3 +31,9 @@ You'll see that this first forwards the reference code on the CPU, then it runs
|
|||
You'll see that this matches all the CPU results but runs much much faster. The typical process from here on is we copy paste the kernel that ran fastest, adjust it manually (e.g. to hardcode the best block size) and drop it into the training code file, e.g. `train_gpt2.cu`.
|
||||
|
||||
To add a new version of a kernel, add the kernel to the corresponding file and adjust the docs. To add a new kernel, add the new file and adjust the Makefile. Run `make clean` to clean up binaries from your directory.
|
||||
|
||||
If you do not have a GPU or is having trouble with CUDA dependencies, you can run the benchmarks on the [Modal platform](http://modal.com). For example, to run the benchmark for the attention forward pass on an A100 GPU with 80GB of memory, you can run the following command:
|
||||
|
||||
```bash
|
||||
GPU_MEM=80 modal run benchmark_on_modal.py --compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" --run-command "./attention_forward 1"
|
||||
```
|
||||
|
|
|
|||
82
dev/cuda/benchmark_on_modal.py
Normal file
82
dev/cuda/benchmark_on_modal.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
"""
|
||||
Script for running benchmarks on the Modal platform.
|
||||
This is useful for folks who do not have access to expensive GPUs locally.
|
||||
|
||||
Example usage:
|
||||
GPU_MEM=80 modal run benchmark_on_modal.py \
|
||||
--compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" \
|
||||
--run-command "./attention_forward 1"
|
||||
|
||||
This will mount the contents of the current directory to the remote container on modal,
|
||||
compile the `attention_forward.cu` file with `nvcc`, and run the resulting binary on a A100 GPU with 80GB of memory.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import sys
|
||||
|
||||
import modal
|
||||
from modal import Image, Stub
|
||||
|
||||
GPU_NAME_TO_MODAL_CLASS_MAP = {
|
||||
"H100": modal.gpu.H100,
|
||||
"A100": modal.gpu.A100,
|
||||
"A10G": modal.gpu.A10G,
|
||||
}
|
||||
|
||||
N_GPUS = int(os.environ.get("N_GPUS", 1))
|
||||
GPU_MEM = int(os.environ.get("GPU_MEM", 40))
|
||||
GPU_NAME = os.environ.get("GPU_NAME", "A100")
|
||||
GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, memory=GPU_MEM)
|
||||
|
||||
APP_NAME = "llm.c benchmark run"
|
||||
|
||||
# We don't actually need to use the Axolotl image here, but it's reliable
|
||||
AXOLOTL_REGISTRY_SHA = (
|
||||
"d5b941ba2293534c01c23202c8fc459fd2a169871fa5e6c45cb00f363d474b6a"
|
||||
)
|
||||
axolotl_image = (
|
||||
Image.from_registry(f"winglian/axolotl@sha256:{AXOLOTL_REGISTRY_SHA}")
|
||||
.run_commands(
|
||||
"git clone https://github.com/OpenAccess-AI-Collective/axolotl /root/axolotl",
|
||||
"cd /root/axolotl && git checkout v0.4.0",
|
||||
)
|
||||
.pip_install("huggingface_hub==0.20.3", "hf-transfer==0.1.5")
|
||||
.env(
|
||||
dict(
|
||||
HUGGINGFACE_HUB_CACHE="/pretrained",
|
||||
HF_HUB_ENABLE_HF_TRANSFER="1",
|
||||
TQDM_DISABLE="true",
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
stub = Stub(APP_NAME)
|
||||
|
||||
|
||||
def execute_command(command: str):
|
||||
command_args = command.split(" ")
|
||||
print(f"{command_args = }")
|
||||
subprocess.run(command_args, stdout=sys.stdout, stderr=subprocess.STDOUT)
|
||||
|
||||
|
||||
@stub.function(
|
||||
gpu=GPU_CONFIG,
|
||||
image=axolotl_image,
|
||||
allow_concurrent_inputs=4,
|
||||
container_idle_timeout=900,
|
||||
# This copies everything in this folder to the remote root folder
|
||||
mounts=[modal.Mount.from_local_dir("./", remote_path="/root/")]
|
||||
)
|
||||
def run_benchmark(compile_command: str, run_command: str):
|
||||
execute_command("pwd")
|
||||
execute_command("ls")
|
||||
execute_command(compile_command)
|
||||
execute_command(run_command)
|
||||
return None
|
||||
|
||||
|
||||
@stub.local_entrypoint()
|
||||
def inference_main(compile_command: str, run_command: str):
|
||||
results = run_benchmark.remote(compile_command, run_command)
|
||||
return results
|
||||
|
|
@ -52,17 +52,15 @@ struct alignas(16) Packed128 {
|
|||
__device__ const ElementType& operator[](int index) const {
|
||||
return payload[index];
|
||||
}
|
||||
__device__ float fp32(int index) {
|
||||
return static_cast<float>(payload[index]);
|
||||
}
|
||||
__device__ int4 get_bits() const {
|
||||
int4 bits;
|
||||
static_assert(sizeof(bits) == sizeof(payload), "Size mismatch.");
|
||||
memcpy(&bits, &payload, sizeof(bits));
|
||||
return bits;
|
||||
}
|
||||
|
||||
static constexpr const size_t size = sizeof(int4) / sizeof(ElementType);
|
||||
// e.g. sizeof(int4) is 16 (4 X 4 bytes), sizeof(bfloat16) = 2, so size = 8
|
||||
// so in the case where ElementType = bfloat16, we store 8 elements in one Packed128
|
||||
static constexpr const int size = sizeof(int4) / sizeof(ElementType);
|
||||
ElementType payload[size];
|
||||
};
|
||||
|
||||
|
|
@ -148,7 +146,7 @@ void validate_result(D* device_result, const T* cpu_reference, const char* name,
|
|||
printf("%f %f\n", cpu_reference[i], (T)out_gpu[i]);
|
||||
}
|
||||
// ensure correctness for all elements. We can set an "ignore" mask by writing NaN
|
||||
if (fabs(cpu_reference[i] - (T)out_gpu[i]) > tolerance && !isnan(cpu_reference[i])) {
|
||||
if (fabs(cpu_reference[i] - (T)out_gpu[i]) > tolerance && isfinite(cpu_reference[i])) {
|
||||
printf("Mismatch of %s at %d: CPU_ref: %f vs GPU: %f\n", name, i, cpu_reference[i], (T)out_gpu[i]);
|
||||
nfaults ++;
|
||||
if (nfaults >= 10) {
|
||||
|
|
@ -169,17 +167,35 @@ void validate_result(D* device_result, const T* cpu_reference, const char* name,
|
|||
template<class Kernel, class... KernelArgs>
|
||||
float benchmark_kernel(int repeats, Kernel kernel, KernelArgs&&... kernel_args) {
|
||||
cudaEvent_t start, stop;
|
||||
// prepare buffer to scrub L2 cache between benchmarks
|
||||
// just memset a large dummy array, recommended by
|
||||
// https://stackoverflow.com/questions/31429377/how-can-i-clear-flush-the-l2-cache-and-the-tlb-of-a-gpu
|
||||
// and apparently used in nvbench.
|
||||
int deviceIdx = 0;
|
||||
cudaCheck(cudaSetDevice(deviceIdx));
|
||||
cudaDeviceProp deviceProp;
|
||||
cudaCheck(cudaGetDeviceProperties(&deviceProp, deviceIdx));
|
||||
void* flush_buffer;
|
||||
cudaCheck(cudaMalloc(&flush_buffer, deviceProp.l2CacheSize));
|
||||
|
||||
cudaCheck(cudaEventCreate(&start));
|
||||
cudaCheck(cudaEventCreate(&stop));
|
||||
cudaCheck(cudaEventRecord(start, nullptr));
|
||||
float elapsed_time = 0.f;
|
||||
for (int i = 0; i < repeats; i++) {
|
||||
// clear L2
|
||||
cudaCheck(cudaMemset(flush_buffer, 0, deviceProp.l2CacheSize));
|
||||
// now we can start recording the timing of the kernel
|
||||
cudaCheck(cudaEventRecord(start, nullptr));
|
||||
kernel(std::forward<KernelArgs>(kernel_args)...);
|
||||
cudaCheck(cudaEventRecord(stop, nullptr));
|
||||
cudaCheck(cudaEventSynchronize(start));
|
||||
cudaCheck(cudaEventSynchronize(stop));
|
||||
float single_call;
|
||||
cudaCheck(cudaEventElapsedTime(&single_call, start, stop));
|
||||
elapsed_time += single_call;
|
||||
}
|
||||
cudaCheck(cudaEventRecord(stop, nullptr));
|
||||
cudaCheck(cudaEventSynchronize(start));
|
||||
cudaCheck(cudaEventSynchronize(stop));
|
||||
float elapsed_time;
|
||||
cudaCheck(cudaEventElapsedTime(&elapsed_time, start, stop));
|
||||
|
||||
cudaCheck(cudaFree(flush_buffer));
|
||||
|
||||
return elapsed_time / repeats;
|
||||
}
|
||||
|
|
@ -9,10 +9,10 @@ If encountering "error: identifier "M_PI" is undefined", add the following lines
|
|||
#define _USE_MATH_DEFINES
|
||||
#include <math.h> OR #include <cmath>
|
||||
|
||||
version 1 is naive port from CPU code to kernel
|
||||
version 1 is naive CPU port
|
||||
./gelu_forward 1
|
||||
|
||||
version 2 uses the Packed128 data structure
|
||||
version 2 is bfloat16 with the Packed128 data structure
|
||||
./gelu_forward 2
|
||||
*/
|
||||
|
||||
|
|
@ -21,6 +21,22 @@ version 2 uses the Packed128 data structure
|
|||
#include <cuda_runtime.h>
|
||||
#include "common.h"
|
||||
|
||||
// turn on bf16 as default, done up here for now
|
||||
#define ENABLE_BF16
|
||||
|
||||
#if defined(ENABLE_BF16)
|
||||
typedef __nv_bfloat16 floatX;
|
||||
typedef __nv_bfloat16 floatN;
|
||||
#elif defined(ENABLE_FP16)
|
||||
typedef half floatX;
|
||||
typedef half floatN;
|
||||
#else
|
||||
typedef float floatX;
|
||||
typedef float floatN;
|
||||
#endif
|
||||
|
||||
typedef Packed128<floatX> x128;
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CPU code reference
|
||||
|
||||
|
|
@ -38,7 +54,7 @@ void gelu_forward_cpu(float* out, const float* inp, int N) {
|
|||
// GPU kernels
|
||||
|
||||
// elementwise ops are nice and ez
|
||||
__global__ void gelu_kernel(float* out, const float* inp, int N) {
|
||||
__global__ void gelu_forward_kernel1(floatX* out, const floatX* inp, int N) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < N) {
|
||||
float xi = inp[i];
|
||||
|
|
@ -48,41 +64,41 @@ __global__ void gelu_kernel(float* out, const float* inp, int N) {
|
|||
}
|
||||
|
||||
// elementwise ops are nice and ez
|
||||
__global__ void gelu_kernel2(float* out, const float* inp, int N) {
|
||||
int i = (blockIdx.x * blockDim.x + threadIdx.x) * f128::size;
|
||||
__global__ void gelu_forward_kernel2(floatX* out, const floatX* inp, int N) {
|
||||
int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
|
||||
if (i < N) {
|
||||
f128 packet_out;
|
||||
f128 packet_in = load128cs(inp + i); // load and do not keep in cache
|
||||
for(int k = 0; k < packet_in.size; ++k) {
|
||||
float xi = packet_in[k];
|
||||
x128 packed_out;
|
||||
x128 packed_inp = load128cs(inp + i); // load and do not keep in cache
|
||||
for(int k = 0; k < packed_inp.size; ++k) {
|
||||
float xi = (float)packed_inp[k];
|
||||
float cube = 0.044715f * xi * xi * xi;
|
||||
packet_out[k] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
|
||||
packed_out[k] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))));
|
||||
}
|
||||
// store instead of storecs (without cache streaming) in case it is useful for the
|
||||
// data to be in the cache for the next operation after this GeLU
|
||||
store128(out + i, packet_out);
|
||||
store128(out + i, packed_out);
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// kernel launcher
|
||||
|
||||
void gelu_forward1(float* out, const float* inp, int N, const int block_size) {
|
||||
void gelu_forward1(floatX* out, const floatX* inp, int N, const int block_size) {
|
||||
const int grid_size = ceil_div(N, block_size);
|
||||
gelu_kernel<<<grid_size, block_size>>>(out, inp, N);
|
||||
gelu_forward_kernel1<<<grid_size, block_size>>>(out, inp, N);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void gelu_forward2(float* out, const float* inp, int N, const int block_size) {
|
||||
const int grid_size = ceil_div(N, block_size) / 4;
|
||||
gelu_kernel2<<<grid_size, block_size>>>(out, inp, N);
|
||||
void gelu_forward2(floatX* out, const floatX* inp, int N, const int block_size) {
|
||||
const int grid_size = ceil_div(N, block_size * x128::size);
|
||||
gelu_forward_kernel2<<<grid_size, block_size>>>(out, inp, N);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
// kernel version dispatch
|
||||
void gelu_forward(int kernel_num,
|
||||
float* out,
|
||||
const float* inp,
|
||||
floatX* out,
|
||||
const floatX* inp,
|
||||
int B, int T, int C,
|
||||
int block_size) {
|
||||
switch (kernel_num) {
|
||||
|
|
@ -100,7 +116,7 @@ void gelu_forward(int kernel_num,
|
|||
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
int main(int argc, const char **argv) {
|
||||
srand(0);
|
||||
|
||||
int B = 8;
|
||||
|
|
@ -114,13 +130,6 @@ int main(int argc, char **argv) {
|
|||
float* out = (float*)malloc(B * T * C * sizeof(float));
|
||||
float* inp = make_random_float(B * T * C);
|
||||
|
||||
// move to GPU
|
||||
float* d_out;
|
||||
float* d_inp;
|
||||
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float)));
|
||||
cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(float)));
|
||||
cudaCheck(cudaMemcpy(d_inp, inp, B * T * C * sizeof(float), cudaMemcpyHostToDevice));
|
||||
|
||||
// read kernel_num from command line
|
||||
int kernel_num = 1;
|
||||
if (argc > 1) {
|
||||
|
|
@ -131,6 +140,19 @@ int main(int argc, char **argv) {
|
|||
// first check the correctness of the kernel
|
||||
gelu_forward_cpu(out, inp, B * T * C);
|
||||
|
||||
// move to GPU
|
||||
floatX* d_out;
|
||||
floatX* d_inp;
|
||||
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(floatX)));
|
||||
cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX)));
|
||||
|
||||
floatX* inpX = (floatX*)malloc(B * T * C * sizeof(floatX));
|
||||
|
||||
for (int i = 0; i < B * T * C; i++) {
|
||||
inpX[i] = (floatX)inp[i];
|
||||
}
|
||||
|
||||
cudaCheck(cudaMemcpy(d_inp, inpX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice));
|
||||
|
||||
// time the kernel at different block sizes
|
||||
int block_sizes[] = {32, 64, 128, 256, 512, 1024};
|
||||
|
|
@ -138,7 +160,12 @@ int main(int argc, char **argv) {
|
|||
int block_size = block_sizes[j];
|
||||
printf("Checking block size %d.\n", block_size);
|
||||
gelu_forward(kernel_num, d_out, d_inp, B, T, C, block_size);
|
||||
validate_result(d_out, out, "out", B * T * C, 1e-5f);
|
||||
#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16)
|
||||
float tol = 1e-5;
|
||||
#else
|
||||
float tol = 1e-2f;
|
||||
#endif
|
||||
validate_result(d_out, out, "out", B * T * C, tol);
|
||||
}
|
||||
|
||||
printf("All results match. Starting benchmarks.\n\n");
|
||||
|
|
@ -155,7 +182,7 @@ int main(int argc, char **argv) {
|
|||
// napkin math: estimate the memory bandwidth achieved
|
||||
// for each (B,T,C) output element, we do 1 read and 1 write, 4 bytes each
|
||||
// and e.g. A100 40GB PCIe is advertised at 1,555GB/s
|
||||
long memory_ops = B * T * C * 2 * 4;
|
||||
long memory_ops = B * T * C * 2 * (int)sizeof(floatX);
|
||||
float memory_bandwidth = memory_ops / elapsed_time / 1e6;
|
||||
|
||||
printf("block_size %4d | time %.4f ms | bandwidth %.2f GB/s\n", block_size, elapsed_time, memory_bandwidth);
|
||||
|
|
@ -164,8 +191,9 @@ int main(int argc, char **argv) {
|
|||
// free memory
|
||||
free(out);
|
||||
free(inp);
|
||||
free(inpX);
|
||||
|
||||
cudaCheck(cudaFree(d_out));
|
||||
cudaCheck(cudaFree(d_inp));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -115,7 +115,7 @@ int main(int argc, char *argv[]) {
|
|||
|
||||
// build the GPT-2 model from a checkpoint
|
||||
GPT2 model;
|
||||
gpt2_build_from_checkpoint(&model, "gpt2_124M_bf16.bin");
|
||||
gpt2_build_from_checkpoint(&model, load_filename);
|
||||
size_t V = model.config.vocab_size;
|
||||
size_t Vp = model.config.padded_vocab_size;
|
||||
size_t maxT = model.config.max_seq_len;
|
||||
|
|
@ -286,7 +286,7 @@ int main(int argc, char *argv[]) {
|
|||
gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.01f, step+1);
|
||||
|
||||
// print the timing information at the end
|
||||
printf("step %d: loss %f (took %f ms)\n", step, model.mean_loss, time_elapsed_s * 1000);
|
||||
printf("step %d: loss %f (took %f ms)\n", step+1, model.mean_loss, time_elapsed_s * 1000);
|
||||
losses[step] = model.mean_loss;
|
||||
}
|
||||
|
||||
|
|
@ -307,10 +307,10 @@ int main(int argc, char *argv[]) {
|
|||
// compare
|
||||
for (int i = 0; i < 10; i++) {
|
||||
if (fabsf(losses[i] - expected_losses[i]) >= loss_diff_threshold) {
|
||||
printf("LOSS MISMATCH AT STEP %d: %f %f\n", i, losses[i], expected_losses[i]);
|
||||
printf("LOSS MISMATCH AT STEP %d: %f %f\n", i+1, losses[i], expected_losses[i]);
|
||||
allok = 0;
|
||||
} else {
|
||||
printf("loss ok at step %d: %f %f\n", i, losses[i], expected_losses[i]);
|
||||
printf("loss ok at step %d: %f %f\n", i+1, losses[i], expected_losses[i]);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -24,6 +24,13 @@ Also we're using TinyStories here for example as it is a bigger dataset
|
|||
|
||||
Example launch using bfloat16 on 4 GPUs, same as above:
|
||||
mpirun -np 4 ./train_gpt2cu -b 8 -v 200 -s 200 -i data/TinyStories
|
||||
|
||||
If you'd like to see train_gpt2.cu produce identical results to
|
||||
`python train_gpt2.py`, you can run it like this:
|
||||
make train_gpt2cu PRECISION=FP32
|
||||
./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -f 0
|
||||
This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200),
|
||||
-a 1 is "overfit single batch", -x 10 is 10 iterations, and -f 0 disables tf32
|
||||
*/
|
||||
#define ENABLE_CUDNN // can be enabled via nvcc "-DENABLE_CUDNN"
|
||||
|
||||
|
|
@ -229,9 +236,6 @@ struct alignas(16) Packed128 {
|
|||
__device__ const ElementType& operator[](int index) const {
|
||||
return payload[index];
|
||||
}
|
||||
__device__ float fp32(int index) {
|
||||
return static_cast<float>(payload[index]);
|
||||
}
|
||||
__device__ int4 get_bits() const {
|
||||
int4 bits;
|
||||
static_assert(sizeof(bits) == sizeof(payload), "Size mismatch.");
|
||||
|
|
@ -245,6 +249,7 @@ struct alignas(16) Packed128 {
|
|||
|
||||
// short-form typedef
|
||||
typedef Packed128<float> f128;
|
||||
typedef Packed128<floatX> x128;
|
||||
|
||||
// load a Packed128 from an aligned memory address
|
||||
template<class ElementType>
|
||||
|
|
@ -839,7 +844,8 @@ __global__ void permute_kernel_backward(floatX* dinp,
|
|||
|
||||
__global__ void unpermute_kernel(floatX* inp, floatX *out, int B, int N, int NH, int d) {
|
||||
// out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int idx = (blockIdx.x * blockDim.x + threadIdx.x);
|
||||
// out[b][n][nh_][d_] <- inp[b][nh_][n][d_]
|
||||
if (idx < B * NH * N * d) {
|
||||
int b = idx / (NH * N * d);
|
||||
|
|
@ -940,12 +946,19 @@ __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2,
|
|||
}
|
||||
|
||||
#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
|
||||
__global__ void gelu_forward_kernel(floatX* out, const floatX* inp, int N) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
__global__ void gelu_forward_kernel2(floatX* out, const floatX* inp, int N) {
|
||||
int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
|
||||
if (i < N) {
|
||||
float xi = (float)inp[i];
|
||||
float cube = 0.044715f * xi * xi * xi;
|
||||
out[i] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))));
|
||||
x128 packed_out;
|
||||
x128 packed_inp = load128cs(inp + i); // load and do not keep in cache
|
||||
for(int k = 0; k < packed_inp.size; ++k) {
|
||||
float xi = (float)packed_inp[k];
|
||||
float cube = 0.044715f * xi * xi * xi;
|
||||
packed_out[k] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))));
|
||||
}
|
||||
// store instead of storecs (without cache streaming) in case it is useful for the
|
||||
// data to be in the cache for the next operation after this GeLU
|
||||
store128(out + i, packed_out);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1464,9 +1477,9 @@ void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) {
|
|||
}
|
||||
|
||||
void gelu_forward(floatX* out, const floatX* inp, int N) {
|
||||
const int block_size = 128;
|
||||
const int grid_size = CEIL_DIV(N, block_size);
|
||||
gelu_forward_kernel<<<grid_size, block_size>>>(out, inp, N);
|
||||
const int block_size = 512;
|
||||
const int grid_size = CEIL_DIV(N, block_size * x128::size);
|
||||
gelu_forward_kernel2<<<grid_size, block_size>>>(out, inp, N);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
|
|
@ -2289,6 +2302,8 @@ void dataloader_init(DataLoader *loader, const MultiGpuConfig* multi_gpu_config,
|
|||
cudaMallocHost((void**)&loader->batch, (B * T + 1) * sizeof(int));
|
||||
loader->inputs = loader->batch;
|
||||
loader->targets = loader->batch + 1; // targets are shifted by one
|
||||
// note: we definitely want to advance by B * T; That is the "stride" by which we move
|
||||
// the window of tokens. We only load B * T + 1 tokens because our targets are offset by 1
|
||||
loader->num_batches = loader->file_size / (loader->num_processes * B * T * sizeof(int));
|
||||
}
|
||||
|
||||
|
|
@ -2307,6 +2322,7 @@ void dataloader_next_batch(DataLoader *loader) {
|
|||
fseekCheck(loader->tokens_file, loader->current_position, SEEK_SET);
|
||||
freadCheck(loader->batch, sizeof(int), B*T+1, loader->tokens_file);
|
||||
// advance the current position by B*T*num_processes integers
|
||||
// note: the "stride" of tokens by which we move each time is definitely B * T
|
||||
loader->current_position += loader->num_processes * B * T * sizeof(int);
|
||||
}
|
||||
|
||||
|
|
@ -2382,10 +2398,13 @@ void error_usage() {
|
|||
fprintf(stderr, " -b <int> batch size B (default = 4)\n");
|
||||
fprintf(stderr, " -t <int> sequence length T (default = 1024)\n");
|
||||
fprintf(stderr, " -l <float> learning rate (default = 3e-4f)\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");
|
||||
fprintf(stderr, " -s <int> sample_every, how often we inference the model (default = 20)\n");
|
||||
fprintf(stderr, " -g <int> genT, how many steps of inference we do (default = 64)\n");
|
||||
fprintf(stderr, " -a <int> overfit a single batch? 0/1. useful for debugging\n");
|
||||
fprintf(stderr, " -f <int> enable_tf32 override (default: 1, set to 0 to disable tf32)\n");
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
|
||||
|
|
@ -2404,6 +2423,9 @@ int main(int argc, char *argv[]) {
|
|||
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?
|
||||
int genT = 64; // number of steps of inference we will do
|
||||
int overfit_single_batch = 0; // useful for debugging, 1 = only load a single data batch once
|
||||
int max_steps = -1;
|
||||
int override_enable_tf32 = 1;
|
||||
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
|
||||
|
|
@ -2414,10 +2436,13 @@ int main(int argc, char *argv[]) {
|
|||
else if (argv[i][1] == 'b') { B = atoi(argv[i+1]); } // Per-GPU batch size
|
||||
else if (argv[i][1] == 't') { T = atoi(argv[i+1]); }
|
||||
else if (argv[i][1] == 'l') { learning_rate = 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]); }
|
||||
else if (argv[i][1] == 's') { sample_every = atoi(argv[i+1]); }
|
||||
else if (argv[i][1] == 'g') { genT = atoi(argv[i+1]); }
|
||||
else if (argv[i][1] == 'a') { overfit_single_batch = atoi(argv[i+1]); }
|
||||
else if (argv[i][1] == 'f') { override_enable_tf32 = atoi(argv[i+1]); }
|
||||
else { error_usage(); }
|
||||
}
|
||||
printf0("+-----------------------+----------------------------------------------------+\n");
|
||||
|
|
@ -2427,11 +2452,13 @@ int main(int argc, char *argv[]) {
|
|||
printf0("| output log file | %-50s |\n", output_log_file == NULL ? "NULL" : output_log_file);
|
||||
printf0("| batch size B | %-50d |\n", B);
|
||||
printf0("| sequence length T | %-50d |\n", T);
|
||||
printf0("| learning rate | %-50f |\n", learning_rate);
|
||||
printf0("| learning rate | %-50e |\n", learning_rate);
|
||||
printf0("| max_steps | %-50d |\n", max_steps);
|
||||
printf0("| val_loss_every | %-50d |\n", val_loss_every);
|
||||
printf0("| val_max_batches | %-50d |\n", val_max_batches);
|
||||
printf0("| sample_every | %-50d |\n", sample_every);
|
||||
printf0("| genT | %-50d |\n", genT);
|
||||
printf0("| overfit_single_batch | %-50d |\n", overfit_single_batch);
|
||||
printf0("+-----------------------+----------------------------------------------------+\n");
|
||||
|
||||
// set up the device
|
||||
|
|
@ -2449,6 +2476,7 @@ int main(int argc, char *argv[]) {
|
|||
|
||||
// TF32 precision is equivalent to torch.set_float32_matmul_precision('high')
|
||||
int enable_tf32 = cuda_arch_major >= 8 ? 1 : 0;
|
||||
if (override_enable_tf32 == 0) { enable_tf32 = 0; } // force to zero via arg
|
||||
cublas_compute_type = enable_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F;
|
||||
cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH;
|
||||
cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode));
|
||||
|
|
@ -2480,13 +2508,16 @@ int main(int argc, char *argv[]) {
|
|||
char train_tokens_filename[128];
|
||||
char val_tokens_filename[128];
|
||||
assert(strlen(input_dataset_prefix) < 100); // being bit lazy here, make sure we don't overflow
|
||||
sprintf(train_tokens_filename, "%s_train.bin", input_dataset_prefix);
|
||||
// if we're only overfitting a single batch for debugging, let's overfit the first batch
|
||||
// from val instead of train split, because val is smaller and a bit faster
|
||||
const char* train_split = (overfit_single_batch == 1) ? "val" : "train";
|
||||
sprintf(train_tokens_filename, "%s_%s.bin", input_dataset_prefix, train_split);
|
||||
sprintf(val_tokens_filename, "%s_val.bin", input_dataset_prefix);
|
||||
DataLoader train_loader;
|
||||
dataloader_init(&train_loader, &multi_gpu_config, train_tokens_filename, B, T);
|
||||
DataLoader val_loader;
|
||||
dataloader_init(&val_loader, &multi_gpu_config, val_tokens_filename, B, T);
|
||||
int train_num_batches = train_loader.num_batches; // let's do 1 epoch by default for now
|
||||
int train_num_batches = (max_steps == -1) ? train_loader.num_batches : max_steps; // default = 1 epoch
|
||||
int val_num_batches = train_loader.num_batches < val_max_batches ? train_loader.num_batches : val_max_batches;
|
||||
printf0("| train_num_batches | %-50d |\n", train_num_batches);
|
||||
printf0("| val_num_batches | %-50d |\n", val_num_batches);
|
||||
|
|
@ -2586,7 +2617,10 @@ int main(int argc, char *argv[]) {
|
|||
|
||||
// do a training step
|
||||
clock_gettime(CLOCK_MONOTONIC, &start);
|
||||
dataloader_next_batch(&train_loader);
|
||||
if (overfit_single_batch == 0 || (step == 0 && overfit_single_batch == 1)) {
|
||||
// if we're overfitting a single batch, we'll only call this at step = 0
|
||||
dataloader_next_batch(&train_loader);
|
||||
}
|
||||
gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T);
|
||||
gpt2_zero_grad(&model);
|
||||
gpt2_backward(&model);
|
||||
|
|
|
|||
|
|
@ -544,10 +544,10 @@ if __name__ == "__main__":
|
|||
with ctx:
|
||||
logits, loss = model(x, y)
|
||||
del logits
|
||||
if not args.inference_only:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
if not args.inference_only:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
# wait on the CPU for all device work to end so we get accurate per-iteration timings below
|
||||
if device == "mps":
|
||||
torch.mps.synchronize()
|
||||
|
|
@ -557,7 +557,7 @@ if __name__ == "__main__":
|
|||
t1 = time.time()
|
||||
# the 0th iteration is often an outlier (much slower) => skip logging it
|
||||
tokens_per_second = ddp_world_size * B * T / (t1-t0)
|
||||
print0(f"iteration {i}, loss: {loss.item():.4f}, time: {(t1-t0)*1000:.3f}ms, tok/s: {tokens_per_second:.2f}")
|
||||
print0(f"iteration {i+1}, loss: {loss.item():.4f}, time: {(t1-t0)*1000:.3f}ms, tok/s: {tokens_per_second:.2f}")
|
||||
if i > 0 and i > args.num_iterations - 20:
|
||||
timings.append(t1-t0)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue