mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-27 20:25:09 -04:00
Merge branch 'stream-2' of https://github.com/ngc92/llm.c into ngc92-stream-2
This commit is contained in:
Andrej Karpathy
commit
8940dcdcb1
12 changed files with 170 additions and 119 deletions
|
|
@ -47,3 +47,19 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg
|
|||
// this will be used in the next update
|
||||
if (master_params_memory != NULL) { master_params_memory[idx] = param; }
|
||||
}
|
||||
|
||||
template <typename Tp, typename Tg>
|
||||
void adamw_update(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, int t, float eps, float weight_decay,
|
||||
float grad_scale, unsigned int seed, cudaStream_t stream) {
|
||||
// AdamW update
|
||||
int block_size = 512;
|
||||
int num_blocks = CEIL_DIV(num_parameters, block_size);
|
||||
float beta1_correction = 1.0f - powf(beta1, t);
|
||||
float beta2_correction = 1.0f - powf(beta2, t);
|
||||
adamw_kernel3<<<num_blocks, block_size, 0, stream>>>(params_memory, master_params_memory, grads_memory,
|
||||
m_memory, v_memory, num_parameters,
|
||||
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay,
|
||||
grad_scale, seed);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
@ -190,7 +190,7 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl
|
|||
|
||||
void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
||||
floatX* inp,
|
||||
int B, int T, int C, int NH) {
|
||||
int B, int T, int C, int NH, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
// 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.
|
||||
|
|
@ -209,9 +209,11 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
v = qkvr + 2 * B * T * C;
|
||||
int total_threads = B * NH * T * HS;
|
||||
int num_blocks = CEIL_DIV(total_threads, block_size);
|
||||
permute_kernel<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
|
||||
permute_kernel<<<num_blocks, block_size, 0, stream>>>(q, k, v, inp, B, T, NH, HS);
|
||||
|
||||
floatX* preatt = inp;
|
||||
cublasCheck(cublasSetStream(cublas_handle, stream));
|
||||
cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
|
||||
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
|
||||
CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
T, T, HS, &alpha,
|
||||
|
|
@ -223,7 +225,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
// multiply all elements of preatt elementwise by scale
|
||||
float scale = 1.0 / sqrtf(HS);
|
||||
int grid_size = CEIL_DIV(B * NH * T * WARP_SIZE, block_size);
|
||||
softmax_forward_kernel5<<<grid_size, block_size>>>(att, scale, preatt, B * NH, T);
|
||||
softmax_forward_kernel5<<<grid_size, block_size, 0, stream>>>(att, scale, preatt, B * NH, T);
|
||||
|
||||
// new approach: first cuBLAS another batched matmul
|
||||
floatX* vaccum = inp;
|
||||
|
|
@ -239,7 +241,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
// now unpermute
|
||||
// y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
|
||||
num_blocks = CEIL_DIV(B * T * C, block_size);
|
||||
unpermute_kernel<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
|
||||
unpermute_kernel<<<num_blocks, block_size, 0, stream>>>(vaccum, out, B, T, NH, HS);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
|
|
@ -248,7 +250,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* datt, floatX* scratch,
|
||||
const floatX* dout,
|
||||
const floatX* qkvr, const floatX* att,
|
||||
int B, int T, int C, int NH) {
|
||||
int B, int T, int C, int NH, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 256;
|
||||
int HS = C / NH; // head size
|
||||
|
|
@ -266,8 +268,10 @@ 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<<<num_blocks, block_size>>>(scratch, dout, B, T, NH, HS);
|
||||
unpermute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(scratch, dout, B, T, NH, HS);
|
||||
// backward into datt
|
||||
cublasCheck(cublasSetStream(cublas_handle, stream));
|
||||
cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
|
||||
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));
|
||||
|
|
@ -278,7 +282,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da
|
|||
// backward into preatt
|
||||
int hs = C / NH; // head size
|
||||
float scale = 1.0f / sqrtf(hs);
|
||||
softmax_autoregressive_backward_kernel<<<dim3(T / 4, B * NH), 256>>>(dpreatt, datt, att, B, T, C, scale);
|
||||
softmax_autoregressive_backward_kernel<<<dim3(T / 4, B * NH), 256, 0, stream>>>(dpreatt, datt, att, B, T, C, scale);
|
||||
// backward into q
|
||||
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,
|
||||
|
|
@ -289,6 +293,6 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da
|
|||
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<<<num_blocks, block_size>>>(dinp, dq, dk, dv, B, T, NH, HS);
|
||||
permute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(dinp, dq, dk, dv, B, T, NH, HS);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ Common utilities for CUDA code.
|
|||
#include <string>
|
||||
#include <cuda_runtime.h>
|
||||
#include <nvtx3/nvToolsExt.h>
|
||||
#include <nvtx3/nvToolsExtCudaRt.h>
|
||||
#include <cuda_profiler_api.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
|
|
|||
|
|
@ -21,9 +21,16 @@ static_assert(false, "cuDNN is not supported in FP32 mode.")
|
|||
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);
|
||||
|
||||
static void checkCudnnFE(fe::error_object e, const char *file, int line) {
|
||||
static void cuDNNCheck(cudnnStatus_t error, const char *file, int line) {
|
||||
if (error != CUDNN_STATUS_SUCCESS) {
|
||||
printf("[CUDNN ERROR] at file %s:%d:\n%s\n", file, line, cudnnGetErrorString(error));
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
};
|
||||
#define cuDNNCheck(err) (cuDNNCheck(err, __FILE__, __LINE__))
|
||||
|
||||
static void checkCudnnFE(const fe::error_object& e, const char *file, int line) {
|
||||
if(!e.is_good()) {
|
||||
printf("[CUDNN ERROR] at file %s:%d:\n%s\n", file, line, e.err_msg.c_str());
|
||||
exit(EXIT_FAILURE);
|
||||
|
|
@ -211,11 +218,13 @@ auto lookup_cache_or_build_graph_bwd(int B, int NH, int T, int HS) {
|
|||
void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
|
||||
float* stats, // output for backward pass: (B, NH, T)
|
||||
floatX* inp, // input: (B, T, 3, NH, HS) QKV
|
||||
int B, int T, int NH, int C) {
|
||||
int B, int T, int NH, int C, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
int HS = C / NH; // number of features per head
|
||||
bool is_inference_only = (stats == nullptr);
|
||||
|
||||
cuDNNCheck(cudnnSetStream(cudnn_handle, stream));
|
||||
|
||||
// Get graph and tensors from cache (or generate it on first use)
|
||||
auto graph = lookup_cache_or_build_graph_fwd(B, NH, T, HS, is_inference_only);
|
||||
|
||||
|
|
@ -242,7 +251,7 @@ void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
|
|||
|
||||
void attention_backward_cudnn(floatX* dqkvr, // output
|
||||
floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs
|
||||
int B, int T, int NH, int C) {
|
||||
int B, int T, int NH, int C, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
int HS = C / NH; // number of features per head
|
||||
|
||||
|
|
@ -269,15 +278,16 @@ void attention_backward_cudnn(floatX* dqkvr,
|
|||
{Attn_scale_UID, &attn_scale_cpu}};
|
||||
|
||||
// Execute graph
|
||||
cuDNNCheck(cudnnSetStream(cudnn_handle, stream));
|
||||
checkCudnnFE(graph->execute(cudnn_handle, variant_pack, cudnn_workspace));
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void create_cudnn() {
|
||||
checkCudnnErr(cudnnCreate(&cudnn_handle));
|
||||
cuDNNCheck(cudnnCreate(&cudnn_handle));
|
||||
}
|
||||
|
||||
void destroy_cudnn() {
|
||||
if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); }
|
||||
checkCudnnErr(cudnnDestroy(cudnn_handle));
|
||||
cuDNNCheck(cudnnDestroy(cudnn_handle));
|
||||
}
|
||||
|
|
@ -12,10 +12,10 @@ void destroy_cudnn();
|
|||
void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
|
||||
float* stats, // output for backward pass: (B, NH, T)
|
||||
floatX* inp, // input: (B, T, 3, NH, HS) QKV
|
||||
int B, int T, int NH, int C);
|
||||
int B, int T, int NH, int C, cudaStream_t stream);
|
||||
|
||||
void attention_backward_cudnn(floatX* dqkvr, // output
|
||||
floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs
|
||||
int B, int T, int NH, int C);
|
||||
int B, int T, int NH, int C, cudaStream_t stream);
|
||||
|
||||
#endif // CUDNN_ATT_H
|
||||
|
|
@ -150,12 +150,12 @@ __global__ void wpe_backward_kernel(floatX* dwpe,
|
|||
|
||||
void encoder_forward(floatX* out,
|
||||
const int* inp, const floatX* wte, const floatX* wpe,
|
||||
int B, int T, int C) {
|
||||
int B, int T, int C, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 256;
|
||||
const int N = B * T * C;
|
||||
const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size));
|
||||
encoder_forward_kernel3<<<grid_size, block_size>>>(out, inp, wte, wpe, B, T, C);
|
||||
encoder_forward_kernel3<<<grid_size, block_size, 0, stream>>>(out, inp, wte, wpe, B, T, C);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
|
|
@ -163,14 +163,14 @@ void encoder_forward(floatX* out,
|
|||
void encoder_backward(floatX* dwte, floatX* dwpe, floatX* scratch, // gpu outputs & scratch
|
||||
int* workload_indices, int4* bucket_info, // cpu scratch buffers
|
||||
const floatX* dout, const int* inp, const int* inputs_cpu, // cpu/gpu inputs
|
||||
int B, int T, int C, unsigned int seed) {
|
||||
int B, int T, int C, unsigned int seed, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
|
||||
// Launch wpe kernel first (so it runs on the GPU in parallel with the CPU pre-processing for wte)
|
||||
const int block_size = 256;
|
||||
const int N = T * C / x128::size;
|
||||
const int grid_size = CEIL_DIV(N, block_size);
|
||||
wpe_backward_kernel<<<grid_size, block_size, 0>>>(dwpe, dout, inp, B, T, C, seed);
|
||||
wpe_backward_kernel<<<grid_size, block_size, 0, stream>>>(dwpe, dout, inp, B, T, C, seed);
|
||||
cudaCheck(cudaGetLastError());
|
||||
|
||||
// check the GPU scratch buffer is large enough to hold the bucket info and workload indices
|
||||
|
|
@ -218,11 +218,11 @@ void encoder_backward(floatX* dwte, floatX* dwpe, floatX* scratch, // gpu output
|
|||
// todo - could use CUDA events (even without streams) to avoid CPU/GPU synchronisation completely
|
||||
int4* d_bucket_info = (int4*)scratch;
|
||||
int* d_workload_indices = (int*)(scratch + B*T*num_c_groups * sizeof(int4));
|
||||
cudaCheck(cudaMemcpyAsync(d_bucket_info, bucket_info, num_buckets * sizeof(int4), cudaMemcpyHostToDevice));
|
||||
cudaCheck(cudaMemcpy(d_workload_indices, workload_indices, total_items * sizeof(int), cudaMemcpyHostToDevice));
|
||||
cudaCheck(cudaMemcpyAsync(d_bucket_info, bucket_info, num_buckets * sizeof(int4), cudaMemcpyHostToDevice, stream));
|
||||
cudaCheck(cudaMemcpyAsync(d_workload_indices, workload_indices, total_items * sizeof(int), cudaMemcpyHostToDevice, stream));
|
||||
|
||||
// Launch wte kernel
|
||||
// todo - profile block sizes on more content (depends on number of buckets and on GPU?)
|
||||
wte_backward_kernel<256><<<num_buckets, 256>>>(dwte, d_bucket_info, d_workload_indices, dout, inp, seed, B, T, C);
|
||||
wte_backward_kernel<256><<<num_buckets, 256, 0, stream>>>(dwte, d_bucket_info, d_workload_indices, dout, inp, seed, B, T, C);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -133,11 +133,11 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
|
|||
template <typename Type>
|
||||
void fused_classifier(Type* logits, Type* losses,
|
||||
const float dloss, const int* targets,
|
||||
int B, int T, int V, int P) {
|
||||
int B, int T, int V, int P, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 1024;
|
||||
const int N = B * T;
|
||||
const int grid_size = N;
|
||||
fused_classifier_kernel5<<<grid_size, block_size>>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P);
|
||||
fused_classifier_kernel5<<<grid_size, block_size, 0, stream>>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -47,20 +47,20 @@ __global__ void gelu_backward_inplace_kernel(floatX* d_in_out, const floatX* inp
|
|||
// ----------------------------------------------------------------------------
|
||||
// kernel launchers
|
||||
|
||||
void gelu_forward(floatX* out, const floatX* inp, int N) {
|
||||
void gelu_forward(floatX* out, const floatX* inp, int N, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 512;
|
||||
assert(N % block_size == 0);
|
||||
const int grid_size = CEIL_DIV(N, block_size * x128::size);
|
||||
gelu_forward_kernel2<<<grid_size, block_size>>>(out, inp);
|
||||
gelu_forward_kernel2<<<grid_size, block_size, 0, stream>>>(out, inp);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void gelu_backward_inplace(floatX* d_in_out, const floatX* inp, const int N) {
|
||||
void gelu_backward_inplace(floatX* d_in_out, const floatX* inp, const int N, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 128;
|
||||
assert(N % block_size == 0);
|
||||
const int grid_size = CEIL_DIV(N, block_size * x128::size);
|
||||
gelu_backward_inplace_kernel<<<grid_size, block_size>>>(d_in_out, inp);
|
||||
gelu_backward_inplace_kernel<<<grid_size, block_size, 0, stream>>>(d_in_out, inp);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ __global__ void global_norm_squared_kernel(float* out, const T* data, size_t cou
|
|||
// kernel launcher
|
||||
|
||||
template<typename T>
|
||||
void global_norm_squared(float* out, const T* values, size_t count) {
|
||||
void global_norm_squared(float* out, const T* values, size_t count, cudaStream_t stream) {
|
||||
const int block_size = 512;
|
||||
// launch just enough blocks to fill the grid. deliberately no DIV_CEIL.
|
||||
// having one block less than possible is a tiny performance hit, having
|
||||
|
|
@ -45,7 +45,7 @@ void global_norm_squared(float* out, const T* values, size_t count) {
|
|||
assert(grid_size > 0); // gives a better error than letting the call below fail
|
||||
// initialize out with zero
|
||||
cudaCheck(cudaMemset(out, 0, sizeof(float)));
|
||||
global_norm_squared_kernel<<<grid_size, block_size>>>(out, values, count);
|
||||
global_norm_squared_kernel<<<grid_size, block_size, 0, stream>>>(out, values, count);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -356,28 +356,28 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with
|
|||
|
||||
void layernorm_forward(floatX* out, floatX* mean, floatX* rstd,
|
||||
floatX* inp, const floatX* weight, const floatX* bias,
|
||||
int B, int T, int C) {
|
||||
int B, int T, int C, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 512;
|
||||
const int N = B * T;
|
||||
const int grid_size = CEIL_DIV(N * WARP_SIZE, block_size);
|
||||
layernorm_forward_kernel3<<<grid_size, block_size>>>(out, mean, rstd, inp, weight, bias, N, C);
|
||||
layernorm_forward_kernel3<<<grid_size, block_size, 0, stream>>>(out, mean, rstd, inp, weight, bias, N, C);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N) {
|
||||
void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 256;
|
||||
assert(N % block_size == 0);
|
||||
const int grid_size = CEIL_DIV(N, block_size * x128::size);
|
||||
residual_forward_kernel<<<grid_size, block_size>>>(out, inp1, inp2);
|
||||
residual_forward_kernel<<<grid_size, block_size, 0, stream>>>(out, inp1, inp2);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void fused_residual_forward5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd,
|
||||
const floatX* inp1, const floatX* inp2,
|
||||
const floatX* weight, const floatX* bias,
|
||||
int N, int C) {
|
||||
int N, int C, cudaStream_t stream) {
|
||||
const int block_size = 256;
|
||||
int block_y = block_size / WARP_SIZE;
|
||||
const int grid_size = CEIL_DIV(N, block_y);
|
||||
|
|
@ -389,18 +389,19 @@ void fused_residual_forward5(floatX* residual, floatX* normed, floatX* mean, flo
|
|||
auto status = cudaFuncSetAttribute(fused_residual_forward_kernel5, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
cudaGetLastError();
|
||||
if(status == cudaSuccess) {
|
||||
fused_residual_forward_kernel5<<<grid_size, dim3(WARP_SIZE, block_y), smem>>>(residual, normed, mean, rstd, inp1, inp2,
|
||||
weight, bias, N, C);
|
||||
fused_residual_forward_kernel5<<<grid_size, dim3(WARP_SIZE, block_y), smem, stream>>>(residual, normed,
|
||||
mean, rstd, inp1, inp2,
|
||||
weight, bias, N, C);
|
||||
} else {
|
||||
residual_forward(residual, inp1, inp2, N*C);
|
||||
layernorm_forward(normed, mean, rstd, residual, weight, bias, N, 1, C);
|
||||
residual_forward(residual, inp1, inp2, N*C, stream);
|
||||
layernorm_forward(normed, mean, rstd, residual, weight, bias, N, 1, C, stream);
|
||||
}
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch,
|
||||
const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd,
|
||||
int B, int T, int C) {
|
||||
int B, int T, int C, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
const int block_size = 512;
|
||||
const int blocks_per_sm = 2; // supported on every architecture and less cache thrashing than 3
|
||||
|
|
@ -409,6 +410,6 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr
|
|||
size_t shared_mem_size = (2 * rounded_C + 2 * (block_size - 32) * f128::size) * sizeof(float);
|
||||
|
||||
cudaCheck(cudaMemset(scratch, 0, 1 * sizeof(float))); // only need to reset the flag to 0
|
||||
layernorm_backward_kernel10<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
|
||||
layernorm_backward_kernel10<<<grid_size, block_size, shared_mem_size, stream>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -105,7 +105,7 @@ __global__ void reduce_add_sum_kernel(floatX* dst, const float* src, size_t n, s
|
|||
// https://docs.nvidia.com/cuda/cublas/#cublasltmatmul
|
||||
void matmul_forward_cublaslt(floatX* out,
|
||||
floatX* inp, floatX* weight, floatX* bias,
|
||||
int B, int T, int C, int OC) {
|
||||
int B, int T, int C, int OC, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
int has_bias = (bias != NULL);
|
||||
|
||||
|
|
@ -162,7 +162,7 @@ void matmul_forward_cublaslt(floatX* out,
|
|||
cublasCheck(cublasLtMatmul(cublaslt_handle, operationDesc,
|
||||
&alpha, weight, weightLayout, inp, inputLayout, &beta,
|
||||
out, outputLayout, out, outputLayout, &heuristic.algo,
|
||||
cublaslt_workspace, cublaslt_workspace_size, 0));
|
||||
cublaslt_workspace, cublaslt_workspace_size, stream));
|
||||
|
||||
// cleanups
|
||||
cublasCheck(cublasLtMatmulPreferenceDestroy(preference));
|
||||
|
|
@ -176,7 +176,7 @@ void matmul_forward_cublaslt(floatX* out,
|
|||
void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
|
||||
floatX* dout, floatX* inp, floatX* weight,
|
||||
float* dbias_buffer,
|
||||
int B, int T, int C, int OC) {
|
||||
int B, int T, int C, int OC, cudaStream_t stream) {
|
||||
NVTX_RANGE_FN();
|
||||
float one = 1.0f, zero = 0.0f;
|
||||
|
||||
|
|
@ -195,18 +195,20 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
|
|||
// If we have enough OC that we don't need cross-block reductions, we can skip the bias_buffer accumulation
|
||||
// and write results directly to the output.
|
||||
if(grid_size_y == 1) {
|
||||
matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim>>>(dbias, dout, B, T, OC, std::bool_constant<false>{});
|
||||
matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim, 0, stream>>>(dbias, dout, B, T, OC, std::bool_constant<false>{});
|
||||
cudaCheck(cudaGetLastError());
|
||||
} else {
|
||||
// kernel 9 overwrites temp buffer, so no need to memset
|
||||
matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim>>>(dbias_buffer, dout, B, T, OC, std::bool_constant<true>{});
|
||||
matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim, 0, stream>>>(dbias_buffer, dout, B, T, OC, std::bool_constant<true>{});
|
||||
cudaCheck(cudaGetLastError());
|
||||
reduce_add_sum_kernel<<<CEIL_DIV(OC, 256 * f128::size), 256>>>(dbias, dbias_buffer, OC, grid_size_y);
|
||||
reduce_add_sum_kernel<<<CEIL_DIV(OC, 256 * f128::size), 256, 0, stream>>>(dbias, dbias_buffer, OC, grid_size_y);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
}
|
||||
|
||||
// backward to input, uses = in the backward pass (set the gradient)
|
||||
cublasCheck(cublasSetStream(cublas_handle, stream));
|
||||
cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
|
||||
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_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
||||
|
|
|
|||
159
train_gpt2.cu
159
train_gpt2.cu
|
|
@ -511,6 +511,7 @@ typedef struct {
|
|||
// todo - if other functions need cpu scratch buffers in the future, reuse as generic scratch?
|
||||
int* workload_indices; // encoder_backward, B*T*num_c_groups (int)
|
||||
int4* bucket_info; // encoder_backward, B*T*num_c_groups (int4) - size for worst case
|
||||
cudaStream_t main_stream;
|
||||
} GPT2;
|
||||
|
||||
void gpt2_init_common(GPT2 *model) {
|
||||
|
|
@ -541,6 +542,8 @@ void gpt2_init_common(GPT2 *model) {
|
|||
model->rng_state = 13371337; // used in stochastic rounding
|
||||
model->use_master_weights = 1; // safe default: do keep master weights in fp32
|
||||
model->recompute = 1; // good default: recompute gelu but not layernorm
|
||||
cudaStreamCreate(&model->main_stream);
|
||||
nvtxNameCudaStreamA(model->main_stream, "main stream");
|
||||
}
|
||||
|
||||
void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) {
|
||||
|
|
@ -633,6 +636,8 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
|
|||
fcloseCheck(model_file);
|
||||
|
||||
gpt2_init_common(model);
|
||||
// only return from this function once we are certain the params are ready on the GPU
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
void gpt2_build_from_random(GPT2 *model, int depth) {
|
||||
|
|
@ -724,7 +729,8 @@ void gpt2_build_from_random(GPT2 *model, int depth) {
|
|||
gpt2_init_common(model);
|
||||
}
|
||||
|
||||
void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, int grad_accum_steps=1) {
|
||||
void gpt2_forward(GPT2 *model, const int* inputs, const int* targets, size_t B, size_t T, int grad_accum_steps=1) {
|
||||
// right now, this function is fully synchronous with the host
|
||||
NVTX_RANGE_FN();
|
||||
// targets are optional and could be NULL
|
||||
// in this function we must be careful and use size_t instead of int, otherwise
|
||||
|
|
@ -737,22 +743,15 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
|
|||
}
|
||||
|
||||
// convenience parameters
|
||||
size_t V = model->config.vocab_size;
|
||||
size_t Vp = model->config.padded_vocab_size;
|
||||
size_t L = model->config.num_layers;
|
||||
size_t NH = model->config.num_heads;
|
||||
size_t C = model->config.channels;
|
||||
|
||||
// validate inputs, all indices must be in the range [0, V)
|
||||
for(int i = 0; i < B * T; i++) {
|
||||
assert(0 <= inputs[i] && inputs[i] < V);
|
||||
if (targets != NULL) {
|
||||
assert(0 <= targets[i] && targets[i] < V);
|
||||
}
|
||||
}
|
||||
const size_t V = model->config.vocab_size;
|
||||
const size_t Vp = model->config.padded_vocab_size;
|
||||
const size_t L = model->config.num_layers;
|
||||
const size_t NH = model->config.num_heads;
|
||||
const size_t C = model->config.channels;
|
||||
|
||||
// allocate space for all the activations if needed (done here, lazily)
|
||||
if(model->acts_memory == NULL) {
|
||||
NvtxRange rng("InitActs");
|
||||
// record the current B,T as well
|
||||
model->batch_size = B;
|
||||
model->seq_len = T;
|
||||
|
|
@ -780,19 +779,29 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
|
|||
}
|
||||
|
||||
// copy inputs/targets to the model
|
||||
// todo - inputs is copied on default stream so this synchronises CPU/GPU for now
|
||||
cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice));
|
||||
cudaCheck(cudaMemcpyAsync(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice, model->main_stream));
|
||||
if (targets != NULL) {
|
||||
cudaCheck(cudaMemcpy(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice));
|
||||
cudaCheck(cudaMemcpyAsync(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice, model->main_stream));
|
||||
}
|
||||
|
||||
cudaStream_t main_stream = model->main_stream;
|
||||
|
||||
// validate inputs, all indices must be in the range [0, V)
|
||||
// we can do this while the copies are already underway
|
||||
for(int i = 0; i < B * T; i++) {
|
||||
assert(0 <= inputs[i] && inputs[i] < V);
|
||||
if (targets != NULL) {
|
||||
assert(0 <= targets[i] && targets[i] < V);
|
||||
}
|
||||
}
|
||||
|
||||
// forward pass
|
||||
ParameterTensors params = model->params; // for brevity
|
||||
ActivationTensors acts = model->acts;
|
||||
encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0]
|
||||
encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C, main_stream); // encoding goes into residual[0]
|
||||
|
||||
// first layernorm isn't fused
|
||||
layernorm_forward(acts.ln1, acts.ln1_mean, acts.ln1_rstd, acts.encoded, params.ln1w, params.ln1b, B, T, C);
|
||||
layernorm_forward(acts.ln1, acts.ln1_mean, acts.ln1_rstd, acts.encoded, params.ln1w, params.ln1b, B, T, C, main_stream);
|
||||
|
||||
for (int l = 0; l < L; l++) {
|
||||
NvtxRange layer_range("Layer", l);
|
||||
|
|
@ -830,22 +839,22 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
|
|||
// now do the forward pass
|
||||
#ifdef ENABLE_CUDNN
|
||||
float* l_att = (float*)acts.att + l * B * NH * T; // cuDNN needs a smaller FP32 tensor
|
||||
matmul_forward_cublaslt(l_qkvr, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C);
|
||||
attention_forward_cudnn(l_atty, (float*)l_att, l_qkvr, B, T, NH, C);
|
||||
matmul_forward_cublaslt(l_qkvr, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C, main_stream);
|
||||
attention_forward_cudnn(l_atty, (float*)l_att, l_qkvr, B, T, NH, C, main_stream);
|
||||
#else
|
||||
floatX* l_att = acts.att + l * B * NH * T * T;
|
||||
// these are only needed as scratchpads for the forward pass, but
|
||||
// need not be stored for backward
|
||||
floatX* scratch = (floatX*)acts.output;
|
||||
matmul_forward_cublaslt(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C);
|
||||
attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH);
|
||||
matmul_forward_cublaslt(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C, main_stream);
|
||||
attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH, main_stream);
|
||||
#endif
|
||||
|
||||
matmul_forward_cublaslt(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C);
|
||||
fused_residual_forward5(l_residual2, l_ln2, l_ln2_mean, l_ln2_rstd, residual, l_attproj, l_ln2w, l_ln2b, B*T, C);
|
||||
matmul_forward_cublaslt(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C);
|
||||
gelu_forward(l_fch_gelu, l_fch, B*T*4*C);
|
||||
matmul_forward_cublaslt(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C);
|
||||
matmul_forward_cublaslt(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C, main_stream);
|
||||
fused_residual_forward5(l_residual2, l_ln2, l_ln2_mean, l_ln2_rstd, residual, l_attproj, l_ln2w, l_ln2b, B*T, C, main_stream);
|
||||
matmul_forward_cublaslt(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C, main_stream);
|
||||
gelu_forward(l_fch_gelu, l_fch, B*T*4*C, main_stream);
|
||||
matmul_forward_cublaslt(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C, main_stream);
|
||||
|
||||
// OK, fusion across blocks.
|
||||
if(l+1 != L) {
|
||||
|
|
@ -855,22 +864,22 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
|
|||
const floatX* l_ln1w = params.ln1w + (l + 1) * C;
|
||||
const floatX* l_ln1b = params.ln1b + (l + 1) * C;
|
||||
fused_residual_forward5(l_residual3, l_ln1, l_ln1_mean, l_ln1_rstd, l_residual2, l_fcproj, l_ln1w, l_ln1b,
|
||||
B * T, C);
|
||||
B * T, C, main_stream);
|
||||
} else {
|
||||
fused_residual_forward5(l_residual3, acts.lnf, acts.lnf_mean, acts.lnf_rstd, l_residual2, l_fcproj,
|
||||
params.lnfw, params.lnfb,
|
||||
B * T, C);
|
||||
B * T, C, main_stream);
|
||||
}
|
||||
}
|
||||
|
||||
matmul_forward_cublaslt(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp);
|
||||
matmul_forward_cublaslt(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp, main_stream);
|
||||
|
||||
// also forward the cross-entropy loss function if we have the targets
|
||||
if (targets != NULL) {
|
||||
NvtxRange classifier_and_loss_range("classifier_and_loss");
|
||||
// fused classifier: does the forward pass and first part of the backward pass
|
||||
const float dloss = 1.0f / (B * T * grad_accum_steps); // results in the uniform average loss over all elements
|
||||
fused_classifier(acts.output, acts.losses, dloss, model->targets, B, T, V, Vp);
|
||||
fused_classifier(acts.output, acts.losses, dloss, model->targets, B, T, V, Vp, main_stream);
|
||||
// for convenience also evaluate the mean loss (TODO re-think this compute+sync point)
|
||||
cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(floatX), cudaMemcpyDeviceToHost));
|
||||
float mean_loss = 0.0f;
|
||||
|
|
@ -885,13 +894,15 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
|
|||
// if we don't have targets, we don't have loss
|
||||
model->mean_loss = -1.0f;
|
||||
}
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
void gpt2_zero_grad(GPT2 *model) {
|
||||
NVTX_RANGE_FN();
|
||||
if (model->grads_memory != NULL) {
|
||||
cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX)));
|
||||
cudaCheck(cudaMemsetAsync(model->grads_memory, 0, model->num_parameters * sizeof(floatX), model->main_stream));
|
||||
}
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
void gpt2_backward(GPT2 *model, int* inputs) {
|
||||
|
|
@ -904,6 +915,7 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
|
||||
// lazily allocate the memory for gradients of the weights and activations, if needed
|
||||
if (model->grads_memory == NULL) {
|
||||
NvtxRange rng("InitGrads");
|
||||
// allocate buffers for weight gradients
|
||||
printf0("allocating %d MiB for parameter gradients\n", (int)round(model->num_parameters * sizeof(floatX) / (1024 * 1024)));
|
||||
model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_elements, model->param_sizeof);
|
||||
|
|
@ -928,12 +940,12 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
}
|
||||
|
||||
// convenience shortcuts, size_t instead of int so that pointer arithmetics don't overflow
|
||||
size_t B = model->batch_size;
|
||||
size_t T = model->seq_len;
|
||||
size_t Vp = model->config.padded_vocab_size;
|
||||
size_t L = model->config.num_layers;
|
||||
size_t NH = model->config.num_heads;
|
||||
size_t C = model->config.channels;
|
||||
const size_t B = model->batch_size;
|
||||
const size_t T = model->seq_len;
|
||||
const size_t Vp = model->config.padded_vocab_size;
|
||||
const size_t L = model->config.num_layers;
|
||||
const size_t NH = model->config.num_heads;
|
||||
const size_t C = model->config.channels;
|
||||
|
||||
// backward pass: go in the reverse order of the forward pass, and call backward() functions
|
||||
ParameterTensors params = model->params; // for brevity
|
||||
|
|
@ -942,22 +954,24 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
GradActTensors grads_acts = model->grads_acts;
|
||||
|
||||
// reset residual stream gradients (put here to work with gradient accumulation)
|
||||
cudaCheck(cudaMemset(model->grads_acts.residual3, 0, B * T * C * sizeof(floatX)));
|
||||
cudaCheck(cudaMemsetAsync(model->grads_acts.residual3, 0, B * T * C * sizeof(floatX), model->main_stream));
|
||||
|
||||
// re-use the output buffer of the forward pass as a scratchpad during backward pass
|
||||
float* scratchF = (float*)acts.output;
|
||||
floatX* scratchX = (floatX*)acts.output;
|
||||
|
||||
cudaStream_t main_stream = model->main_stream;
|
||||
|
||||
// we kick off the chain rule by filling in dlosses with 1.0f/(B*T)
|
||||
// this was done in the fused classifier kernel as last step of forward pass
|
||||
// technically that is a small, inline backward() pass of calculating
|
||||
// total, final loss as the mean over all losses over all (B,T) positions in the batch
|
||||
// next: backward the classifier matmul
|
||||
matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp);
|
||||
matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp, main_stream);
|
||||
// backward the final layernorm
|
||||
floatX* residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3
|
||||
floatX* dresidual = (floatX*)grads_acts.residual3; // the main buffer holding the gradient in the backward pass
|
||||
layernorm_backward(dresidual, grads.lnfw, grads.lnfb, scratchF, grads_acts.bt4c, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C);
|
||||
layernorm_backward(dresidual, grads.lnfw, grads.lnfb, scratchF, grads_acts.bt4c, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C, main_stream);
|
||||
|
||||
// from this point on, we no longer need the values stored in the last residual, so we can reuse that memory as generic
|
||||
// scratch for backward computations
|
||||
|
|
@ -1013,40 +1027,42 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
if(model->recompute >= 1) {
|
||||
// recompute >= 1 means we recompute gelu. in this case,
|
||||
// l_fch_gelu is just a buffer, so re-compute the gelu from l_fch here
|
||||
gelu_forward(l_fch_gelu, l_fch, B*T*4*C);
|
||||
gelu_forward(l_fch_gelu, l_fch, B*T*4*C, main_stream);
|
||||
}
|
||||
matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C);
|
||||
gelu_backward_inplace(dl_bt4c, l_fch, B*T*4*C);
|
||||
matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C, main_stream);
|
||||
gelu_backward_inplace(dl_bt4c, l_fch, B*T*4*C, main_stream);
|
||||
if(model->recompute >= 2) {
|
||||
// same as gelu above, l_ln1 and l_ln2 are just buffers if recompute >= 2, recompute them here on demand
|
||||
layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C);
|
||||
layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C, main_stream);
|
||||
}
|
||||
matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, scratchF, B, T, C, 4 * C);
|
||||
matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, scratchF, B, T, C, 4 * C, main_stream);
|
||||
// layernorm backward does += to the dresidual, so it correctly accumulates grad from the MLP block above
|
||||
layernorm_backward(dresidual, dl_ln2w, dl_ln2b, scratchF, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C);
|
||||
matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, scratchF, B, T, C, C);
|
||||
layernorm_backward(dresidual, dl_ln2w, dl_ln2b, scratchF, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C, main_stream);
|
||||
matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, scratchF, B, T, C, C, main_stream);
|
||||
|
||||
#ifdef ENABLE_CUDNN
|
||||
float* l_att = (float*)acts.att + l * B * NH * T; // cuDNN needs a smaller FP32 tensor
|
||||
attention_backward_cudnn(dl_bt4c, dl_btc, l_qkvr, l_atty, (float*)l_att, B, T, NH, C);
|
||||
attention_backward_cudnn(dl_bt4c, dl_btc, l_qkvr, l_atty, (float*)l_att, B, T, NH, C, main_stream);
|
||||
#else
|
||||
floatX* l_att = acts.att + l * B * NH * T * T;
|
||||
// we need B x T x (4)C buffers. l_atty and l_fch aren't needed anymore at this point, so reuse their memory
|
||||
floatX* buffer_a = l_atty;
|
||||
floatX* buffer_b = l_fch; // this is B x T x 4C, so even larger than what we need
|
||||
floatX* dl_preatt = (floatX*)grads_acts.preatt; // dedicated scratchpad allocation
|
||||
attention_backward(dl_bt4c, buffer_b, dl_preatt, scratchX, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH);
|
||||
attention_backward(dl_bt4c, buffer_b, dl_preatt, scratchX, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH, main_stream);
|
||||
#endif
|
||||
if(model->recompute >= 2) {
|
||||
layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C);
|
||||
layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C, main_stream);
|
||||
}
|
||||
// QKV parameter gradients
|
||||
matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, scratchF, B, T, C, 3 * C);
|
||||
matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, scratchF, B, T, C, 3 * C, main_stream);
|
||||
// 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);
|
||||
layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C, main_stream);
|
||||
}
|
||||
encoder_backward(grads.wte, grads.wpe, scratchX, model->workload_indices, model->bucket_info,
|
||||
dresidual, model->inputs, inputs, B, T, C, random_u32(&model->rng_state));
|
||||
dresidual, model->inputs, inputs, B, T, C, random_u32(&model->rng_state), main_stream);
|
||||
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
// Compute sum of a single CPU value across all GPU processes. No-op when multi-GPU is disabled.
|
||||
|
|
@ -1076,16 +1092,17 @@ void gpt2_multi_gpu_loss_and_grad_reduce(GPT2* model, MultiGpuConfig* multi_gpu_
|
|||
ncclCheck(ncclAllReduce(model->grads_memory, model->grads_memory,
|
||||
model->num_parameters,
|
||||
ncclFloatX, ncclAvg,
|
||||
multi_gpu_config->nccl_comm, 0));
|
||||
multi_gpu_config->nccl_comm, model->main_stream));
|
||||
} else if (multi_gpu_config->zero_stage == 1) {
|
||||
// ZERO-1: Get the average gradient only for local shard
|
||||
floatX* local_grads_memory = (floatX*) model->grads_memory + multi_gpu_config->shard_offset;
|
||||
ncclCheck(ncclReduceScatter(model->grads_memory, local_grads_memory,
|
||||
multi_gpu_config->shard_num_parameters,
|
||||
ncclFloatX, ncclAvg,
|
||||
multi_gpu_config->nccl_comm, 0));
|
||||
multi_gpu_config->nccl_comm, model->main_stream));
|
||||
}
|
||||
#endif
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_clip, int t, MultiGpuConfig* multi_gpu_config) {
|
||||
|
|
@ -1103,18 +1120,19 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
|
||||
// lazily allocate m,v memory and master weights (usually on the first iteration)
|
||||
if (model->m_memory == NULL) {
|
||||
NvtxRange rng("InitOpt");
|
||||
printf0("allocating %zu MiB for AdamW optimizer state m\n", (shard_num_parameters * sizeof(float)) >> 20);
|
||||
printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20);
|
||||
cudaCheck(cudaMalloc((void**)&model->m_memory, shard_num_parameters * sizeof(float)));
|
||||
cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float)));
|
||||
cudaCheck(cudaMemset(model->m_memory, 0, shard_num_parameters * sizeof(float)));
|
||||
cudaCheck(cudaMemset(model->v_memory, 0, shard_num_parameters * sizeof(float)));
|
||||
cudaCheck(cudaMemsetAsync(model->m_memory, 0, shard_num_parameters * sizeof(float), model->main_stream));
|
||||
cudaCheck(cudaMemsetAsync(model->v_memory, 0, shard_num_parameters * sizeof(float), model->main_stream));
|
||||
}
|
||||
if (model->use_master_weights == 1 && model->master_weights == NULL) {
|
||||
printf0("allocating %zu MiB for master copy of params\n", (shard_num_parameters * sizeof(float)) >> 20);
|
||||
cudaCheck(cudaMalloc((void**)&model->master_weights, shard_num_parameters * sizeof(float)));
|
||||
size_t grid_size = CEIL_DIV(shard_num_parameters, 512);
|
||||
copy_and_cast_kernel<<<grid_size, 512>>>(model->master_weights, params_memory + shard_offset, shard_num_parameters);
|
||||
copy_and_cast_kernel<<<grid_size, 512, 0, model->main_stream>>>(model->master_weights, params_memory + shard_offset, shard_num_parameters);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
|
|
@ -1125,11 +1143,11 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
// ^1 because of the ncclReduceScatter() in gpt2_multi_gpu_loss_and_grad_reduce,
|
||||
// grads_memory only contains the averaged gradients at the local shard
|
||||
// so we only calculate the grad norm at the grads_memory belonging to the local shard
|
||||
global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters);
|
||||
global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters, model->main_stream);
|
||||
} else {
|
||||
// the ncclAllReduce() in gpt2_multi_gpu_loss_and_grad_reduce has averaged the gradients across all GPUs
|
||||
// so each GPU can compute the squared norm over the whole grad vector, with no added comms needed
|
||||
global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters);
|
||||
global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters, model->main_stream);
|
||||
}
|
||||
// transfer the gradient norm to CPU
|
||||
float grad_norm_squared_cpu = 0.0f;
|
||||
|
|
@ -1149,9 +1167,6 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
float grad_scale = (grad_norm_cpu > grad_clip) ? grad_clip / grad_norm_cpu : 1.0f;
|
||||
|
||||
// AdamW update
|
||||
int block_size = 512;
|
||||
float beta1_correction = 1.0f - powf(beta1, t);
|
||||
float beta2_correction = 1.0f - powf(beta2, t);
|
||||
unsigned int seed = random_u32(&model->rng_state);
|
||||
// individually call the adamw_kernel3 on all parameter tensors separately
|
||||
size_t offset = 0;
|
||||
|
|
@ -1202,16 +1217,16 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
// - the position embeddings actively participate at every forward/backward pass
|
||||
float wd = (i == 0 || i == 1 || i == 4 || i == 6 || i == 10 || i == 12) ? weight_decay : 0.0f;
|
||||
// ok finally call the kernel
|
||||
size_t num_blocks = CEIL_DIV(local_params, block_size);
|
||||
adamw_kernel3<<<num_blocks, block_size>>>(params_ptr, master_ptr, grad_ptr,
|
||||
m_ptr, v_ptr, local_params, learning_rate,
|
||||
beta1, beta2, beta1_correction, beta2_correction,
|
||||
eps, wd, grad_scale, seed);
|
||||
adamw_update(params_ptr, master_ptr, grad_ptr,
|
||||
m_ptr, v_ptr, local_params, learning_rate,
|
||||
beta1, beta2, t, eps, wd, grad_scale, seed, model->main_stream);
|
||||
}
|
||||
// advance the offset pointer to the next parameter tensor
|
||||
offset += num_parameters;
|
||||
}
|
||||
cudaCheck(cudaGetLastError());
|
||||
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
return grad_norm_cpu;
|
||||
}
|
||||
|
||||
|
|
@ -1223,10 +1238,11 @@ void gpt2_multi_gpu_param_gather(GPT2 *model, MultiGpuConfig* multi_gpu_config)
|
|||
// gather updated shards of model->params_memory from each process
|
||||
ncclCheck(ncclAllGather((floatX*)model->params_memory + multi_gpu_config->shard_offset, (floatX*)model->params_memory,
|
||||
multi_gpu_config->shard_num_parameters, ncclFloatX,
|
||||
multi_gpu_config->nccl_comm, 0));
|
||||
multi_gpu_config->nccl_comm, model->main_stream));
|
||||
}
|
||||
cudaCheck(cudaGetLastError());
|
||||
#endif
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
float gpt2_estimate_mfu(GPT2 *model, int num_tokens, float dt) {
|
||||
|
|
@ -1260,6 +1276,7 @@ void gpt2_free(GPT2 *model) {
|
|||
cudaCheck(cudaFree(model->grads_acts_memory));
|
||||
cudaCheck(cudaFree(model->inputs));
|
||||
cudaCheck(cudaFree(model->targets));
|
||||
cudaCheck(cudaStreamDestroy(model->main_stream));
|
||||
cudaCheck(cudaFreeHost(model->cpu_losses));
|
||||
cudaCheck(cudaFreeHost(model->cpu_losses_fp32));
|
||||
free(model->workload_indices);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue