diff --git a/profile_gpt2.cu b/profile_gpt2.cu index eab7fc5..fd1d78d 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -28,44 +28,7 @@ the profile.ncu-rep from a cloud box to local to pretty view. #include "train_gpt2.cu" int main() { - - // set up the device - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, deviceIdx); - printf("[System]\n"); - printf("Device %d: %s\n", deviceIdx, deviceProp.name); - - cuda_num_SMs = deviceProp.multiProcessorCount; - cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; - cuda_arch_major = deviceProp.major; - cuda_arch_minor = deviceProp.minor; - - cudaCheck(cudaStreamCreate(&main_stream)); - cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); - cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); - for (int i = 0; i < num_parallel_streams; i++) { - cudaCheck(cudaStreamCreate(¶llel_streams[i])); - cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); - } - - // setup cuBLAS and cuBLASLt - cublasCheck(cublasCreate(&cublas_handle)); - cublasCheck(cublasSetStream(cublas_handle, main_stream)); - cublasCheck(cublasLtCreate(&cublaslt_handle)); - // TF32 precision is equivalent to torch.set_float32_matmul_precision('high') - int enable_tf32 = deviceProp.major >= 8 ? 1 : 0; - printf("enable_tf32: %d\n", enable_tf32); - 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)); - // setup the (global) cuBLASLt workspace - cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - - #ifdef ENABLE_CUDNN - checkCudnnErr(cudnnCreate(&cudnn_handle)); - #endif + common_start(); // build the GPT-2 model from a checkpoint GPT2 model; @@ -91,16 +54,8 @@ int main() { gpt2_backward(&model); gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1); cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings + // free - gpt2_free(&model); - - #ifdef ENABLE_CUDNN - if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); - #endif - cudaCheck(cudaFree(cublaslt_workspace)); - cublasCheck(cublasDestroy(cublas_handle)); - cublasCheck(cublasLtDestroy(cublaslt_handle)); - + common_free(model); return 0; } diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index b3eec86..8e15b7d 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -31,7 +31,7 @@ metrics = [ "dram__bytes_write.sum", # DRAM writes "lts__t_sectors_srcunit_tex_op_read.sum", # L2 reads (sectors -- 32B) "lts__t_sectors_srcunit_tex_op_write.sum", # L2 reads (sectors -- 32B) - "smsp__inst_executed.sum", # instructions + "sm__pipe_tensor_op_hmma_cycles_active.avg.pct_of_peak_sustained_active", # todo - tensor core % ] cmd = [NCU, "-i", "profile.ncu-rep", "--csv", "--page", "raw", "--metrics", ",".join(metrics)] result = subprocess.check_output(cmd, text=True).strip() @@ -55,11 +55,11 @@ print("Kernel calls:") for rid, row in enumerate(reader): if rid == 0: # headings - print(f"id pass {'name':<40} {'time':>8} {'RAM rd':>8} {'RAM wt':>8} {'L2 rd':>8} {'L2 wt':>8} {'inst':>8}") + print(f"id pass {'name':<70} {'time':>8} {'RAM BW':>8} {'RAM rd':>8} {'RAM wt':>8} {'L2 rd':>8} {'L2 wt':>8} {'inst':>8}") continue if rid == 1: # units - units = f" {'':<40} {'ms':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'MInst':>8}" + units = f" {'':<70} {'ms':>8} {'GB/s':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'MInst':>8}" print(units) print("." * len(units)) continue @@ -74,7 +74,7 @@ for rid, row in enumerate(reader): write = float(row[12]) l2_read = float(row[14]) l2_write = float(row[15]) - inst = float(row[16]) / 1e6 + inst = float(row[16]) kid = rid - 2 @@ -118,18 +118,21 @@ for rid, row in enumerate(reader): total['l2_write'] += l2_write total['inst'] += inst - print(f"{kid:02} {pass_name:4} {fn_name:<40} {time:8.2f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}") + dram_bw = (read + write) / (time / 1000.0); + + print(f"{kid:02} {pass_name:4} {fn_name:<70} {time:8.2f} {dram_bw:8.1f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}") total_time = total['time'] +total_dram_bw = (total['read'] + total['write']) / (total_time / 1000.0); print("." * len(units)) -print(f" {'Total':<40} {total['time']:8.2f} {total['read']:8.2f} {total['write']:8.2f} {total['l2_read']:8.2f} {total['l2_write']:8.2f} {total['inst']:8.2f}") +print(f" {'Total':<70} {total['time']:8.2f} {total_dram_bw:8.1f} {total['read']:8.2f} {total['write']:8.2f} {total['l2_read']:8.2f} {total['l2_write']:8.2f} {total['inst']:8.2f}") print() print("Kernel type summaries:") -print(f" {'name':<40} {'time':>6} {'frac':>6}") +print(f" {'name':<70} {'time':>6} {'frac':>6}") ordered = sorted(summaries.items(), key=lambda x: x[1], reverse=True) for entry, value in ordered: - print(f" {entry:<40} {value:6.2f} {100*value / total_time:6.2f}%") + print(f" {entry:<70} {value:6.2f} {100*value / total_time:6.2f}%") ts = total_time / 1000 diff --git a/test_gpt2.cu b/test_gpt2.cu index 9c98f56..3fc6b6f 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -83,44 +83,7 @@ float* float_cpu_malloc_and_point_parameters(FloatParameterTensors* params, size } int main(int argc, char *argv[]) { - - // set up the device - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, deviceIdx); - printf("[System]\n"); - printf("Device %d: %s\n", deviceIdx, deviceProp.name); - - cuda_num_SMs = deviceProp.multiProcessorCount; - cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; - cuda_arch_major = deviceProp.major; - cuda_arch_minor = deviceProp.minor; - - cudaCheck(cudaStreamCreate(&main_stream)); - cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); - cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); - for (int i = 0; i < num_parallel_streams; i++) { - cudaCheck(cudaStreamCreate(¶llel_streams[i])); - cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); - } - - // setup cuBLAS and cuBLASLt - cublasCheck(cublasCreate(&cublas_handle)); - cublasCheck(cublasSetStream(cublas_handle, main_stream)); - cublasCheck(cublasLtCreate(&cublaslt_handle)); - // TF32 precision is equivalent to torch.set_float32_matmul_precision('high') - int enable_tf32 = cuda_arch_major >= 8 ? 1 : 0; - enable_tf32 = 0; // NOTE: disable TF32 for testing!!! - printf("enable_tf32: %d\n", enable_tf32); - 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)); - cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - - #ifdef ENABLE_CUDNN - checkCudnnErr(cudnnCreate(&cudnn_handle)); - #endif + common_start(false); // build the GPT-2 model from a checkpoint GPT2 model; @@ -327,6 +290,7 @@ int main(int argc, char *argv[]) { printf("overall okay: %d\n", allok); // free everything + common_free(model); free(x); free(y); free(logits_cpu_raw); @@ -336,14 +300,5 @@ int main(int argc, char *argv[]) { free(expected_grads_memory); free(grads_memory_cpu); free(grads_memory_cpu_float); - gpt2_free(&model); - #ifdef ENABLE_CUDNN - if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); - #endif - cudaCheck(cudaFree(cublaslt_workspace)); - cublasCheck(cublasDestroy(cublas_handle)); - cublasCheck(cublasLtDestroy(cublaslt_handle)); - return 0; } diff --git a/train_gpt2.cu b/train_gpt2.cu index 971e9be..2a86fde 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -11,7 +11,7 @@ sure that those parts work out ok and that we do a += as necessary. E.g., the layernorms are connected to the residuals so we += in layernorm backward. In this file we are using Mixed Precision training, so different activations, -paramaters, grads and buffers may be kept at different precisions, to take +parameters, grads and buffers may be kept at different precisions, to take advantage of the fast low-precision hardware in the latest GPUs (bf16/fp16), and fp8 (coming soon^TM). @@ -33,26 +33,15 @@ 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 */ -#include - #include -#include #include -#include -#include -#include -#include -#include -#include -#include +#include // GPU / CUDA related -#include -#include #include +#include #include -#include #include - +#include // Multi-GPU related #ifdef MULTI_GPU #include @@ -73,20 +62,12 @@ enum PrecisionMode { PRECISION_BF16 }; -// Default Properties -typedef float floatN; -#define CUBLAS_LOWP_COMPUTE cublas_compute_type -#ifdef MULTI_GPU -const ncclDataType_t ncclFloatN = ncclFloat; -#endif - // Specific configurations based on the enabled precision #if defined(ENABLE_FP32) typedef float floatX; #define CUBLAS_LOWP CUDA_R_32F #define PRECISION_MODE PRECISION_FP32 const char* load_filename = "gpt2_124M.bin"; -const char* precision_mode_str = "fp32"; #ifdef MULTI_GPU const ncclDataType_t ncclFloatX = ncclFloat; #endif @@ -97,7 +78,6 @@ typedef half floatX; #define CUBLAS_LOWP CUDA_R_16F #define PRECISION_MODE PRECISION_FP16 const char* load_filename = "gpt2_124M.bin"; -const char* precision_mode_str = "fp16"; #ifdef MULTI_GPU const ncclDataType_t ncclFloatX = ncclHalf; #endif @@ -107,7 +87,6 @@ typedef __nv_bfloat16 floatX; #define CUBLAS_LOWP CUDA_R_16BF #define PRECISION_MODE PRECISION_BF16 const char* load_filename = "gpt2_124M_bf16.bin"; // bf16 weights specific filename -const char* precision_mode_str = "bf16"; #ifdef MULTI_GPU const ncclDataType_t ncclFloatX = ncclBfloat16; #endif @@ -121,11 +100,9 @@ namespace fe = cudnn_frontend; #else #define CUDNN_16BIT fe::DataType_t::HALF #endif - static cudnnHandle_t cudnn_handle; static size_t cudnn_workspace_size = 0; // dynamically allocated as needed (up to 256MiB!) static void* cudnn_workspace = NULL; -#define checkCudnnErr(err) assert((int)err == 0); #endif // ENABLE_CUDNN // ---------------------------------------------------------------------------- @@ -144,17 +121,14 @@ class NvtxRange { #define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__) // cuBLAS workspace. Hardcoding to 32MiB but only Hopper needs 32, for others 4 is OK -static size_t cublaslt_workspace_size = 32 * 1024 * 1024; -static void* cublaslt_workspace = NULL; -static cublasComputeType_t cublas_compute_type; -cublasHandle_t cublas_handle; +const size_t cublaslt_workspace_size = 32 * 1024 * 1024; +void* cublaslt_workspace = NULL; +cublasComputeType_t cublas_compute = CUBLAS_COMPUTE_32F; cublasLtHandle_t cublaslt_handle; -int cuda_arch_major = 0; -int cuda_arch_minor = 0; -int cuda_num_SMs = 0; // for persistent threads where we want 1 threadblock per SM -int cuda_threads_per_SM = 0; +cublasHandle_t cublas_handle; +cudaDeviceProp deviceProp; -// CUDA streams & events (note: non-timing events, use separate event for timing/profiling!) +// CUDA streams & events (note: non-timing events, use separate events for timing/profiling!) constexpr int num_parallel_streams = 2; // + 1 primary "main_stream" (+ default stream) cudaStream_t parallel_streams[num_parallel_streams]; cudaEvent_t parallel_events[num_parallel_streams]; @@ -168,8 +142,7 @@ cudaEvent_t loss_event; // to make sure fused_classifier has written the losses // CUDA error checking void cudaCheck(cudaError_t error, const char *file, int line) { if (error != cudaSuccess) { - printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, - cudaGetErrorString(error)); + printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, cudaGetErrorString(error)); exit(EXIT_FAILURE); } }; @@ -206,35 +179,6 @@ void mpi_check(int status, const char *file, int line) { #define mpiCheck(err) (mpi_check(err, __FILE__, __LINE__)) #endif -// GPU helper functions for atomicAdd on smaller than 32-bit types -#ifdef ENABLE_BF16 -__device__ void atomicAddX(__nv_bfloat16* addr, __nv_bfloat16 val) { - uintptr_t ptr_val = reinterpret_cast(addr); - __nv_bfloat162* ptr_bf16 = reinterpret_cast<__nv_bfloat162*>(ptr_val & ~uintptr_t(0x3)); - - // Prepare the value to add, setting the other half to zero - __nv_bfloat162 add_val = (ptr_val & 0x3) ? __halves2bfloat162(__ushort_as_bfloat16(0), val) - : __halves2bfloat162(val, __ushort_as_bfloat16(0)); - atomicAdd(ptr_bf16, add_val); -} -#endif - -#ifdef ENABLE_FP16 -__device__ void atomicAddX(half* addr, half val) { - uintptr_t ptr_val = reinterpret_cast(addr); - half2* ptr_fp16 = reinterpret_cast(ptr_val & ~uintptr_t(0x3)); - - // Prepare the value to add, setting the other half to zero - half2 add_val = (ptr_val & 0x3) ? __halves2half2(__ushort_as_half(0), val) - : __halves2half2(val, __ushort_as_half(0)); - atomicAdd(ptr_fp16, add_val); -} -#endif - -__device__ void atomicAddX(float* addr, float val) { - atomicAdd(addr, val); -} - // warp-level reduction for summing values __device__ float warpReduceSum(float val) { for (int offset = 16; offset > 0; offset /= 2) { @@ -242,7 +186,6 @@ __device__ float warpReduceSum(float val) { } return val; } - // warp-level reduction for finding the maximum value __device__ float warpReduceMax(float val) { for (int offset = 16; offset > 0; offset /= 2) { @@ -250,16 +193,6 @@ __device__ float warpReduceMax(float val) { } return val; } - -#if defined(ENABLE_BF16) || defined(ENABLE_FP16) -__device__ floatX warpReduceSum(floatX val) { - for (int offset = 16; offset > 0; offset /= 2) { - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - } - return val; -} -#endif - // requires all 32 threads in the warp to be active, but should work for any block size // uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes // the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end @@ -270,14 +203,13 @@ __device__ float blockReduce(float val, bool final_sync=false, float out_of_boun // two reductions of up to 1024 threads: // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) __shared__ float shared_val[32]; - int lane_id = threadIdx.x % 32; - int warp_id = threadIdx.x / 32; - int num_warps = blockDim.x / 32; + const int lane_id = threadIdx.x % 32; + const int warp_id = threadIdx.x / 32; + const int num_warps = blockDim.x / 32; float warp_val = warp_reduction(val); if (lane_id == 0) { shared_val[warp_id] = warp_val; } __syncthreads(); - // same strategy, now reduce across warps warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds; float block_val = warp_reduction(warp_val); @@ -287,7 +219,6 @@ __device__ float blockReduce(float val, bool final_sync=false, float out_of_boun return block_val; } - // ---------------------------------------------------------------------------- // Packed128 data structure, which forces the compiler to use 128-bit loads/stores // in GPUs that support (the LDG.128 and STS.128 instructions) @@ -296,12 +227,11 @@ __device__ float blockReduce(float val, bool final_sync=false, float out_of_boun template struct alignas(16) Packed128 { - __device__ Packed128() = default; + Packed128() = default; __device__ explicit Packed128(int4 bits) { static_assert(sizeof(bits) == sizeof(payload), "Size mismatch."); memcpy(&payload, &bits, sizeof(bits)); } - __device__ ElementType& operator[](int index) { return payload[index]; } @@ -314,39 +244,35 @@ struct alignas(16) Packed128 { memcpy(&bits, &payload, sizeof(bits)); return bits; } - static constexpr const size_t size = sizeof(int4) / sizeof(ElementType); ElementType payload[size]; }; -// short-form typedef -typedef Packed128 f128; -typedef Packed128 x128; - // load a Packed128 from an aligned memory address template __device__ Packed128 load128(const ElementType* address) { return Packed128{*reinterpret_cast(address)}; } - // load a Packed128 from an aligned memory address with streaming cache hint template __device__ Packed128 load128cs(const ElementType* address) { return Packed128{__ldcs(reinterpret_cast(address))}; } - // store a Packed128 to an aligned memory address template __device__ void store128(ElementType* target, Packed128 value) { *reinterpret_cast(target) = value.get_bits(); } - // store a Packed128 to an aligned memory address with streaming cache hint template __device__ void store128cs(ElementType* target, Packed128 value) { __stcs(reinterpret_cast(target), value.get_bits()); } +// short-form typedefs +typedef Packed128 f128; +typedef Packed128 x128; + // ---------------------------------------------------------------------------- // Random Number Generatiom @@ -387,25 +313,11 @@ __device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigne mangledBits ^= (mangledBits >> 17); return mangledBits; } -__device__ __host__ constexpr unsigned int Get1dNoiseUint(int positionX, unsigned int seed) -{ - return SquirrelNoise5(positionX, seed); -} __device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY, unsigned int seed) { constexpr int PRIME_NUMBER = 198491317; // Large prime number with non-boring bits return SquirrelNoise5(indexX + (PRIME_NUMBER * indexY), seed); } -__device__ __host__ constexpr float Get1dNoiseZeroToOne(int index, unsigned int seed) -{ - constexpr double ONE_OVER_MAX_UINT = (1.0 / (double) 0xFFFFFFFF); - return (float)(ONE_OVER_MAX_UINT * (double) SquirrelNoise5(index, seed)); -} -__device__ __host__ constexpr float Get2dNoiseZeroToOne(int indexX, int indexY, unsigned int seed) -{ - constexpr double ONE_OVER_MAX_UINT = (1.0 / (double) 0xFFFFFFFF); - return (float)(ONE_OVER_MAX_UINT * (double) Get2dNoiseUint(indexX, indexY, seed)); -} // stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift) __device__ __forceinline__ void stochastic_rounding(float in, __nv_bfloat16 *out, unsigned int seed) { @@ -564,20 +476,16 @@ auto lookup_cache_or_build_graph_fwd(Args... args) { .set_compute_data_type(fe::DataType_t::FLOAT); // QKV is (B, T, 3, NH, HS) which cuDNN can handle directly without an external permute - auto Q = graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") + auto Q = graph->tensor(fe::graph::Tensor_attributes().set_name("Q") .set_dim({B, H, T, HS}) .set_stride({3 * H * HS * T, HS, 3 * H * HS, 1})); - auto K = graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") + auto K = graph->tensor(fe::graph::Tensor_attributes().set_name("K") .set_dim({B, H, T, HS}) .set_stride({3 * H * HS * T, HS, 3 * H * HS, 1})); - auto V = graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") + auto V = graph->tensor(fe::graph::Tensor_attributes().set_name("V") .set_dim({B, H, T, HS}) .set_stride({3 * H * HS * T, HS, 3 * H * HS, 1})); - auto attn_scale = graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + auto attn_scale = graph->tensor(fe::graph::Tensor_attributes().set_name("attn_scale") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_is_pass_by_value(true) @@ -613,6 +521,7 @@ auto lookup_cache_or_build_graph_fwd(Args... args) { auto plans = graph->create_execution_plans({fe::HeurMode_t::A}); assert(graph->check_support(cudnn_handle).is_good()); assert(graph->build_plans(cudnn_handle).is_good()); + assert(graph->get_workspace_size() <= cudnn_workspace_size); // fwd shouldn't need workspace auto tuple = std::make_tuple(graph, Q, K, V, attn_scale, O, stats); user_maintained_cache_fwd.insert({key, tuple}); @@ -631,40 +540,32 @@ auto lookup_cache_or_build_graph_bwd(Args... args) { // (B, N, 3, NH, HS) // must come from inp (which means we also need to convert THAT to FP16) - auto Q = graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") + auto Q = graph->tensor(fe::graph::Tensor_attributes().set_name("Q") .set_dim({B, NH, T, HS}) .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); - auto K = graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") + auto K = graph->tensor(fe::graph::Tensor_attributes().set_name("K") .set_dim({B, NH, T, HS}) .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); - auto V = graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") + auto V = graph->tensor(fe::graph::Tensor_attributes().set_name("V") .set_dim({B, NH, T, HS}) .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); - auto O = graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") + auto O = graph->tensor(fe::graph::Tensor_attributes().set_name("O") .set_dim({B, NH, T, HS}) .set_stride({NH * HS * T, HS, NH * HS, 1})); - auto dO = graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") + auto dO = graph->tensor(fe::graph::Tensor_attributes().set_name("dO") .set_dim({B, NH, T, HS}) .set_stride({NH * HS * T, HS, NH * HS, 1})); - auto stats = graph->tensor(fe::graph::Tensor_attributes() - .set_name("stats") + auto stats = graph->tensor(fe::graph::Tensor_attributes().set_name("stats") .set_dim({B, NH, T, 1}) .set_stride({NH * T, T, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); - auto attn_scale = graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + auto attn_scale = graph->tensor(fe::graph::Tensor_attributes().set_name("attn_scale") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_is_pass_by_value(true) .set_data_type(fe::DataType_t::FLOAT)); - auto sdpa_backward_options = fe::graph::SDPA_backward_attributes() - .set_name("flash_attention_backward") + auto sdpa_backward_options = fe::graph::SDPA_backward_attributes().set_name("flash_attention_backward") .set_causal_mask(true) .set_attn_scale(attn_scale); @@ -688,6 +589,16 @@ auto lookup_cache_or_build_graph_bwd(Args... args) { assert(graph->check_support(cudnn_handle).is_good()); assert(graph->build_plans(cudnn_handle).is_good()); + // Reallocate the workspace if the required size is greater than the current workspace + // By default, cuDNN uses up to 256MiB of workspace, so we don't want to just allocate the maximum + if (graph->get_workspace_size() > cudnn_workspace_size) { + if (cudnn_workspace_size > 0) { + cudaCheck(cudaFree(cudnn_workspace)); + } + cudnn_workspace_size = graph->get_workspace_size(); + cudaCheck(cudaMalloc(&cudnn_workspace, cudnn_workspace_size)); + } + auto tuple = std::make_tuple(graph, Q, K, V, O, dO, stats, attn_scale, dQ, dK, dV); user_maintained_cache_bwd.insert({key, tuple}); return tuple; @@ -721,16 +632,6 @@ void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) variant_pack[softmax_stats] = stats; } - // Reallocate the workspace if the required size is greater than the current workspace - // By default, cuDNN uses up to 256MiB of workspace, so we don't want to just allocate the maximum - if (graph->get_workspace_size() > cudnn_workspace_size) { - if (cudnn_workspace_size > 0) { - cudaCheck(cudaFree(cudnn_workspace)); - } - cudnn_workspace_size = graph->get_workspace_size(); - cudaCheck(cudaMalloc(&cudnn_workspace, cudnn_workspace_size)); - } - // Execute graph assert(graph->execute(cudnn_handle, variant_pack, cudnn_workspace).is_good()); cudaCheck(cudaGetLastError()); @@ -765,16 +666,6 @@ void attention_backward_cudnn(floatX* dqkvr, {dQ, devPtrdQ}, {dK, devPtrdK}, {dV, devPtrdV}, {attn_scale, &attn_scale_cpu}}; - // Reallocate the workspace if the required size is greater than the current workspace - // By default, cuDNN uses up to 256MiB of workspace, so we don't want to just allocate the maximum - if (graph->get_workspace_size() > cudnn_workspace_size) { - if (cudnn_workspace_size > 0) { - cudaCheck(cudaFree(cudnn_workspace)); - } - cudnn_workspace_size = graph->get_workspace_size(); - cudaCheck(cudaMalloc(&cudnn_workspace, cudnn_workspace_size)); - } - // Execute graph assert(graph->execute(cudnn_handle, variant_pack, cudnn_workspace).is_good()); cudaCheck(cudaGetLastError()); @@ -789,51 +680,70 @@ __global__ void encoder_forward_kernel3(floatX* out, int B, int T, int C) { int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; int N = B * T * C; - if (idx < N) { - int bt = idx / C; - int b = bt / T; - int t = bt % T; - int c = idx % C; + if (idx >= N) { return; } - int ix = inp[b * T + t]; + int bt = idx / C; + int b = bt / T; + int t = bt % T; + int c = idx % C; - floatX* out_btc = out + b * T * C + t * C + c; - const floatX* wte_ix = wte + ix * C + c; - const floatX* wpe_tc = wpe + t * C + c; + int ix = inp[b * T + t]; - x128 packed_out; - x128 wte = load128cs(wte_ix); - x128 wpe = load128cs(wpe_tc); - #pragma unroll - for (int k = 0; k < wte.size; k++) { - packed_out[k] = (floatX)((float)wte[k] + (float)wpe[k]); - } - store128(out_btc, packed_out); + floatX* out_btc = out + b * T * C + t * C + c; + const floatX* wte_ix = wte + ix * C + c; + const floatX* wpe_tc = wpe + t * C + c; + + x128 packed_out; + x128 wte128 = load128cs(wte_ix); + x128 wpe128 = load128cs(wpe_tc); + for (int k = 0; k < x128::size; k++) { + packed_out[k] = (floatX)((float)wte128[k] + (float)wpe128[k]); } + store128(out_btc, packed_out); +} + +__device__ void atomicStochasticAdd(__nv_bfloat16* address, float val0, float val1, uint seed) { + float2 val = make_float2(val0, val1); + uint* address_as_uint = (uint*)address; + uint old = *address_as_uint, assumed; + uint random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); + do { + assumed = old; + float2 old_fp32 = __bfloat1622float2(*(__nv_bfloat162*)&old); + float2 new_fp32 = make_float2(old_fp32.x + val.x, old_fp32.y + val.y); + __nv_bfloat162 new_bf16; + stochastic_rounding(new_fp32.x, &new_bf16.x, random); + stochastic_rounding(new_fp32.y, &new_bf16.y, random >> 16); + old = atomicCAS(address_as_uint, assumed, *(uint*)&new_bf16); + } while (assumed != old); +} +__device__ void atomicStochasticAdd(float* address, float val0, float val1, uint seed) { + atomicAdd(address, val0); + atomicAdd(address + 1, val1); } -// really bad naive kernel with atomicAdd __global__ void encoder_backward_kernel(floatX* dwte, floatX* dwpe, const floatX* dout, const int* inp, - int B, int T, int C) { + int B, int T, int C, uint seed) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int N = B * T * C; + idx *= 2; // 2 elements per thread + if (idx >= N) { return; } - if (idx < N) { - int bt = idx / C; - int b = bt / T; - int t = bt % T; - int c = idx % C; + int bt = idx / C; + int b = bt / T; + int t = bt % T; + int c = idx % C; - int ix = inp[b * T + t]; + int ix = inp[b * T + t]; - const floatX* dout_btc = dout + b * T * C + t * C + c; - floatX* dwte_ix = dwte + ix * C + c; - floatX* dwpe_tc = dwpe + t * C + c; + const floatX* dout_btc = dout + b * T * C + t * C + c; + floatX* dwte_ix = dwte + ix * C + c; + floatX* dwpe_tc = dwpe + t * C + c; - atomicAddX(dwte_ix, (floatX)*dout_btc); - atomicAddX(dwpe_tc, (floatX)*dout_btc); - } + float2 dout_data = make_float2(dout_btc[0], dout_btc[1]); + atomicStochasticAdd(dwte_ix, dout_data.x, dout_data.y, seed); + atomicStochasticAdd(dwpe_tc, dout_data.x, dout_data.y, seed ^ 0xFFFFFFFF); } __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, @@ -891,38 +801,38 @@ __global__ void permute_kernel(floatX* q, floatX* k, floatX* v, // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d) // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d) int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * NH * N * d) { return; } + // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_] - if (idx < B * NH * N * d) { - int b = idx / (NH * N * d); - int rest = idx % (NH * N * d); - int nh_ = rest / (N * d); - rest = rest % (N * d); - int n = rest / d; - int d_ = rest % d; - int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; - q[idx] = __ldcs(&inp[inp_idx]); - k[idx] = __ldcs(&inp[inp_idx + NH * d]); - v[idx] = __ldcs(&inp[inp_idx + 2 * (NH * d)]); - } + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; + q[idx] = __ldcs(&inp[inp_idx]); + k[idx] = __ldcs(&inp[inp_idx + NH * d]); + v[idx] = __ldcs(&inp[inp_idx + 2 * (NH * d)]); } __global__ void permute_kernel_backward(floatX* dinp, const floatX* dq, const floatX* dk, const floatX* dv, int B, int N, int NH, int d) { int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < B * NH * N * d) { - int b = idx / (NH * N * d); - int rest = idx % (NH * N * d); - int nh_ = rest / (N * d); - rest = rest % (N * d); - int n = rest / d; - int d_ = rest % d; + if (idx >= B * NH * N * d) { return; } - int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; - dinp[inp_idx] = dq[idx]; - dinp[inp_idx + NH * d] = dk[idx]; - dinp[inp_idx + 2 * (NH * d)] = dv[idx]; - } + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + + int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; + dinp[inp_idx] = dq[idx]; + dinp[inp_idx + NH * d] = dk[idx]; + dinp[inp_idx + 2 * (NH * d)] = dv[idx]; } __global__ void unpermute_kernel(floatX* inp, floatX *out, int B, int N, int NH, int d) { @@ -930,30 +840,30 @@ __global__ void unpermute_kernel(floatX* inp, floatX *out, int B, int N, int NH, 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); - int rest = idx % (NH * N * d); - int nh_ = rest / (N * d); - rest = rest % (N * d); - int n = rest / d; - int d_ = rest % d; - int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; - out[other_idx] = __ldcs(&inp[idx]); - } + if (idx >= B * NH * N * d) { return; } + + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; + out[other_idx] = __ldcs(&inp[idx]); } __global__ void unpermute_kernel_backward(floatX* dinp, const floatX *dout, int B, int N, int NH, int d) { int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < B * NH * N * d) { - int b = idx / (NH * N * d); - int rest = idx % (NH * N * d); - int nh_ = rest / (N * d); - rest = rest % (N * d); - int n = rest / d; - int d_ = rest % d; - int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; - dinp[idx] = (floatX)dout[other_idx]; - } + if (idx >= B * NH * N * d) { return; } + + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; + dinp[idx] = (floatX)dout[other_idx]; } __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, const floatX* inp, int N, int T) { @@ -983,13 +893,13 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons const floatX* x = inp + idx * T; // not INF, so we don't get NaNs accidentally when subtracting two values. + const float FLT_MAX = 340282346638528859811704183484516925440.0f; // to avoid including float.h float maxval = -FLT_MAX; float sumval = 0.0f; const floatX* x_aligned = reinterpret_cast(__builtin_assume_aligned(x, 16)); for (int i = lane_id; i < pos_by_4; i += warp_size) { float regarray[4]; - #pragma unroll for (int k = 0; k < 4; ++k) { regarray[k] = (float)x_aligned[4*i + k]; } @@ -1026,67 +936,61 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2, int N) { int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; - if (idx < N) { - x128 packed_out; - x128 packed_inp1 = load128cs(inp1 + idx); - x128 packed_inp2 = load128cs(inp2 + idx); - #pragma unroll - for (int k = 0; k < packed_inp1.size; k++) { - packed_out[k] = (floatX)((float)packed_inp1[k] + (float)packed_inp2[k]); - } - store128(out + idx, packed_out); + if (idx >= N) { return; } + + x128 packed_out; + x128 packed_inp1 = load128cs(inp1 + idx); + x128 packed_inp2 = load128cs(inp2 + idx); + for (int k = 0; k < packed_inp1.size; k++) { + packed_out[k] = (floatX)((float)packed_inp1[k] + (float)packed_inp2[k]); } + store128(out + idx, packed_out); } #define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI) __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) { - 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); + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + if (idx >= N) { return; } + + x128 packed_out; + x128 packed_inp = load128cs(inp + idx); // 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 + idx, packed_out); } __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { - int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; - if (i < N) { - x128 packed_dinp; - x128 packed_inp = load128cs(inp + i); - x128 packed_dout = load128cs(dout + i); - for (int k = 0; k < packed_inp.size; ++k) { - float x = (float)packed_inp[k]; - float cube = 0.044715f * x * x * x; - float tanh_arg = GELU_SCALING_FACTOR * (x + cube); - float tanh_out = tanhf(tanh_arg); - float coshf_out = coshf(tanh_arg); - float sech_out = 1.0f / (coshf_out * coshf_out); - float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); - packed_dinp[k] = (floatX)(local_grad * (float)packed_dout[k]); - } + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + if (idx >= N) { return; } - store128(dinp + i, packed_dinp); + x128 packed_dinp; + x128 packed_inp = load128cs(inp + idx); + x128 packed_dout = load128cs(dout + idx); + for (int k = 0; k < packed_inp.size; ++k) { + float x = (float)packed_inp[k]; + float cube = 0.044715f * x * x * x; + float tanh_arg = GELU_SCALING_FACTOR * (x + cube); + float tanh_out = tanhf(tanh_arg); + float coshf_out = coshf(tanh_arg); + float sech_out = 1.0f / (coshf_out * coshf_out); + float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); + packed_dinp[k] = (floatX)(local_grad * (float)packed_dout[k]); } + store128(dinp + idx, packed_dinp); } __global__ void matmul_backward_bias_kernel6(float* dbias, const floatX* dout, int B, int T, int OC) { // note: this kernel reads in floatX, but it writes to float! // this is because we're using atomics, which are super slow in < fp32 precision on < H100 GPUs // so the trick is do fp32 atomics to a buffer, and then copy_and_cast the result to floatX + // (this also results in higher accuracy than doing doing accumulation directly in floatX) - // Each warp is responsible for 32 * "x128::size" = 256 OCs at BF16 (OC must be a multiple of 256!) - // Block size is 512 threads (16 warps) and we reduce those 16 values into 1 at the end - // blockDim.x is 32 --> single warp being responsible for those 256 OCs - // blockDim.y is 16 --> 16 parallel independent warps processing the same OCs for different BTs - // gridDim.x is OC / 256 --> each block processes 256 OCs - // grimDim.y is max(1, (cuda_num_SMs * cuda_threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU! + // see comments in matmul_backward() for an explanation of block/grid dimensions etc. const int block_size = 512; const int block_size_x = 32; const int block_size_y = block_size / block_size_x; // 16 @@ -1101,26 +1005,23 @@ __global__ void matmul_backward_bias_kernel6(float* dbias, const floatX* dout, i accumulators[k] = 0.0f; } int thread_id = threadIdx.y * block_size_x + threadIdx.x; - for (int i = thread_id; i < OC_per_warp; i += block_size) { - shared[i] = 0.0f; + for (int idx = thread_id; idx < OC_per_warp; idx += block_size) { + shared[idx] = 0.0f; } __syncthreads(); - for (int i = blockIdx.y*block_size_y + threadIdx.y; i < B * T; i += gridDim.y*block_size_y) { - x128 packed_dout = load128(dout + global_oc + i*OC); + for (int idx = blockIdx.y*block_size_y + threadIdx.y; idx < B * T; idx += gridDim.y*block_size_y) { + x128 packed_dout = load128(dout + global_oc + idx*OC); for (int k = 0; k < x128::size; k++) { - //printf("%d: %f + %f\n", oc, accumulators[k], (float)packed_dout[k]); accumulators[k] += (float)packed_dout[k]; } - //__syncthreads(); // keep block synchronised to maximise memory locality (?) } for (int k = 0; k < x128::size; k++) { atomicAdd(shared + local_oc + k, accumulators[k]); } __syncthreads(); if (threadIdx.y == 0) { - for (int i = threadIdx.x; i < OC_per_warp; i += block_size_x) { - //printf("%d => %f\n", i, shared[i]); - atomicAdd(dbias + i + blockIdx.x*OC_per_warp, shared[i]); + for (int idx = threadIdx.x; idx < OC_per_warp; idx += block_size_x) { + atomicAdd(dbias + idx + blockIdx.x*OC_per_warp, shared[idx]); } } } @@ -1140,7 +1041,6 @@ __global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX float* dweight_shared = shared + C; // init shared memory to zero - #pragma unroll 4 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; @@ -1167,12 +1067,11 @@ __global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } - dnorm_mean = warpReduceSum(dnorm_mean); - dnorm_norm_mean = warpReduceSum(dnorm_norm_mean); - dnorm_mean = dnorm_mean / C; - dnorm_norm_mean = dnorm_norm_mean / C; + dnorm_mean = warpReduceSum(dnorm_mean) / C; + dnorm_norm_mean = warpReduceSum(dnorm_norm_mean) / C; // now iterate again and accumulate all the gradients + // todo - use x128 for this loop to improve performance for (int i = warpThreadIdx; i < C; i += warpSize) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; @@ -1193,7 +1092,7 @@ __global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX // Accumulate into a FP32 scratchpad // BF16 atomics are potentially much slower... and this is more precise! - // todo - could potentially avoid the extra copy if floatX is FP32, fairly negligible though + // todo - could avoid the extra copy if floatX is FP32, fairly negligible though __syncthreads(); float* scratch_dbias = scratch; float* scratch_dweight = scratch + C; @@ -1221,9 +1120,9 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl constexpr const int BlockSize = 256; constexpr int T_per_block = 4; - int idx = blockIdx.y; // go through blocks in reverse order, so the slowest block starts first int t0 = T - 1 - T_per_block*blockIdx.x; + int idx = blockIdx.y; att += idx * T * T; datt += idx * T * T; @@ -1254,41 +1153,40 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl // Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation). // Reference: https://developer.nvidia.com/blog/lerp-faster-cuda -__device__ inline float lerp(float start, float end, float weight) { +__device__ float lerp(float start, float end, float weight) { return fma(weight, end, fma(-weight, start, start)); } -// Termplate type T instead of floatx template __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters, float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay, unsigned int seed) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= num_parameters) return; // guard + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_parameters) { return; } // guard // get the gradient, m, and v for this parameter - float grad = (float)grads_memory[i]; - float m = m_memory[i]; - float v = v_memory[i]; + float grad = (float)grads_memory[idx]; + float m = m_memory[idx]; + float v = v_memory[idx]; // update the first moment (momentum) m = lerp(grad, m, beta1); - m_memory[i] = m; + m_memory[idx] = m; // update the second moment (RMSprop) v = lerp(grad * grad, v, beta2); - v_memory[i] = v; + v_memory[idx] = v; m /= beta1_correction; // m_hat v /= beta2_correction; // v_hat // fetch the old value of this parameter as a float, from either source - float old_param = (master_params_memory != NULL) ? master_params_memory[i] : (float)params_memory[i]; + float old_param = (master_params_memory != NULL) ? master_params_memory[idx] : (float)params_memory[idx]; // update this parameter float param = old_param - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * old_param)); // update our low precision version of the parameters using stochastic rounding // this will be used in the next forward pass // TODO: simply doing `params_memory[i] = (floatX)param;` breaks everything (why?) unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); - stochastic_rounding(param, ¶ms_memory[i], random); + stochastic_rounding(param, ¶ms_memory[idx], random); // write the full, float version of the param into our master copy, if we maintain one // this will be used in the next update - if (master_params_memory != NULL) { master_params_memory[i] = param; } + if (master_params_memory != NULL) { master_params_memory[idx] = param; } } struct SoftmaxParams { @@ -1378,14 +1276,14 @@ __global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX* __global__ void copy_and_cast_kernel(float* dst, const floatX* src, size_t n) { // a small kernel to copy and cast, i.e. `dst <- (float) src` - const size_t i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { dst[i] = (float)src[i]; } + const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { dst[idx] = (float)src[idx]; } } __global__ void cast_and_add_kernel(floatX* dst, const float* src, size_t n) { // used only for matmul_backward_bias kernel, a little bit embarassing TODO delete later - const size_t i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { dst[i] += (floatX)src[i]; } // have to += because dbias is a paramater + const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { dst[idx] += (floatX)src[idx]; } // have to += because dbias is a paramater } // ---------------------------------------------------------------------------- @@ -1404,12 +1302,12 @@ void encoder_forward(floatX* out, void encoder_backward(floatX* dwte, floatX* dwpe, const floatX* dout, const int* inp, - int B, int T, int C) { + int B, int T, int C, uint seed) { NVTX_RANGE_FN(); const int N = B * T * C; const int block_size = 256; - const int grid_size = CEIL_DIV(N, block_size); - encoder_backward_kernel<<>>(dwte, dwpe, dout, inp, B, T, C); + const int grid_size = CEIL_DIV(N, block_size * 2); // each thread handles 2 elements + encoder_backward_kernel<<>>(dwte, dwpe, dout, inp, B, T, C, seed); cudaCheck(cudaGetLastError()); } @@ -1424,9 +1322,7 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, cudaCheck(cudaGetLastError()); } -// uses cuBLASLt to fuse the bias and gelu. does not work with OC = 50257 (last layer) // https://docs.nvidia.com/cuda/cublas/#cublasltmatmul -// https://github.com/NVIDIA/CUDALibrarySamples/blob/master/cuBLASLt/LtSgemm/sample_cublasLt_LtSgemm.cu void matmul_forward_cublaslt(floatX* out, floatX* inp, floatX* weight, floatX* bias, int B, int T, int C, int OC) { @@ -1439,13 +1335,8 @@ void matmul_forward_cublaslt(floatX* out, exit(EXIT_FAILURE); } - // FP16 alpha/beta need to be used if and only if CUBLAS_COMPUTE_16F + // these need to be in FP16 if and only if alpha/beta are CUBLAS_COMPUTE_16F const float alpha = 1.0f, beta = 0.0f; - const half alpha_fp16 = (half)alpha, beta_fp16 = (half)beta; - const void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? - (const void*)&alpha_fp16 : (const void*)α - const void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? - (const void*)&beta_fp16 : (const void*)β int returnedResults = 0; cublasLtMatmulDesc_t operationDesc; @@ -1459,16 +1350,12 @@ void matmul_forward_cublaslt(floatX* out, // create the operation descriptor cublasOperation_t opNoTranspose = CUBLAS_OP_N; cublasOperation_t opTranspose = CUBLAS_OP_T; - cublasLtEpilogue_t epilogueBias = CUBLASLT_EPILOGUE_BIAS; + cublasLtEpilogue_t epilogueBias = has_bias ? CUBLASLT_EPILOGUE_BIAS : CUBLASLT_EPILOGUE_DEFAULT; - cudaDataType_t scale_type = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? CUDA_R_16F : CUDA_R_32F; - cublasCheck(cublasLtMatmulDescCreate(&operationDesc, CUBLAS_LOWP_COMPUTE, scale_type)); + cublasCheck(cublasLtMatmulDescCreate(&operationDesc, cublas_compute, CUDA_R_32F)); // FP16 if CUBLAS_COMPUTE_16F cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opTranspose, sizeof(opTranspose))); cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opNoTranspose, sizeof(opNoTranspose))); - if(has_bias) { - cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogueBias, - sizeof(epilogueBias))); - } + cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogueBias, sizeof(epilogueBias))); cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias))); // define matrix layouts @@ -1480,8 +1367,7 @@ void matmul_forward_cublaslt(floatX* out, // create a preference handle with specified max workspace cublasCheck(cublasLtMatmulPreferenceCreate(&preference)); cublasCheck(cublasLtMatmulPreferenceSetAttribute(preference, - CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, - &cublaslt_workspace_size, sizeof(cublaslt_workspace_size))); + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &cublaslt_workspace_size, sizeof(cublaslt_workspace_size))); // find a suitable algorithm cublasCheck(cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc, @@ -1494,7 +1380,7 @@ void matmul_forward_cublaslt(floatX* out, // call the matmul cublasCheck(cublasLtMatmul(cublaslt_handle, operationDesc, - alpha_ptr, weight, weightLayout, inp, inputLayout, beta_ptr, + &alpha, weight, weightLayout, inp, inputLayout, &beta, out, outputLayout, out, outputLayout, &heuristic.algo, cublaslt_workspace, cublaslt_workspace_size, main_stream)); @@ -1514,7 +1400,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. // Its contents will be overwritten by this function. const int block_size = 256; - const int softmax_block_size = 256; + const float alpha = 1.0f, beta = 0.0f; // inp is (B, T, 3C) QKV // preatt, att are (B, NH, T, T) @@ -1529,50 +1415,32 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, int total_threads = B * NH * T * HS; int num_blocks = CEIL_DIV(total_threads, block_size); permute_kernel<<>>(q, k, v, inp, B, T, NH, HS); - cudaCheck(cudaGetLastError()); - // IMPORTANT: alpha/beta are FP32 for CUBLAS_COMPUTE_32F even if FP16 inputs/outputs - // But need FP16 scale for CUBLAS_COMPUTE_16F (no errors otherwise, just garbage results *sigh*) - const float alpha = 1.0f; - const float beta = 0.0f; - const floatX alpha_lowp = (floatX)alpha; - const floatX beta_lowp = (floatX)beta; - void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (void*)&alpha_lowp : (void*)α - void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (void*)&beta_lowp : (void*)β floatX* preatt = inp; cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, - T, T, HS, - alpha_ptr, + T, T, HS, &alpha, k, CUBLAS_LOWP, HS, T * HS, q, CUBLAS_LOWP, HS, T * HS, - beta_ptr, - preatt, CUBLAS_LOWP, T, T * T, - B * NH, - CUBLAS_LOWP_COMPUTE, - CUBLAS_GEMM_DEFAULT)); + &beta, preatt, CUBLAS_LOWP, T, T * T, + B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); // multiply all elements of preatt elementwise by scale float scale = 1.0 / sqrtf(HS); - int grid_size = CEIL_DIV(B * NH * T * 32, softmax_block_size); - softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); - cudaCheck(cudaGetLastError()); + int grid_size = CEIL_DIV(B * NH * T * 32, block_size); + softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); // new approach: first cuBLAS another batched matmul floatX* vaccum = inp; // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs) cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, - HS, T, T, - alpha_ptr, + HS, T, T, &alpha, v, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, - beta_ptr, - vaccum, CUBLAS_LOWP, HS, T * HS, - B * NH, - CUBLAS_LOWP_COMPUTE, - CUBLAS_GEMM_DEFAULT)); + &beta, vaccum, CUBLAS_LOWP, HS, T * HS, + B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); // now unpermute // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side @@ -1610,8 +1478,7 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* dbias_buffer, int B, int T, int C, int OC) { NVTX_RANGE_FN(); - float one = 1.0f; - float zero = 0.0f; + float one = 1.0f, zero = 0.0f; // backward to bias, if given, does a += if (dbias != NULL) { @@ -1620,14 +1487,15 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, // blockDim.x is 32 --> single warp being responsible for those 256 OCs // blockDim.y is 16 --> 16 parallel independent warps processing the same OCs for different BTs // gridDim.x is OC / 256 --> each block processes 256 OCs - // grimDim.y is max(1, (cuda_num_SMs * cuda_threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU! + // grimDim.y is max(1, (cuda_num_SMs * threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU! const int warp_size = 32; const int block_size = 512; const int OC_per_warp = warp_size * x128::size; // 256 at BF16 const int block_size_x = 32; const int block_size_y = block_size / block_size_x; // 16 const int grid_size_x = OC / OC_per_warp; // e.g. 3 horizontal blocks for 768 OCs at BF16 - const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / (block_size * grid_size_x)); // full GPU! + const int grid_size_y = max(1, deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount + / (block_size * grid_size_x)); // full GPU! assert((OC % OC_per_warp) == 0); // there is no bounds checking in the kernel to maximise performance @@ -1636,17 +1504,17 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, dim3(block_size_x, block_size_y), OC_per_warp * sizeof(float), main_stream>>>(dbias_buffer, dout, B, T, OC); cast_and_add_kernel<<>>(dbias, dbias_buffer, OC); - cudaCheck(cudaGetLastError()); } // backward to input, uses = in the backward pass (set the gradient) cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one, weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero, - dinp, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - // backward to weight, uses += in the backward pass (accumulate the gradient) + dinp, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + // backward to weight, uses += in the backward pass (accumulate the gradient) by setting alpha=one cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &one, - dweight, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + dweight, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + cudaCheck(cudaGetLastError()); } void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, @@ -1654,11 +1522,10 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr int B, int T, int C) { NVTX_RANGE_FN(); const int block_size = 1024; - const int grid_size = 1 * cuda_num_SMs; + const int grid_size = deviceProp.multiProcessorCount; size_t shared_mem_size = (2 * C + 1) * sizeof(float); cudaMemsetAsync(scratch, 0, (2 * C + 1) * sizeof(float), main_stream); - layernorm_backward_kernel7<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); cudaCheck(cudaGetLastError()); } @@ -1672,14 +1539,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da NVTX_RANGE_FN(); const int block_size = 256; int HS = C / NH; // head size - - // FP16 alpha/beta need to be used if and only if CUBLAS_COMPUTE_16F const float alpha = 1.0f, beta = 0.0f; - const half alpha_fp16 = (half)alpha, beta_fp16 = (half)beta; - const void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? - (const void*)&alpha_fp16 : (const void*)α - const void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? - (const void*)&beta_fp16 : (const void*)β // unpack convenience pointers into q, k, v const floatX *q, *k, *v; @@ -1694,31 +1554,26 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da // backward through the unpermute operation int num_blocks = CEIL_DIV(B * T * C, block_size); unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); - cudaCheck(cudaGetLastError()); // backward into datt - - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, alpha_ptr, - v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, beta_ptr, - datt, CUBLAS_LOWP, T, T * T, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); - + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, + v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, &beta, + datt, CUBLAS_LOWP, T, T * T, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); // backward into dv - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, alpha_ptr, - scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, beta_ptr, - dv, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); - + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, + scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, &beta, + dv, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); // backward into preatt int hs = C / NH; // head size float scale = 1.0f / sqrtf(hs); softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); - cudaCheck(cudaGetLastError()); // backward into q - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, alpha_ptr, - k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, beta_ptr, - dq, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, + k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, + dq, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); // backward into k - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, alpha_ptr, - q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, beta_ptr, - dk, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, + q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, + dk, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); // backward into inp num_blocks = CEIL_DIV(B * NH * T * HS, block_size); permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); @@ -1814,9 +1669,9 @@ void* malloc_and_point_parameters(ParameterTensors* params, size_t* param_elemen cudaCheck(cudaMalloc((void**)¶ms_memory, num_parameters_bytes)); // assign all the tensors their place in the array floatX** ptrs[] = { - ¶ms->wte, ¶ms->wpe, (floatX**)¶ms->ln1w, (floatX**)¶ms->ln1b, ¶ms->qkvw, ¶ms->qkvb, - ¶ms->attprojw, ¶ms->attprojb, (floatX**)¶ms->ln2w, (floatX**)¶ms->ln2b, ¶ms->fcw, ¶ms->fcb, - ¶ms->fcprojw, ¶ms->fcprojb, (floatX**)¶ms->lnfw, (floatX**)¶ms->lnfb + ¶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 }; char* params_memory_iterator = (char*)params_memory; for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { @@ -1826,7 +1681,7 @@ void* malloc_and_point_parameters(ParameterTensors* params, size_t* param_elemen return params_memory; } -#define NUM_ACTIVATION_TENSORS 21 +#define NUM_ACTIVATION_TENSORS 20 typedef struct { floatX* encoded; // (B, T, C) floatX* ln1; // (L, B, T, C) @@ -1846,7 +1701,6 @@ typedef struct { floatX* lnf; // (B, T, C) floatX* lnf_mean; // (B, T) floatX* lnf_rstd; // (B, T) - floatX* losses; // (B, T) // todo - no longer used as GPU writes directly to cpu_losses // adding these two compared to the CPU .c code, needed for attention kernel as buffers floatX* qkvr; // (L, B, T, 3*C) // in inference mode, this buffer will store the logits @@ -1885,9 +1739,8 @@ void fill_in_activation_sizes(size_t* act_sizes, size_t B, size_t T, GPT2Config act_sizes[15] = B * T * C; // lnf act_sizes[16] = B * T; // lnf_mean act_sizes[17] = B * T; // lnf_rstd - 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 + act_sizes[18] = L * B * T * 3*C; // qkvr + act_sizes[19] = B * T * max(3*C, max(NH*T, Vp)); // output / scratch } // Backward pass is conceptually quite different from forward, because we can discard @@ -1938,7 +1791,7 @@ void* malloc_and_point_activations(ActivationTensors* acts, const size_t* act_si &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 + &acts->lnf_mean, &acts->lnf_rstd, &acts->qkvr, &acts->output }; return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS); } @@ -2358,7 +2211,7 @@ void gpt2_backward(GPT2 *model) { // layernorm backward does += to dresidual, so it correctly accumulates gradient for the Attention block above layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); } - encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C); + encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C, random_u32(&model->rng_state)); // accumulate the loss, this was calculated at the end of gpt2_forward() cudaCheck(cudaEventSynchronize(loss_event)); // hopefully finished long ago @@ -2441,9 +2294,61 @@ void gpt2_free(GPT2 *model) { cudaFreeHost(model->cpu_losses); } +// ---------------------------------------------------------------------------- +// common init & free code for train/test/profile +void common_start(bool override_enable_tf32 = true) { + int deviceIdx = 0; + cudaCheck(cudaSetDevice(deviceIdx)); + cudaGetDeviceProperties(&deviceProp, deviceIdx); + printf("[System]\n"); + printf("Device %d: %s\n", deviceIdx, deviceProp.name); + + cudaCheck(cudaStreamCreate(&main_stream)); + cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); + cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); + for (int i = 0; i < num_parallel_streams; i++) { + cudaCheck(cudaStreamCreate(¶llel_streams[i])); + cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); + } + + // set up cuBLAS and cuBLASLt (and cuDNN if enabled) + cublasCheck(cublasCreate(&cublas_handle)); + cublasCheck(cublasSetStream(cublas_handle, main_stream)); + cublasCheck(cublasLtCreate(&cublaslt_handle)); + cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); + #ifdef ENABLE_CUDNN + checkCudnnErr(cudnnCreate(&cudnn_handle)); + #endif + // TF32 precision is equivalent to torch.set_float32_matmul_precision('high') + bool enable_tf32 = PRECISION_MODE == PRECISION_FP32 && deviceProp.major >= 8 && override_enable_tf32; + cublasCheck(cublasSetMathMode(cublas_handle, enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH)); + cublas_compute = enable_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F; + + // setup the (global) cuBLASLt workspace + cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); +} + +void common_free(GPT2 &model) { + cudaCheck(cudaEventDestroy(main_event)); + cudaCheck(cudaEventDestroy(loss_event)); + for (int i = 0; i < num_parallel_streams; i++) { + cudaCheck(cudaStreamDestroy(parallel_streams[i])); + cudaCheck(cudaEventDestroy(parallel_events[i])); + } + cudaCheck(cudaStreamDestroy(main_stream)); + + gpt2_free(&model); + #ifdef ENABLE_CUDNN + if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } + checkCudnnErr(cudnnDestroy(cudnn_handle)); + #endif + cudaCheck(cudaFree(cublaslt_workspace)); + cublasCheck(cublasDestroy(cublas_handle)); + cublasCheck(cublasLtDestroy(cublaslt_handle)); +} + #ifndef TESTING // if we are TESTING (see test_gpt2.cu), we'll skip the int main below - // ---------------------------------------------------------------------------- // data loader lite: returns random batches of data from a file of integers @@ -2656,45 +2561,14 @@ int main(int argc, char *argv[]) { printf0("| use_master_weights | %-50s |\n", use_master_weights ? "enabled" : "disabled"); printf0("+-----------------------+----------------------------------------------------+\n"); - // set up the device - cudaCheck(cudaSetDevice(multi_gpu_config.local_device_idx)); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, multi_gpu_config.local_device_idx); - cuda_num_SMs = deviceProp.multiProcessorCount; - cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; - cuda_arch_major = deviceProp.major; - cuda_arch_minor = deviceProp.minor; + common_start(override_enable_tf32); // common init code for train/test/profile - cudaCheck(cudaStreamCreate(&main_stream)); - cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); - cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); - for (int i = 0; i < num_parallel_streams; i++) { - cudaCheck(cudaStreamCreate(¶llel_streams[i])); - cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); - } - - // set up cuBLAS and cuBLASLt - cublasCheck(cublasCreate(&cublas_handle)); - cublasCheck(cublasSetStream(cublas_handle, main_stream)); - cublasCheck(cublasLtCreate(&cublaslt_handle)); - cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - // setup compute precision settings for cublas - // 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)); - if(cublas_compute_type); // unused in BF16 mode, avoid warning - - // set up cuDNN - #ifdef ENABLE_CUDNN - checkCudnnErr(cudnnCreate(&cudnn_handle)); - #endif + const char* precision_str = (PRECISION_MODE == PRECISION_FP32) + ? (cublas_compute == CUBLAS_COMPUTE_32F_FAST_TF32 ? "TF32" : "FP32") + : (PRECISION_MODE == PRECISION_FP16 ? "FP16" : "BF16"); printf0("| device | %-50s |\n", deviceProp.name); - printf0("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); - printf0("| precision | %-50s |\n", precision_mode_str); + printf0("| precision | %-50s |\n", precision_str); printf0("+-----------------------+----------------------------------------------------+\n"); // build the GPT-2 model from a checkpoint @@ -2712,17 +2586,15 @@ int main(int argc, char *argv[]) { printf0("+-----------------------+----------------------------------------------------+\n"); // build DataLoaders for both train and val - char train_tokens_filename[128]; - char val_tokens_filename[128]; + char train_tokens_filename[128], val_tokens_filename[128]; assert(strlen(input_dataset_prefix) < 100); // being bit lazy here, make sure we don't overflow // 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 train_loader, val_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 = (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; @@ -2738,11 +2610,9 @@ int main(int argc, char *argv[]) { printf0("num_parameters: %zu ==> bytes: %zu\n", model.num_parameters, model.num_parameters_bytes); printf0("allocated %d MiB for model parameters\n", (int)round(model.num_parameters_bytes / (1024 * 1024))); - // set up the Logger + // set up the Logger & Tokenizer Logger logger; logger_init(&logger, output_log_file); - - // build the Tokenizer Tokenizer tokenizer; tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); @@ -2870,20 +2740,13 @@ int main(int argc, char *argv[]) { dataloader_free(&train_loader); dataloader_free(&val_loader); tokenizer_free(&tokenizer); - gpt2_free(&model); free(cpu_logits_raw); free(cpu_logits); free(gen_tokens); - #ifdef ENABLE_CUDNN - if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); - #endif - cudaCheck(cudaFree(cublaslt_workspace)); - cublasCheck(cublasDestroy(cublas_handle)); - cublasCheck(cublasLtDestroy(cublaslt_handle)); logger_free(&logger); multi_gpu_config_free(&multi_gpu_config); + common_free(model); return 0; } #endif