mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-27 20:25:09 -04:00
1110 lines
48 KiB
Text
1110 lines
48 KiB
Text
/*
|
|
Kernels for layernorm backward pass.
|
|
|
|
Compile example:
|
|
nvcc -O3 --use_fast_math -lcublas -lcublasLt layernorm_backward.cu -o layernorm_backward
|
|
|
|
version 1 is naive port from CPU code to kernel: parallelizes over B,T, loops over C
|
|
./layernorm_backward 1
|
|
|
|
version 2 moves a lot of reduction to shared memory over global memory
|
|
./layernorm_backward 2
|
|
*/
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <cuda_runtime.h>
|
|
#include <assert.h>
|
|
#include <cooperative_groups.h>
|
|
#include <cooperative_groups/reduce.h>
|
|
|
|
#define ENABLE_BF16
|
|
#include "common.h"
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// CPU code reference
|
|
|
|
void layernorm_forward_cpu(float* out, float* mean, float* rstd,
|
|
const float* inp, const float* weight, const float* bias,
|
|
int B, int T, int C) {
|
|
// reference: https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
|
|
// both inp and out are (B,T,C) of the activations
|
|
// mean and rstd are (B,T) buffers, to be used later in backward pass
|
|
// at each position (b,t) of the input, the C-dimensional vector
|
|
// of activations gets normalized, then scaled and shifted
|
|
float eps = 1e-5f;
|
|
for (int b = 0; b < B; b++) {
|
|
for (int t = 0; t < T; t++) {
|
|
// seek to the input position inp[b,t,:]
|
|
const float* x = inp + b * T * C + t * C;
|
|
// calculate the mean
|
|
float m = 0.0f;
|
|
for (int i = 0; i < C; i++) {
|
|
m += x[i];
|
|
}
|
|
m = m/C;
|
|
// calculate the variance (without any bias correction)
|
|
float v = 0.0f;
|
|
for (int i = 0; i < C; i++) {
|
|
float xshift = x[i] - m;
|
|
v += xshift * xshift;
|
|
}
|
|
v = v/C;
|
|
// calculate the rstd (reciprocal standard deviation)
|
|
float s = 1.0f / sqrtf(v + eps);
|
|
// seek to the output position in out[b,t,:]
|
|
float* out_bt = out + b * T * C + t * C;
|
|
for (int i = 0; i < C; i++) {
|
|
float n = (s * (x[i] - m)); // normalize
|
|
float o = n * weight[i] + bias[i]; // scale and shift
|
|
out_bt[i] = o; // write
|
|
}
|
|
// cache the mean and rstd for the backward pass later
|
|
mean[b * T + t] = m;
|
|
rstd[b * T + t] = s;
|
|
}
|
|
}
|
|
}
|
|
|
|
void layernorm_backward_cpu(float* dinp, float* dweight, float* dbias,
|
|
const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd,
|
|
int B, int T, int C) {
|
|
for (int b = 0; b < B; b++) {
|
|
for (int t = 0; t < T; t++) {
|
|
const float* dout_bt = dout + b * T * C + t * C;
|
|
const float* inp_bt = inp + b * T * C + t * C;
|
|
float* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = mean[b * T + t];
|
|
const float rstd_bt = rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = 0; i < C; i++) {
|
|
float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = weight[i] * dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = 0; i < C; i++) {
|
|
float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = weight[i] * dout_bt[i];
|
|
// gradient contribution to bias
|
|
dbias[i] += dout_bt[i];
|
|
// gradient contribution to weight
|
|
dweight[i] += norm_bti * dout_bt[i];
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] += dval;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// GPU kernels
|
|
|
|
// 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<uintptr_t>(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<uintptr_t>(addr);
|
|
half2* ptr_fp16 = reinterpret_cast<half2*>(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);
|
|
}
|
|
|
|
// super naive kernel that just parallelizes over B,T and loops over C
|
|
__global__ void layernorm_backward_kernel1(float* dinp, float* dweight, float* dbias,
|
|
const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd,
|
|
int B, int T, int C) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx >= B*T) return;
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const float* dout_bt = dout + b * T * C + t * C;
|
|
const float* inp_bt = inp + b * T * C + t * C;
|
|
float* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = mean[b * T + t];
|
|
const float rstd_bt = rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = 0; i < C; i++) {
|
|
float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = weight[i] * dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = 0; i < C; i++) {
|
|
float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = weight[i] * dout_bt[i];
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias[i], dout_bt[i]);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight[i], norm_bti * dout_bt[i]);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] += dval;
|
|
}
|
|
}
|
|
|
|
// uses shared memory instead for the reduces
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
__global__ void layernorm_backward_kernel2(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C) {
|
|
extern __shared__ float shared[]; // size = 2 * C
|
|
|
|
namespace cg = cooperative_groups;
|
|
cg::thread_block block = cg::this_thread_block();
|
|
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
|
|
int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
|
|
int N = B * T;
|
|
if(idx >= N) { return; } // thread guards
|
|
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const Tdout* dout_bt = dout + b * T * C + t * C;
|
|
const Trest* inp_bt = inp + b * T * C + t * C;
|
|
Tdinp* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
float* dweight_shared = shared + C;
|
|
|
|
// init shared memory to zero
|
|
#pragma unroll
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x){
|
|
dbias_shared[i] = 0.0f;
|
|
dweight_shared[i] = 0.0f;
|
|
}
|
|
__syncthreads();
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
|
|
dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias_shared[i], (float)dout_bt[i]);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight_shared[i], norm_bti * (float)dout_bt[i]);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval);
|
|
}
|
|
__syncthreads();
|
|
|
|
// write to global memory
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
atomicAddX(&dbias[i], (Tparams)dbias_shared[i]);
|
|
atomicAddX(&dweight[i], (Tparams)dweight_shared[i]);
|
|
}
|
|
}
|
|
|
|
// kernel2 is 1 threadblock for all Cs on 32 BTs (assuming threadblock size of 1024 threads = 32 warps)
|
|
// To minimise the amount of atomicAdds, we will aim for 1 threadblock per SM, processing (total BTs / threadblocks) BTs
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
__global__ void layernorm_backward_kernel3(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C) {
|
|
extern __shared__ float shared[]; // size = 2 * C
|
|
|
|
namespace cg = cooperative_groups;
|
|
cg::thread_block block = cg::this_thread_block();
|
|
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
|
|
int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
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;
|
|
}
|
|
__syncthreads();
|
|
|
|
int warps_in_grid = gridDim.x * warp.meta_group_size();
|
|
for (int idx = base_idx; idx < B * T; idx += warps_in_grid) {
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const Tdout* dout_bt = dout + b * T * C + t * C;
|
|
const Trest* inp_bt = inp + b * T * C + t * C;
|
|
Tdinp* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
|
|
dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float dout_i = (float)__ldcs(&dout_bt[i]);
|
|
float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * dout_i;
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias_shared[i], dout_i);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight_shared[i], norm_bti * dout_i);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
atomicAddX(&dbias[i], (Tparams)dbias_shared[i]);
|
|
atomicAddX(&dweight[i], (Tparams)dweight_shared[i]);
|
|
}
|
|
}
|
|
|
|
// atomicCAS version of kernel3
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
__global__ void layernorm_backward_kernel4(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C) {
|
|
extern __shared__ float shared[]; // size = 2 * C
|
|
|
|
namespace cg = cooperative_groups;
|
|
cg::thread_block block = cg::this_thread_block();
|
|
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
|
|
int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
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;
|
|
}
|
|
__syncthreads();
|
|
|
|
int warps_in_grid = gridDim.x * warp.meta_group_size();
|
|
for (int idx = base_idx; idx < B * T; idx += warps_in_grid) {
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const Tdout* dout_bt = dout + b * T * C + t * C;
|
|
const Trest* inp_bt = inp + b * T * C + t * C;
|
|
Tdinp* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
|
|
dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float dout_i = (float)__ldcs(&dout_bt[i]);
|
|
float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * dout_i;
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias_shared[i], dout_i);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight_shared[i], norm_bti * dout_i);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
__nv_bfloat162* dbiasVec2 = reinterpret_cast<__nv_bfloat162*>(dbias);
|
|
__nv_bfloat162* dweightVec2 = reinterpret_cast<__nv_bfloat162*>(dweight);
|
|
|
|
// write to global memory
|
|
for(int i = threadIdx.x; i < C/2; i+= blockDim.x) {
|
|
__nv_bfloat162 add_dbias = __halves2bfloat162((__nv_bfloat16)dbias_shared[i*2], (__nv_bfloat16)dbias_shared[i*2+1]);
|
|
__nv_bfloat162 add_dweight = __halves2bfloat162((__nv_bfloat16)dweight_shared[i*2], (__nv_bfloat16)dweight_shared[i*2+1]);
|
|
|
|
// Get the current value from L2 cache
|
|
__nv_bfloat162 current_dbias = __ldcg(&dbiasVec2[i]);
|
|
__nv_bfloat162 current_dweight = __ldcg(&dweightVec2[i]);
|
|
|
|
// Add the two values
|
|
__nv_bfloat162 new_dbias = add_dbias + current_dbias;
|
|
__nv_bfloat162 new_dweight = add_dweight + current_dweight;
|
|
|
|
// Write the result back to L2 cache using 32-bit integer atomic compare and exchange
|
|
uint current_dbias32b = *reinterpret_cast<uint*>(¤t_dbias);
|
|
uint current_dweight32b = *reinterpret_cast<uint*>(¤t_dweight);
|
|
|
|
uint new_dbias32b = *reinterpret_cast<uint*>(&new_dbias);
|
|
uint new_dweight32b = *reinterpret_cast<uint*>(&new_dweight);
|
|
|
|
uint old_dbias32b = atomicCAS((uint*)&dbiasVec2[i], current_dbias32b, new_dbias32b);
|
|
uint old_dweight32b = atomicCAS((uint*)&dweightVec2[i], current_dweight32b, new_dweight32b);
|
|
|
|
// If the value has changed between read and atomic, we need to try again
|
|
while (old_dbias32b != current_dbias32b) {
|
|
current_dbias32b = old_dbias32b;
|
|
new_dbias = *reinterpret_cast<__nv_bfloat162*>(¤t_dbias32b) + add_dbias;
|
|
new_dbias32b = *reinterpret_cast<uint*>(&new_dbias);
|
|
old_dbias32b = atomicCAS((uint*)&dbiasVec2[i], current_dbias32b, new_dbias32b);
|
|
}
|
|
|
|
while (old_dweight32b != current_dweight32b) {
|
|
current_dweight32b = old_dweight32b;
|
|
new_dweight = *reinterpret_cast<__nv_bfloat162*>(¤t_dweight32b) + add_dweight;
|
|
new_dweight32b = *reinterpret_cast<uint*>(&new_dweight);
|
|
old_dweight32b = atomicCAS((uint*)&dweightVec2[i], current_dweight32b, new_dweight32b);
|
|
}
|
|
}
|
|
}
|
|
|
|
// FP32 scratchpad per threadgroup, zero atomics except atomicAdd on uint for the flag (based on kernel3)
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
__global__ void layernorm_backward_kernel5(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C) {
|
|
extern __shared__ float shared[]; // size = 2 * C + 1
|
|
|
|
namespace cg = cooperative_groups;
|
|
cg::thread_block block = cg::this_thread_block();
|
|
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
|
|
int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
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;
|
|
}
|
|
uint *tmp_flag = (uint*)(shared + C*2);
|
|
__syncthreads();
|
|
|
|
int warps_in_grid = gridDim.x * warp.meta_group_size();
|
|
for (int idx = base_idx; idx < B * T; idx += warps_in_grid) {
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const Tdout* dout_bt = dout + b * T * C + t * C;
|
|
const Trest* inp_bt = inp + b * T * C + t * C;
|
|
Tdinp* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
|
|
dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float dout_i = (float)__ldcs(&dout_bt[i]);
|
|
float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * dout_i;
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias_shared[i], dout_i);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight_shared[i], norm_bti * dout_i);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
float* scratch_dbias = scratch;
|
|
float* scratch_dweight = scratch + C * gridDim.x;
|
|
uint* scratchFlag = (uint*)(scratch + (2 * C * gridDim.x));
|
|
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
scratch_dbias[i + C*blockIdx.x] = dbias_shared[i];
|
|
scratch_dweight[i + C*blockIdx.x] = dweight_shared[i];
|
|
}
|
|
__threadfence();
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
*tmp_flag = atomicAdd(scratchFlag, 1);
|
|
}
|
|
__syncthreads();
|
|
if (*tmp_flag == gridDim.x-1) {
|
|
// last block to finish, accumulate the scratchpad
|
|
for (int i = threadIdx.x; i < C; i += blockDim.x) {
|
|
float dbias_sum = 0.0f;
|
|
float dweight_sum = 0.0f;
|
|
#pragma unroll 8
|
|
for (int j = 0; j < gridDim.x; j++) {
|
|
dbias_sum += scratch_dbias[i + j*C];
|
|
dweight_sum += scratch_dweight[i + j*C];
|
|
}
|
|
dbias[i] = (Tparams)((float)dbias[i] + dbias_sum);
|
|
dweight[i] = (Tparams)((float)dweight[i] + dweight_sum);
|
|
}
|
|
}
|
|
}
|
|
|
|
// single FP32 scratchpad shared by all the threadblocks (based on kernels 3 & 5)
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
__global__ void layernorm_backward_kernel6(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C) {
|
|
extern __shared__ float shared[]; // size = 2 * C + 1
|
|
|
|
namespace cg = cooperative_groups;
|
|
cg::thread_block block = cg::this_thread_block();
|
|
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
|
|
int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
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;
|
|
}
|
|
uint *tmp_flag = (uint*)(shared + C*2);
|
|
__syncthreads();
|
|
|
|
int warps_in_grid = gridDim.x * warp.meta_group_size();
|
|
for (int idx = base_idx; idx < B * T; idx += warps_in_grid) {
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const Tdout* dout_bt = dout + b * T * C + t * C;
|
|
const Trest* inp_bt = inp + b * T * C + t * C;
|
|
Tdinp* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
|
|
dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
|
|
dnorm_mean = dnorm_mean / C;
|
|
dnorm_norm_mean = dnorm_norm_mean / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
|
|
float dout_i = (float)__ldcs(&dout_bt[i]);
|
|
float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * dout_i;
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias_shared[i], dout_i);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight_shared[i], norm_bti * dout_i);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval);
|
|
}
|
|
}
|
|
|
|
// Accumulate into a FP32 scratchpad
|
|
// BF16 atomics are potentially much slower... and this is more precise!
|
|
__syncthreads();
|
|
float* scratch_dbias = scratch;
|
|
float* scratch_dweight = scratch + C;
|
|
uint* scratchFlag = (uint*)(scratch + (2 * C));
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
atomicAdd(&scratch_dbias[i], dbias_shared[i]);
|
|
atomicAdd(&scratch_dweight[i], dweight_shared[i]);
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
*tmp_flag = atomicAdd(scratchFlag, 1);
|
|
}
|
|
__syncthreads();
|
|
if (*tmp_flag == gridDim.x-1) {
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
// todo - potentially do stochastic rounding here as well
|
|
dbias[i] = (Tparams)scratch_dbias[i];
|
|
dweight[i] = (Tparams)scratch_dweight[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// Same as kernel 6 but without cooperative groups or templates
|
|
__global__ void layernorm_backward_kernel7(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) {
|
|
extern __shared__ float shared[]; // size = 2 * C + 1
|
|
int warpId = threadIdx.x / warpSize; // warp index within a block
|
|
int warpsInBlock = blockDim.x / warpSize;
|
|
int base_idx = blockIdx.x * warpsInBlock + warpId;
|
|
int warpThreadIdx = threadIdx.x % warpSize; // Thread index within the warp
|
|
int warps_in_grid = gridDim.x * warpsInBlock;
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
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;
|
|
}
|
|
uint *tmp_flag = (uint*)(shared + C*2);
|
|
__syncthreads();
|
|
|
|
for (int idx = base_idx; idx < B * T; idx += warps_in_grid) {
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const floatX* dout_bt = dout + b * T * C + t * C;
|
|
const floatX* inp_bt = inp + b * T * C + t * C;
|
|
floatX* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warpThreadIdx; i < C; i += warpSize) {
|
|
float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight[i] * (float)dout_bt[i];
|
|
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;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
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;
|
|
float dnorm_i = (float)weight[i] * dout_i;
|
|
// gradient contribution to bias
|
|
atomicAdd(&dbias_shared[i], dout_i);
|
|
// gradient contribution to weight
|
|
atomicAdd(&dweight_shared[i], norm_bti * dout_i);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp_bt[i] = (floatX)((float)dinp_bt[i] + dval);
|
|
}
|
|
}
|
|
|
|
// Accumulate into a FP32 scratchpad
|
|
// BF16 atomics are potentially much slower... and this is more precise!
|
|
__syncthreads();
|
|
float* scratch_dbias = scratch;
|
|
float* scratch_dweight = scratch + C;
|
|
uint* scratchFlag = (uint*)(scratch + (2 * C));
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
atomicAdd(&scratch_dbias[i], dbias_shared[i]);
|
|
atomicAdd(&scratch_dweight[i], dweight_shared[i]);
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
*tmp_flag = atomicAdd(scratchFlag, 1);
|
|
}
|
|
__syncthreads();
|
|
if (*tmp_flag == gridDim.x-1) {
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
// todo - potentially do stochastic rounding here as well
|
|
dbias[i] = (floatX)scratch_dbias[i];
|
|
dweight[i] = (floatX)scratch_dweight[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
__global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
|
|
layernorm_backward_kernel8(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) {
|
|
extern __shared__ float shared[]; // size = 2 * C + 1
|
|
int warpId = threadIdx.x / warpSize; // warp index within a block
|
|
int warpsInBlock = blockDim.x / warpSize; //number of warps in block
|
|
int baseIdx = blockIdx.x * warpsInBlock + warpId;
|
|
int warpThreadIdx = threadIdx.x % warpSize; // Thread index within the warp
|
|
int warpsInGrid = gridDim.x * warpsInBlock;
|
|
int C_per_iteration = warpSize * x128::size;
|
|
int iterations_C = C / C_per_iteration;
|
|
|
|
// the first half of shared memory is bias, second is weight
|
|
float* dbias_shared = shared;
|
|
float* dweight_shared = shared + C;
|
|
|
|
// init shared memory to zero
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x){
|
|
dbias_shared[i] = 0.0f;
|
|
dweight_shared[i] = 0.0f;
|
|
}
|
|
unsigned int *tmp_flag = (unsigned int*)(shared + C*2);
|
|
__syncthreads();
|
|
|
|
for (int idx = baseIdx; idx < B * T; idx += warpsInGrid) {
|
|
int b = idx / T;
|
|
int t = idx % T;
|
|
|
|
const floatX* dout_bt = dout + b * T * C + t * C;
|
|
const floatX* inp_bt = inp + b * T * C + t * C;
|
|
floatX* dinp_bt = dinp + b * T * C + t * C;
|
|
const float mean_bt = (float)mean[b * T + t];
|
|
const float rstd_bt = (float)rstd[b * T + t];
|
|
|
|
// first: two reduce operations
|
|
float dnorm_mean = 0.0f;
|
|
float dnorm_norm_mean = 0.0f;
|
|
for (int i = warpThreadIdx * x128::size; i < C; i += warpSize * x128::size) {
|
|
x128 dout128_i = load128(dout_bt + i);
|
|
x128 inp128_i = load128(inp_bt + i);
|
|
x128 weight128_i = load128(weight + i);
|
|
for (int k = 0; k < x128::size; k++) {
|
|
float norm_bti = ((float)inp128_i[k] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight128_i[k] * (float)dout128_i[k];
|
|
dnorm_mean += dnorm_i;
|
|
dnorm_norm_mean += dnorm_i * norm_bti;
|
|
}
|
|
}
|
|
dnorm_mean = warpReduceSum(dnorm_mean) / C;
|
|
dnorm_norm_mean = warpReduceSum(dnorm_norm_mean) / C;
|
|
|
|
// now iterate again and accumulate all the gradients
|
|
// unfortunately we cannot use the same index for x128 arrays and shared memory
|
|
// as atomics can only be 32-bit rather than 128-bit (at least pre-SM90/Hopper)
|
|
// so this would result in an 8-way bank conflict, and kill performance
|
|
// so instead, we use a shared memory friendly index, and reorder before the final write
|
|
for (int i = 0; i < iterations_C; i++) {
|
|
int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration);
|
|
int shared_index = warpThreadIdx + (i * C_per_iteration);
|
|
x128 dout128 = load128cs(dout_bt + global_index);
|
|
x128 inp128 = load128cs(inp_bt + global_index);
|
|
x128 dinp128 = load128(dinp_bt + global_index);
|
|
x128 weight128 = load128(weight + global_index);
|
|
|
|
for (int x = 0; x < x128::size; x++) {
|
|
float dout_i = (float)dout128[x];
|
|
float norm_bti = ((float)inp128[x] - mean_bt) * rstd_bt;
|
|
float dnorm_i = (float)weight128[x] * dout_i;
|
|
// gradient contribution to bias (using shared memory friendly index)
|
|
atomicAdd(&dbias_shared[shared_index + x*warpSize], dout_i);
|
|
// gradient contribution to weight (using shared memory friendly index)
|
|
atomicAdd(&dweight_shared[shared_index + x*warpSize], norm_bti * dout_i);
|
|
// gradient contribution to input
|
|
float dval = 0.0f;
|
|
dval += dnorm_i; // term 1
|
|
dval -= dnorm_mean; // term 2
|
|
dval -= norm_bti * dnorm_norm_mean; // term 3
|
|
dval *= rstd_bt; // final scale
|
|
dinp128[x] = (floatX)((float)dinp128[x] + dval);
|
|
}
|
|
// cache in L2 as this is read by the next kernel, but bypass L1 to minimise thrashing
|
|
store128cg(dinp_bt + global_index, dinp128);
|
|
}
|
|
}
|
|
// 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
|
|
__syncthreads();
|
|
float* scratch_dbias = scratch;
|
|
float* scratch_dweight = scratch + C;
|
|
unsigned int* scratchFlag = (unsigned int*)(scratch + (2 * C));
|
|
for(int i = threadIdx.x; i < C; i+= blockDim.x) {
|
|
// global atomics in the same "shared memory banking friendly" order
|
|
atomicAdd(&scratch_dbias[i], dbias_shared[i]);
|
|
atomicAdd(&scratch_dweight[i], dweight_shared[i]);
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
*tmp_flag = atomicInc(scratchFlag, gridDim.x);
|
|
}
|
|
__syncthreads();
|
|
if (*tmp_flag == gridDim.x-1) {
|
|
for (int i = warpId; i < iterations_C; i += warpsInBlock) {
|
|
// reorder from atomic/shared memory-friendly index to real global memory index
|
|
// and convert from float/FP32 to floatX/BF16 for the final write
|
|
int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration);
|
|
int shared_index = warpThreadIdx + (i * C_per_iteration);
|
|
|
|
x128 dbias128 = load128(dbias + global_index);
|
|
x128 dweight128 = load128(dweight + global_index);
|
|
for (int x = 0; x < x128::size; x++) {
|
|
float s_db = scratch_dbias[shared_index + x*warpSize];
|
|
float s_dw = scratch_dweight[shared_index + x*warpSize];
|
|
dbias128[x] = (floatX)(s_db + (float)dbias128[x]);
|
|
dweight128[x] = (floatX)(s_dw + (float)dweight128[x]);
|
|
}
|
|
store128(dbias + global_index, dbias128);
|
|
store128(dweight + global_index, dweight128);
|
|
}
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// kernel launchers
|
|
|
|
void layernorm_backward1(float* dinp, float* dweight, float* dbias,
|
|
const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd,
|
|
int B, int T, int C, const int block_size) {
|
|
const int N = B * T;
|
|
const int grid_size = ceil_div(N, block_size);
|
|
layernorm_backward_kernel1<<<grid_size, block_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward2(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int N = B * T;
|
|
const int grid_size = ceil_div(32*N, block_size);
|
|
size_t shared_mem_size = 2 * C * sizeof(float);
|
|
layernorm_backward_kernel2<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward3(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int grid_size = (1024/block_size) * cuda_num_SMs;
|
|
size_t shared_mem_size = 2 * C * sizeof(float);
|
|
layernorm_backward_kernel3<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward4(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int grid_size = (1024/block_size) * cuda_num_SMs;
|
|
size_t shared_mem_size = 2 * C * sizeof(float);
|
|
layernorm_backward_kernel4<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward5(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int grid_size = 1 * cuda_num_SMs; // only support 1 block per SM for simplicity, 1024 threads is best anyway
|
|
size_t shared_mem_size = (2 * C + 1) * sizeof(float);
|
|
cudaMemset(scratch, 0, (grid_size * 2 * C + 1) * sizeof(float));
|
|
layernorm_backward_kernel5<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward6(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int grid_size = (1024/block_size) * cuda_num_SMs;
|
|
size_t shared_mem_size = (2 * C + 1) * sizeof(float);
|
|
|
|
// Including this as part of the timing until we can parallelise it
|
|
// It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams
|
|
cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float));
|
|
|
|
layernorm_backward_kernel6<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward7(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int grid_size = (1024/block_size) * cuda_num_SMs;
|
|
size_t shared_mem_size = (2 * C + 1) * sizeof(float);
|
|
|
|
// Including this as part of the timing until we can parallelise it
|
|
// It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams
|
|
cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float));
|
|
|
|
layernorm_backward_kernel7<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
|
|
void layernorm_backward8(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch,
|
|
const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
|
|
int B, int T, int C, int block_size) {
|
|
const int grid_size = (1024/block_size) * cuda_num_SMs;
|
|
size_t shared_mem_size = (2 * C + 1) * sizeof(float);
|
|
|
|
// Including this as part of the timing until we can parallelise it
|
|
// It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams
|
|
cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float));
|
|
|
|
layernorm_backward_kernel8<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
|
|
}
|
|
|
|
// kernel version dispatch
|
|
void layernorm_backward(int kernel_num,
|
|
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,
|
|
const int block_size) {
|
|
switch (kernel_num) {
|
|
#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16)
|
|
case 1:
|
|
layernorm_backward1(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
#endif
|
|
case 2:
|
|
layernorm_backward2(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
case 3:
|
|
layernorm_backward3(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
#if defined(ENABLE_BF16)
|
|
case 4:
|
|
layernorm_backward4(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
#endif
|
|
case 5:
|
|
layernorm_backward5(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
case 6:
|
|
layernorm_backward6(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
case 7:
|
|
layernorm_backward7(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
case 8:
|
|
layernorm_backward8(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size);
|
|
break;
|
|
default:
|
|
printf("Invalid kernel number\n");
|
|
exit(1);
|
|
}
|
|
cudaCheck(cudaGetLastError());
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
int main(int argc, char **argv) {
|
|
setup_main();
|
|
|
|
int B = 8;
|
|
int T = 1024;
|
|
int C = 768;
|
|
|
|
// first do the forward pass in CPU
|
|
float* out = (float*)malloc(B * T * C * sizeof(float));
|
|
float* mean = (float*)malloc(B * T * sizeof(float));
|
|
float* rstd = (float*)malloc(B * T * sizeof(float));
|
|
float* inp = make_random_float(B * T * C);
|
|
float* weight = make_random_float(C);
|
|
float* bias = make_random_float(C);
|
|
layernorm_forward_cpu(out, mean, rstd, inp, weight, bias, B, T, C);
|
|
|
|
// now do the backward pass, again on CPU
|
|
float *dout = make_random_float(B * T * C);
|
|
float *dinp = make_zeros_float(B * T * C);
|
|
float *dweight = make_zeros_float(C);
|
|
float *dbias = make_zeros_float(C);
|
|
layernorm_backward_cpu(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
|
|
|
|
// the above calculations act as the reference
|
|
// now let's do the same on the GPU
|
|
|
|
// read kernel_num from command line
|
|
int kernel_num = 2;
|
|
if (argc > 1) {
|
|
kernel_num = atoi(argv[1]);
|
|
}
|
|
printf("Using kernel %d\n", kernel_num);
|
|
|
|
// move all the variables we need for backward pass onto the GPU
|
|
floatX* d_dinp;
|
|
floatX* d_dweight;
|
|
floatX* d_dbias;
|
|
floatX* d_dout;
|
|
floatX* d_inp;
|
|
floatX* d_weight;
|
|
floatX* d_mean;
|
|
floatX* d_rstd;
|
|
float* d_scratch;
|
|
cudaCheck(cudaMalloc(&d_dinp, B * T * C * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_dweight, C * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_dbias, C * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_dout, B * T * C * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_weight, C * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_mean, B * T * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_rstd, B * T * sizeof(floatX)));
|
|
cudaCheck(cudaMalloc(&d_scratch, cuda_num_SMs * (2 * C + 1) * sizeof(float)));
|
|
// copy over the "inputs" to the backward call
|
|
cudaCheck(memcpy_convert(d_dout, dout, B * T * C));
|
|
cudaCheck(memcpy_convert(d_inp, inp, B * T * C));
|
|
cudaCheck(memcpy_convert(d_weight, weight, C));
|
|
cudaCheck(memcpy_convert(d_mean, mean, B * T));
|
|
cudaCheck(memcpy_convert(d_rstd, rstd, B * T));
|
|
|
|
// launch the kernel
|
|
int block_sizes[] = {32, 64, 128, 256, 512, 768, 1024};
|
|
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
|
|
int block_size = block_sizes[j];
|
|
// init the "outputs" of the backward call to zeros
|
|
cudaCheck(cudaMemset(d_dinp, 0, B * T * C * sizeof(floatX)));
|
|
cudaCheck(cudaMemset(d_dweight, 0, C * sizeof(floatX)));
|
|
cudaCheck(cudaMemset(d_dbias, 0, C * sizeof(floatX)));
|
|
|
|
layernorm_backward(kernel_num, d_dinp, d_dweight, d_dbias, d_scratch, d_dout, d_inp, d_weight, d_mean, d_rstd,
|
|
B, T, C, block_size);
|
|
|
|
// check the correctness of the kernel
|
|
float error_threshold_dinp = sizeof(floatX) == 4 ? 1e-3f : 1e-1f; // allow larger errors for BF16/FP16
|
|
float error_threshold_dparams = sizeof(floatX) == 4 ? 1e-3f : 5e-1f; // much, much larger...
|
|
printf("Checking correctness...\n");
|
|
printf("dinp:\n");
|
|
validate_result(d_dinp, dinp, "dinp", B * T * C, error_threshold_dinp);
|
|
printf("dweight:\n");
|
|
validate_result(d_dweight, dweight, "dweight", C, error_threshold_dparams);
|
|
printf("dbias:\n");
|
|
validate_result(d_dbias, dbias, "dbias", C, error_threshold_dparams);
|
|
|
|
printf("All results match for block_size=%d.\n\n", block_size);
|
|
}
|
|
|
|
// now time the kernel
|
|
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
|
|
int block_size = block_sizes[j];
|
|
int repeat_times = 100;
|
|
float elapsed_time = benchmark_kernel(repeat_times, layernorm_backward, kernel_num,
|
|
d_dinp, d_dweight, d_dbias, d_scratch, d_dout, d_inp, d_weight, d_mean, d_rstd,
|
|
B, T, C, block_size);
|
|
printf("block_size %4d time %.4f ms\n", block_size, elapsed_time);
|
|
}
|
|
|
|
// cleanups
|
|
free(out);
|
|
free(mean);
|
|
free(rstd);
|
|
free(inp);
|
|
free(weight);
|
|
free(bias);
|
|
free(dout);
|
|
free(dinp);
|
|
free(dweight);
|
|
free(dbias);
|
|
cudaCheck(cudaFree(d_dinp));
|
|
cudaCheck(cudaFree(d_dweight));
|
|
cudaCheck(cudaFree(d_dbias));
|
|
cudaCheck(cudaFree(d_dout));
|
|
cudaCheck(cudaFree(d_inp));
|
|
cudaCheck(cudaFree(d_weight));
|
|
cudaCheck(cudaFree(d_mean));
|
|
cudaCheck(cudaFree(d_rstd));
|
|
cudaCheck(cudaFree(d_scratch));
|
|
return 0;
|
|
}
|