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

update norm computation: multi-process reduce on GPU

Browse source
This commit is contained in:
Erik Schultheis 2024-06-24 23:57:42 +03:00
parent 8d994782cf
commit 75d4610442
2 changed files with 7 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

7
llmc/global_norm.cuh
View file

@ -87,10 +87,3 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st
global_norm_squared_kernel<<<dim3(gx, gy), block_size, 0, stream>>>(out, values, count, stride);
cudaCheck(cudaGetLastError());
}
void global_norm_squared_aggregate(float* out, int max_num_block_sums, cudaStream_t stream) {
assert(max_num_block_sums > 0 && max_num_block_sums < 1024); // we need to accumulate the block sums in a single block
// important to use 1024 here for determinism, otherwise blockreduce might introduce errors
global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, max_num_block_sums);
cudaCheck(cudaGetLastError());
}

13
train_gpt2.cu
View file

@ -1015,18 +1015,19 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
max_num_block_sums, is_first_pass, main_stream);
}
}
global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream);
cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost));
// further sum the (partial) squared norm across all GPUs (see comment ^1 above)
grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu, multi_gpu_config);
global_sum_deterministic(grad_norm_squared, grad_norm_squared, max_num_block_sums, main_stream);
#if MULTI_GPU
// further sum the (partial) squared norm across all GPUs
ncclCheck(ncclAllReduce(grad_norm_squared, grad_norm_squared, sizeof(float), ncclFloat, ncclAvg, multi_gpu_config->nccl_comm, main_stream));
#endif
} else {
// in regular DDP, backward 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, 0, 1, max_num_block_sums, true, main_stream);
global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream);
cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost));
global_sum_deterministic(grad_norm_squared, grad_norm_squared, max_num_block_sums, main_stream);
}
cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost));
if(!isfinite(grad_norm_squared_cpu)) {
// may happen due to some issue (e.g. overflow?)
// TODO: later may want to keep a global counter of instabilities like this