Merge remote-tracking branch 'karpathy/master' into cudnn_try2

dev/cuda/README.md

@@ -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"
+```

dev/cuda/benchmark_on_modal.py

@@ -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

dev/cuda/common.h

@@ -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;
 }

dev/cuda/gelu_forward.cu

@@ -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;
 }