aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-27 20:25:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge branch 'karpathy:master' into feature/dataloader_test

Browse source
This commit is contained in:
rosslwheeler 2024-06-17 18:18:23 -07:00 • committed by GitHub
commit 4e0221de0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 217 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

18
dev/cuda/fused_residual_forward.cu
View file

@ -133,7 +133,7 @@ __global__ void fused_residual_forward2(floatX* residual, floatX* normed, floatX
for(int c = 0; c < C; ++c) {
float out = (float)inp1[c] + (float)inp2[c];
m += out;
residual[c] = out;
residual[c] = (floatX)out;
}
m = m / C;
@ -149,11 +149,11 @@ __global__ void fused_residual_forward2(floatX* residual, floatX* normed, floatX
for (int c = 0; c < C; c++) {
float n = (s * ((float)residual[c] - m)); // normalized output
float o = n * (float)weight[c] + (float)bias[c]; // scale and shift it
normed[c] = o; // write
normed[c] = (floatX)o; // write
}
// cache the mean and rstd for the backward pass later
mean[idx] = m;
rstd[idx] = s;
mean[idx] = (floatX)m;
rstd[idx] = (floatX)s;
}
// handle one token per warp for coalesced access
@ -232,7 +232,7 @@ __global__ void fused_residual_forward_kernel4(floatX* residual, floatX* normed,
const x128 in2 = load128cs(inp2 + c);
x128 out;
for(int k = 0; k < x128::size; ++k) {
out[k] = (float)in1[k] + (float)in2[k];
out[k] = (floatX)((float)in1[k] + (float)in2[k]);
sum += (float)out[k];
sum_sq += (float)out[k] * (float)out[k];
}
@ -309,7 +309,7 @@ __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed,
const x128 in2 = load128cs(inp2 + c);
x128 out;
for(int k = 0; k < x128::size; ++k) {
out[k] = (float)in1[k] + (float)in2[k];
out[k] = (floatX)((float)in1[k] + (float)in2[k]);
sum += (float)out[k];
}
store128cs(residual + c, out);
@ -372,8 +372,8 @@ __global__ void fused_residual_forward_kernel6(floatX* residual, floatX* normed,
// weights and biases are shared among all tokens
x128* s_weight = reinterpret_cast<x128*>(params);
x128* s_bias = reinterpret_cast<x128*>(params + C * sizeof(floatX));
// residual output (input to layernorm) is indpendent for each sub-block indicates by threadIdx.z
x128* s_res = reinterpret_cast<x128*>(params + (2 + threadIdx.z) * C * sizeof(floatX) );
// residual output (input to layernorm) is independent for each sub-block indicates by threadIdx.z
x128* s_res = reinterpret_cast<x128*>(params + (2 + threadIdx.z) * C * sizeof(floatX));
// similarly, each sub-block needs its own reduction buffers
float* s_mean = reinterpret_cast<float*>(params + (2 + blockDim.z) * C * sizeof(floatX) + threadIdx.z * 32 * sizeof(float));
float* s_var = reinterpret_cast<float*>(params + (2 + blockDim.z) * C * sizeof(floatX) + 32 * sizeof(float) * (blockDim.z + threadIdx.z));
@ -385,10 +385,10 @@ __global__ void fused_residual_forward_kernel6(floatX* residual, floatX* normed,
s_weight[c / x128::size] = load128(weight + c);
s_bias[c / x128::size] = load128(bias + c);
}
// the block-level reductions will cause sync before the first time we read these
// => no syncthreads needed here
// loop over all tokens
for(int tidx = blockIdx.x * blockDim.z + threadIdx.z; tidx < N; tidx += gridDim.x * blockDim.z) {
// adjust pointers to current token

106
dev/cuda/layernorm_forward.cu
View file

@ -29,6 +29,8 @@ verstion 5 allocates blocks per row instead of warps per row, same alg as 4 othe
#include <cooperative_groups/reduce.h>
#include "common.h"
#define WARP_SIZE 32
// ----------------------------------------------------------------------------
// CPU code reference
@ -337,6 +339,82 @@ __global__ void layernorm_forward_kernel5(float* __restrict__ out, float* __rest
}
}
// Inspired by `fused_residual_forward_kernel5` in fused_residual_forward.cu
__global__ void layernorm_forward_kernel6(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd,
const float* __restrict__ inp, const float* __restrict__ weight,
const float* __restrict__ bias, int N, int C) {
assert(blockDim.x == WARP_SIZE);
// load weights and biases into shared memory
// do this before we allow any threads to exit!
extern __shared__ char params[];
// load128/store128 sometimes generated multiple instructions when the types here were floatX*, so
// let's keep everything as x128
x128* s_weight = reinterpret_cast<x128*>(params);
x128* s_bias = reinterpret_cast<x128*>(params) + (C / x128::size);
x128* s_in = reinterpret_cast<x128*>(params) + ((2 + threadIdx.y) * C / x128::size);
int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size;
for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) {
s_weight[i/x128::size] = load128(weight + i);
s_bias[i/x128::size] = load128(bias + i);
}
__syncthreads();
int idx = blockIdx.x * blockDim.y + threadIdx.y;
if(idx >= N) { return; } // guard
// adjust pointers to current token
inp += idx * C;
out += idx * C;
const float eps = 1e-5f;
float sum = 0.0f;
for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) {
const x128 in_data = load128cs(inp + c);
for(int k = 0; k < x128::size; ++k) {
sum += (float)in_data[k];
}
s_in[c / x128::size] = in_data;
}
sum = warpReduceSum(sum);
float m = sum / C;
float v = 0.f;
for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) {
const x128 in_data = s_in[c / x128::size];
for(int k = 0; k < x128::size; ++k) {
v += ((float)in_data[k] - m) * ((float)in_data[k] - m);
}
}
v = warpReduceSum(v) / C;
float s = rsqrtf(v + eps);
for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) {
const x128 in_data = s_in[c / x128::size];
const x128 w = s_weight[c / x128::size];
const x128 b = s_bias[c / x128::size];
x128 out_data;
for(int k = 0; k < x128::size; ++k) {
float n = s * ((float)in_data[k] - m); // normalized output
float o = n * (float)w[k] + (float)b[k]; // scale and shift it
out_data[k] = o;
}
store128cs(out + c, out_data);
}
// cache the mean and rstd for the backward pass later
if(threadIdx.x == 0 && mean != nullptr) {
__stcs(mean + idx, m);
}
// store the rstd, no need to cache it
if(threadIdx.x == 0 && rstd != nullptr) {
__stcs(rstd + idx, s);
}
}
// ----------------------------------------------------------------------------
// kernel launcher
@ -401,6 +479,31 @@ void layernorm_forward5(float* out, float* mean, float* rstd,
cudaCheck(cudaGetLastError());
}
void layernorm_forward6(float* out, float* mean, float* rstd,
const float* inp, const float* weight, const float* bias,
int B, int T, int C,
int block_size) {
assert(block_size % 32 == 0);
const int N = B * T;
int block_y = block_size / WARP_SIZE;
const int grid_size = ceil_div(N, block_y);
size_t smem = (2 + block_y) * C * sizeof(float);
// in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute
// this may fail, in which case we fall back to the smem free implementation.
cudaCheck(cudaGetLastError());
auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
cudaGetLastError();
if (status == cudaSuccess) {
layernorm_forward_kernel6<<<grid_size, dim3(32, block_y), smem>>>(out, mean, rstd, inp, weight, bias, N, C);
} else {
const int grid_size = N;
// fall back to the version without shared memory
layernorm_forward_kernel5<<<grid_size, block_size>>>(out, mean, rstd, inp, weight, bias, N, C);
}
cudaCheck(cudaGetLastError());
}
// kernel version dispatch
void layernorm_forward(int kernel_num,
float* out, float* mean, float* rstd,
@ -423,6 +526,9 @@ void layernorm_forward(int kernel_num,
case 5:
layernorm_forward5(out, mean, rstd, inp, weight, bias, B, T, C, block_size);
break;
case 6:
layernorm_forward6(out, mean, rstd, inp, weight, bias, B, T, C, block_size);
break;
default:
printf("Invalid kernel number\n");
exit(1);

9
llmc/dataloader.h
View file

@ -110,7 +110,8 @@ void prepare_intra_shard_indices_(DataLoader *loader) {
free(loader->intra_shard_indices);
}
loader->intra_shard_indices = (int*)mallocCheck(loader->shard_num_samples * sizeof(int));
random_permutation_with_init(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng, 1);
init_identity_permutation(loader->intra_shard_indices, loader->shard_num_samples);
random_permutation(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng);
}
void dataloader_reset(DataLoader *loader) {
@ -118,7 +119,7 @@ void dataloader_reset(DataLoader *loader) {
loader->current_sample_idx = 0;
if (loader->should_shuffle) { // shuffle the shards
random_permutation_with_init(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng, 0);
random_permutation(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng);
}
dataloader_load_shard_(loader, loader->current_shard_idx);
@ -178,9 +179,7 @@ void dataloader_init(DataLoader *loader,
manual_seed(&shuffle_rng, 42 + process_rank);
loader->shuffle_rng = shuffle_rng;
loader->shard_indices = (int*)mallocCheck(loader->glob_result.gl_pathc * sizeof(int));
for (int i = 0; i < loader->glob_result.gl_pathc; i++) {
loader->shard_indices[i] = i; // start with identity permutation
}
init_identity_permutation(loader->shard_indices, loader->glob_result.gl_pathc);
loader->intra_shard_indices = NULL; // dynamically allocated allowing different shard sizes
}

96
llmc/layernorm.cuh
View file

@ -64,6 +64,81 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __re
}
}
__global__ void layernorm_forward_kernel6(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd,
const floatX* __restrict__ inp, const floatX* __restrict__ weight,
const floatX* __restrict__ bias, int N, int C) {
assert(blockDim.x == WARP_SIZE);
// load weights and biases into shared memory
// do this before we allow any threads to exit!
extern __shared__ char* params[];
// load128/store128 sometimes generated multiple instructions when the types here were floatX*, so
// let's keep everything as x128
x128* s_weight = reinterpret_cast<x128*>(params);
x128* s_bias = reinterpret_cast<x128*>(params) + (C / x128::size);
x128* s_in = reinterpret_cast<x128*>(params) + ((2 + threadIdx.y) * C / x128::size);
int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size;
for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) {
s_weight[i/x128::size] = load128(weight + i);
s_bias[i/x128::size] = load128(bias + i);
}
__syncthreads();
int idx = blockIdx.x * blockDim.y + threadIdx.y;
if(idx >= N) { return; } // guard
// adjust pointers to current token
inp += idx * C;
out += idx * C;
const float eps = 1e-5f;
float sum = 0.0f;
for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) {
const x128 in_data = load128cs(inp + c);
for(int k = 0; k < x128::size; ++k) {
sum += (float)in_data[k];
}
s_in[c / x128::size] = in_data;
}
sum = warpReduceSum(sum);
float m = sum / C;
float v = 0.f;
for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) {
const x128 in_data = s_in[c / x128::size];
for(int k = 0; k < x128::size; ++k) {
v += ((float)in_data[k] - m) * ((float)in_data[k] - m);
}
}
v = warpReduceSum(v) / C;
float s = rsqrtf(v + eps);
for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) {
const x128 in_data = s_in[c / x128::size];
const x128 w = s_weight[c / x128::size];
const x128 b = s_bias[c / x128::size];
x128 out_data;
for(int k = 0; k < x128::size; ++k) {
float n = s * ((float)in_data[k] - m); // normalized output
float o = n * (float)w[k] + (float)b[k]; // scale and shift it
out_data[k] = (floatX)o;
}
store128cs(out + c, out_data);
}
// cache the mean and rstd for the backward pass later
if(threadIdx.x == 0 && mean != nullptr) {
__stcs(mean + idx, (floatX)m);
}
// store the rstd, no need to cache it
if(threadIdx.x == 0 && rstd != nullptr) {
__stcs(rstd + idx, (floatX)s);
}
}
__global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd,
const floatX* inp1, const floatX* inp2,
const floatX* weight, const floatX* bias,
@ -354,14 +429,29 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with
// ----------------------------------------------------------------------------
// kernel launchers
// similar to `fused_residual_forward5`
void layernorm_forward(floatX* out, floatX* mean, floatX* rstd,
floatX* inp, const floatX* weight, const floatX* bias,
int B, int T, int C, cudaStream_t stream) {
NVTX_RANGE_FN();
const int block_size = 512;
const int block_size = 256;
int block_y = block_size / WARP_SIZE;
const int N = B * T;
const int grid_size = CEIL_DIV(N * WARP_SIZE, block_size);
layernorm_forward_kernel3<<<grid_size, block_size, 0, stream>>>(out, mean, rstd, inp, weight, bias, N, C);
const int grid_size = CEIL_DIV(N, block_y);
size_t smem = (2 + block_y) * C * sizeof(floatX);
// in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute
// this may fail, in which case we fall back to the smem free implementation.
cudaCheck(cudaGetLastError());
auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
cudaGetLastError();
if (status == cudaSuccess) {
layernorm_forward_kernel6<<<grid_size, dim3(WARP_SIZE, block_y), smem, stream>>>(out, mean, rstd, inp, weight, bias, N, C);
} else {
// fall back to the version without shared memory
const int grid_size_fb = CEIL_DIV(N * WARP_SIZE, block_size);
layernorm_forward_kernel3<<<grid_size_fb, block_size, 0, stream>>>(out, mean, rstd, inp, weight, bias, N, C);
}
cudaCheck(cudaGetLastError());
}

10
llmc/rand.h
View file

@ -220,13 +220,13 @@ void normal_(float* data, unsigned int numel, float mean, float std, mt19937_sta
}
}
void random_permutation_with_init(int* data, int numel, mt19937_state* state, int should_init) {
if (should_init) {
for (int i = 0; i < numel; i++) {
data[i] = i;
}
void init_identity_permutation(int *data, int numel) {
for (int i = 0; i < numel; i++) {
data[i] = i;
}
}
void random_permutation(int* data, int numel, mt19937_state* state) {
for (int i = numel - 1; i > 0; i--) {
// pick an index j in [0, i] with equal probability
int j = randint32(state) % (i + 1);