mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-28 20:35:09 -04:00
Merge branch 'karpathy:master' into FLT_MAX-fix
This commit is contained in:
Vyom Sharma
GitHub
commit
ee4b056ddf
12 changed files with 508 additions and 376 deletions
2
Makefile
2
Makefile
|
|
@ -267,4 +267,4 @@ profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN)
|
|||
|
||||
clean:
|
||||
$(REMOVE_FILES) $(TARGETS)
|
||||
$(REMOVE_BUILD_OBJECT_FILES)
|
||||
$(REMOVE_BUILD_OBJECT_FILES)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ LLMs in simple, pure C/CUDA with no need for 245MB of PyTorch or 107MB of cPytho
|
|||
|
||||
The best introduction to the llm.c repo today is reproducing the GPT-2 (124M) model. [Discussion #481](https://github.com/karpathy/llm.c/discussions/481) steps through this in detail. We can reproduce other models from the GPT-2 and GPT-3 series in both llm.c and in the parallel implementation of PyTorch. Have a look at the [scripts README](scripts/README.md).
|
||||
|
||||
debugging tip: when you run the `make` command to build the binary, modify it by replacing `-O3` with `-g` so you can step through the code in your favorite IDE (e.g. vscode).
|
||||
|
||||
## quick start (1 GPU, fp32 only)
|
||||
|
||||
If you won't be training on multiple nodes, aren't interested in mixed precision, and are interested in learning CUDA, the fp32 (legacy) files might be of interest to you. These are files that were "checkpointed" early in the history of llm.c and frozen in time. They are simpler, more portable, and possibly easier to understand. Run the 1 GPU, fp32 code like this:
|
||||
|
|
|
|||
|
|
@ -1,46 +1,53 @@
|
|||
"""
|
||||
Script for running benchmarks on the Modal platform.
|
||||
This is useful for folks who do not have access to expensive GPUs locally.
|
||||
|
||||
Example usage:
|
||||
Example usage for cuda kernels:
|
||||
GPU_MEM=80 modal run benchmark_on_modal.py \
|
||||
--compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" \
|
||||
--run-command "./attention_forward 1"
|
||||
OR if you want to use cuDNN etc.
|
||||
|
||||
This will mount the contents of the current directory to the remote container on modal,
|
||||
compile the `attention_forward.cu` file with `nvcc`, and run the resulting binary on a A100 GPU with 80GB of memory.
|
||||
|
||||
For training the gpt2 model with cuDNN use:
|
||||
GPU_MEM=80 modal run dev/cuda/benchmark_on_modal.py \
|
||||
--compile-command "make train_gpt2cu USE_CUDNN=1"
|
||||
--run-command "./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin -v 250 -s 250 -g 144 -f shakespeare.log -b 4"
|
||||
|
||||
|
||||
For profiling using nsight system:
|
||||
GPU_MEM=80 modal run dev/cuda/benchmark_on_modal.py \
|
||||
--compile-command "make train_gpt2cu USE_CUDNN=1" \
|
||||
--run-command "nsys profile --cuda-graph-trace=graph --python-backtrace=cuda --cuda-memory-usage=true \
|
||||
./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
|
||||
-j dev/data/tinyshakespeare/tiny_shakespeare_val.bin -v 250 -s 250 -g 144 -f shakespeare.log -b 4"
|
||||
|
||||
For more nsys profiling specifics and command options, take a look at: https://docs.nvidia.com/nsight-systems/2024.2/UserGuide/
|
||||
-> To profile the report using a GUI, download NVIDIA NSight System GUI version (this software can run on all OS, so you download it locally)
|
||||
|
||||
NOTE: Currently there is a bug in the profiling using nsight system which produces a unrecognized GPU UUId error on the command line but it
|
||||
does not actually interfere with the model training and validation. The report (that you download) is still generated and can be viewed from Nsight Systems
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import sys
|
||||
import datetime
|
||||
|
||||
import modal
|
||||
from modal import Image, Stub
|
||||
|
||||
GPU_NAME_TO_MODAL_CLASS_MAP = {
|
||||
"H100": modal.gpu.H100,
|
||||
"A100": modal.gpu.A100,
|
||||
"A10G": modal.gpu.A10G,
|
||||
}
|
||||
|
||||
N_GPUS = int(os.environ.get("N_GPUS", 1))
|
||||
GPU_MEM = int(os.environ.get("GPU_MEM", 40))
|
||||
GPU_NAME = os.environ.get("GPU_NAME", "A100")
|
||||
GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, size=str(GPU_MEM)+'GB')
|
||||
GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, size=str(GPU_MEM) + 'GB')
|
||||
|
||||
APP_NAME = "llm.c benchmark run"
|
||||
|
||||
# We don't actually need to use the Axolotl image here, but it's reliable
|
||||
AXOLOTL_REGISTRY_SHA = (
|
||||
"d5b941ba2293534c01c23202c8fc459fd2a169871fa5e6c45cb00f363d474b6a"
|
||||
)
|
||||
axolotl_image = (
|
||||
Image.from_registry(f"winglian/axolotl@sha256:{AXOLOTL_REGISTRY_SHA}")
|
||||
.run_commands(
|
||||
"git clone https://github.com/OpenAccess-AI-Collective/axolotl /root/axolotl",
|
||||
"cd /root/axolotl && git checkout v0.4.0",
|
||||
)
|
||||
image = (
|
||||
Image.from_registry("totallyvyom/cuda-env:latest-2")
|
||||
.pip_install("huggingface_hub==0.20.3", "hf-transfer==0.1.5")
|
||||
.env(
|
||||
dict(
|
||||
|
|
@ -49,34 +56,66 @@ axolotl_image = (
|
|||
TQDM_DISABLE="true",
|
||||
)
|
||||
)
|
||||
.run_commands(
|
||||
"wget -q https://github.com/Kitware/CMake/releases/download/v3.28.1/cmake-3.28.1-Linux-x86_64.sh",
|
||||
"bash cmake-3.28.1-Linux-x86_64.sh --skip-license --prefix=/usr/local",
|
||||
"rm cmake-3.28.1-Linux-x86_64.sh",
|
||||
"ln -s /usr/local/bin/cmake /usr/bin/cmake",)
|
||||
.run_commands(
|
||||
"apt-get install -y --allow-change-held-packages libcudnn8 libcudnn8-dev",
|
||||
"apt-get install -y openmpi-bin openmpi-doc libopenmpi-dev kmod sudo",
|
||||
"git clone https://github.com/NVIDIA/cudnn-frontend.git /root/cudnn-frontend",
|
||||
"cd /root/cudnn-frontend && mkdir build && cd build && cmake .. && make"
|
||||
)
|
||||
.run_commands(
|
||||
"wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin && \
|
||||
mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600 && \
|
||||
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \
|
||||
add-apt-repository \"deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /\" && \
|
||||
apt-get update"
|
||||
).run_commands(
|
||||
"apt-get install -y nsight-systems-2023.3.3"
|
||||
)
|
||||
)
|
||||
|
||||
stub = modal.App(APP_NAME)
|
||||
|
||||
|
||||
def execute_command(command: str):
|
||||
command_args = command.split(" ")
|
||||
print(f"{command_args = }")
|
||||
subprocess.run(command_args, stdout=sys.stdout, stderr=subprocess.STDOUT)
|
||||
|
||||
|
||||
@stub.function(
|
||||
gpu=GPU_CONFIG,
|
||||
image=axolotl_image,
|
||||
image=image,
|
||||
allow_concurrent_inputs=4,
|
||||
container_idle_timeout=900,
|
||||
# This copies everything in this folder to the remote root folder
|
||||
mounts=[modal.Mount.from_local_dir("./", remote_path="/root/")]
|
||||
mounts=[modal.Mount.from_local_dir("./", remote_path="/root/")],
|
||||
# Instead of 'cuda-env' put your volume name that you create from 'modal volume create {volume-name}'
|
||||
# This enables the profiling reports to be saved on the volume that you can download by using:
|
||||
# 'modal volume get {volume-name} {/output_file_name}
|
||||
# For example right now, when profiling using this command "nsys profile --trace=cuda,nvtx --cuda-graph-trace=graph --python-backtrace=cuda --cuda-memory-usage=true" you would get your report
|
||||
# using in a directory in your volume, where the name contains the timestamp unique id.
|
||||
# This script will generate a "report1_{timestamp} folder in volume"
|
||||
# and you can download it with 'modal volume get {volume-name} report1_{timestamp}
|
||||
volumes={"/cuda-env": modal.Volume.from_name("cuda-env")},
|
||||
)
|
||||
def run_benchmark(compile_command: str, run_command: str):
|
||||
execute_command("pwd")
|
||||
execute_command("ls")
|
||||
execute_command(compile_command)
|
||||
execute_command(run_command)
|
||||
return None
|
||||
# Use this section if you want to profile using nsight system and install the reports on your volume to be locally downloaded
|
||||
timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
|
||||
|
||||
execute_command("mkdir report1_" + timestamp)
|
||||
execute_command("mv /root/report1.nsys-rep /root/report1_" + timestamp + "/")
|
||||
execute_command("mv /root/report1.qdstrm /root/report1_" + timestamp + "/")
|
||||
execute_command("mv /root/report1_" + timestamp + "/" + " /cuda-env/")
|
||||
|
||||
return None
|
||||
|
||||
@stub.local_entrypoint()
|
||||
def inference_main(compile_command: str, run_command: str):
|
||||
results = run_benchmark.remote(compile_command, run_command)
|
||||
return results
|
||||
return results
|
||||
|
|
@ -16,9 +16,9 @@ __device__ float lerp(float start, float end, float weight) {
|
|||
}
|
||||
|
||||
template <typename Tp, typename Tg>
|
||||
__global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters,
|
||||
float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay,
|
||||
float grad_scale, unsigned int seed) {
|
||||
__device__ 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, float beta1_correction, float beta2_correction, float eps, float weight_decay,
|
||||
float grad_scale, unsigned int seed) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= num_parameters) { return; } // guard
|
||||
|
||||
|
|
@ -41,24 +41,39 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg
|
|||
// update our low precision version of the parameters using stochastic rounding
|
||||
// this will be used in the next forward pass
|
||||
// TODO: simply doing `params_memory[i] = (floatX)param;` breaks everything (why?)
|
||||
unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed);
|
||||
unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x + blockDim.y * gridDim.y, seed);
|
||||
stochastic_rounding(param, ¶ms_memory[idx], random);
|
||||
// write the full, float version of the param into our master copy, if we maintain one
|
||||
// this will be used in the next update
|
||||
if (master_params_memory != NULL) { master_params_memory[idx] = param; }
|
||||
}
|
||||
|
||||
template <typename Tp, typename Tg>
|
||||
__global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters,
|
||||
ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride,
|
||||
float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay,
|
||||
float grad_scale, unsigned int seed) {
|
||||
adamw_update(params_memory + blockIdx.y * w_stride,
|
||||
master_params_memory ? master_params_memory + blockIdx.y * s_stride : NULL,
|
||||
grads_memory + blockIdx.y * g_stride,
|
||||
m_memory + blockIdx.y * s_stride,
|
||||
v_memory + blockIdx.y * s_stride,
|
||||
num_parameters, learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, grad_scale,
|
||||
seed
|
||||
);
|
||||
}
|
||||
|
||||
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,
|
||||
ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride, int num_slices, 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,
|
||||
adamw_kernel3<<<dim3(num_blocks, num_slices), block_size, 0, stream>>>(params_memory, master_params_memory, grads_memory,
|
||||
m_memory, v_memory, num_parameters, w_stride, g_stride, s_stride,
|
||||
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay,
|
||||
grad_scale, seed);
|
||||
cudaCheck(cudaGetLastError());
|
||||
|
|
|
|||
|
|
@ -149,8 +149,8 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons
|
|||
}
|
||||
}
|
||||
|
||||
__global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const floatX* datt, const floatX* att,
|
||||
int B, int T, int C, float scale) {
|
||||
__global__ void softmax_autoregressive_backward_inplace_kernel(floatX* datt, const floatX* att,
|
||||
int B, int T, int C, float scale) {
|
||||
constexpr const int BlockSize = 256;
|
||||
constexpr int T_per_block = 4;
|
||||
|
||||
|
|
@ -160,14 +160,13 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl
|
|||
|
||||
att += idx * T * T;
|
||||
datt += idx * T * T;
|
||||
dpreatt += idx * T * T;
|
||||
|
||||
for(int to = 0; to < T_per_block; ++to) {
|
||||
int t = t0 - to;
|
||||
if(t < 0) return;
|
||||
const floatX* att_bth = att + t * T;
|
||||
const floatX* datt_bth = datt + t * T;
|
||||
floatX* dpreatt_bth = dpreatt + t * T;
|
||||
floatX* dpreatt_bth = datt + t * T;
|
||||
|
||||
float local_sum = 0;
|
||||
for (int t2 = threadIdx.x; t2 <= t; t2 += BlockSize) {
|
||||
|
|
@ -176,11 +175,16 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl
|
|||
|
||||
local_sum = blockReduce<warpReduceSum>(local_sum);
|
||||
|
||||
for (int t3 = threadIdx.x; t3 <= t; t3 += BlockSize) {
|
||||
for (int t3 = threadIdx.x; t3 < T; t3 += BlockSize) {
|
||||
// don't touch the cache. Some parts will still be here from the previous loop, and
|
||||
// we want to exploit those.
|
||||
float acc = (float)__ldcs(att_bth + t3) * ((float)__ldcs(datt_bth + t3) - local_sum);
|
||||
__stcs(dpreatt_bth + t3, (floatX)(scale * acc));
|
||||
if(t3 <= t) {
|
||||
float acc = (float) __ldcs(att_bth + t3) * ((float) __ldcs(datt_bth + t3) - local_sum);
|
||||
__stcs(dpreatt_bth + t3, (floatX) (scale * acc));
|
||||
} else {
|
||||
// explicitly set non-causal elements to zero
|
||||
__stcs(dpreatt_bth + t3, (floatX)0.f);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -200,7 +204,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
// inp is (B, T, 3C) QKV
|
||||
// preatt, att are (B, NH, T, T)
|
||||
// output is (B, T, C)
|
||||
int HS = C / NH; // head size
|
||||
const int HS = C / NH; // head size
|
||||
|
||||
// permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)
|
||||
floatX *q, *k, *v;
|
||||
|
|
@ -223,7 +227,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
|
||||
|
||||
// multiply all elements of preatt elementwise by scale
|
||||
float scale = 1.0 / sqrtf(HS);
|
||||
float scale = 1.f / sqrtf(HS);
|
||||
int grid_size = CEIL_DIV(B * NH * T * WARP_SIZE, block_size);
|
||||
softmax_forward_kernel5<<<grid_size, block_size, 0, stream>>>(att, scale, preatt, B * NH, T);
|
||||
|
||||
|
|
@ -247,13 +251,13 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
|
|||
|
||||
// the sequence of transformations in this compound op is:
|
||||
// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C)
|
||||
void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* datt, floatX* scratch,
|
||||
void attention_backward(floatX* dinp, floatX* dqkvr, floatX* datt, floatX* scratch,
|
||||
const floatX* dout,
|
||||
const floatX* qkvr, const floatX* att,
|
||||
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
|
||||
const int HS = C / NH; // head size
|
||||
const float alpha = 1.0f, beta = 0.0f;
|
||||
|
||||
// unpack convenience pointers into q, k, v
|
||||
|
|
@ -279,10 +283,10 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da
|
|||
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha,
|
||||
scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, &beta,
|
||||
dv, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
|
||||
// backward into preatt
|
||||
int hs = C / NH; // head size
|
||||
float scale = 1.0f / sqrtf(hs);
|
||||
softmax_autoregressive_backward_kernel<<<dim3(T / 4, B * NH), 256, 0, stream>>>(dpreatt, datt, att, B, T, C, scale);
|
||||
const float scale = 1.0f / sqrtf((float)HS);
|
||||
// backward into preatt. this is an in-place operation; datt turns into dpreatt here
|
||||
softmax_autoregressive_backward_inplace_kernel<<<dim3(T / 4, B * NH), 256>>>(datt, att, B, T, C, scale);
|
||||
const floatX* dpreatt = datt;
|
||||
// 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,
|
||||
|
|
|
|||
|
|
@ -102,11 +102,11 @@ __device__ float cast_value<float, __nv_bfloat16>(__nv_bfloat16 val) {
|
|||
}
|
||||
|
||||
template<typename Td, typename Ts>
|
||||
__global__ void copy_and_cast_kernel(Td* dst, const Ts* src, size_t n) {
|
||||
__global__ void copy_and_cast_kernel(Td* dst, const Ts* src, size_t n, ptrdiff_t stride_dst, ptrdiff_t stride_src) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
// need to try grid stride looping for more perf later
|
||||
if (idx < n) {
|
||||
dst[idx] = cast_value<Td, Ts>(src[idx]);
|
||||
dst[idx + stride_dst * blockIdx.y] = cast_value<Td, Ts>(src[idx + stride_src * blockIdx.y]);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
Global norm, used in gradient clipping
|
||||
*/
|
||||
#include <assert.h>
|
||||
#include <stddef.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
// llmc internal imports
|
||||
#include "cuda_common.h"
|
||||
#include "cuda_utils.cuh"
|
||||
|
|
@ -10,7 +12,7 @@ Global norm, used in gradient clipping
|
|||
// CUDA kernels
|
||||
|
||||
template<class T>
|
||||
__global__ void global_norm_squared_kernel(float* out, const T* data, size_t count) {
|
||||
__device__ float global_norm_squared_for_range(const T* data, size_t count) {
|
||||
// we want as few atomics as possible, so each block tries to do
|
||||
// the maximum amount of work (so no fixed chunk, but instead iterating
|
||||
// until we run out of data), and then we reduce inside the block
|
||||
|
|
@ -23,8 +25,13 @@ __global__ void global_norm_squared_kernel(float* out, const T* data, size_t cou
|
|||
for(size_t i = index; i < count; i += grid_width) {
|
||||
accumulator += (float)data[i] * (float)data[i];
|
||||
}
|
||||
// warp-level reduce
|
||||
float block_sum = blockReduce<warpReduceSum>(accumulator);
|
||||
// block-level reduce
|
||||
return blockReduce<warpReduceSum>(accumulator);
|
||||
}
|
||||
|
||||
template<class T>
|
||||
__global__ void global_norm_squared_kernel(float* out, const T* data, size_t count, ptrdiff_t stride) {
|
||||
float block_sum = global_norm_squared_for_range(data + blockIdx.y * stride, count);
|
||||
if(threadIdx.x == 0) {
|
||||
atomicAdd(out, block_sum);
|
||||
}
|
||||
|
|
@ -34,7 +41,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, cudaStream_t stream) {
|
||||
void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t stride, int num_slices, bool reset, 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
|
||||
|
|
@ -44,8 +51,12 @@ void global_norm_squared(float* out, const T* values, size_t count, cudaStream_t
|
|||
const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size;
|
||||
assert(grid_size > 0); // gives a better error than letting the call below fail
|
||||
// initialize out with zero
|
||||
cudaCheck(cudaMemsetAsync(out, 0, sizeof(float), stream));
|
||||
global_norm_squared_kernel<<<grid_size, block_size, 0, stream>>>(out, values, count);
|
||||
if(reset) {
|
||||
cudaCheck(cudaMemsetAsync(out, 0, sizeof(float), stream));
|
||||
}
|
||||
const int gx = CEIL_DIV(grid_size, num_slices);
|
||||
const int gy = num_slices;
|
||||
global_norm_squared_kernel<<<dim3(gx, gy), block_size, 0, stream>>>(out, values, count, stride);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
|
|
|
|||
232
llmc/zero.cuh
Normal file
232
llmc/zero.cuh
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
/*
|
||||
Utilities for ZeRO sharding
|
||||
*/
|
||||
|
||||
#ifndef LLMC_ZERO_CUH
|
||||
#define LLMC_ZERO_CUH
|
||||
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
#include <stdio.h>
|
||||
#include <stddef.h>
|
||||
|
||||
#ifdef MULTI_GPU
|
||||
#include <mpi.h>
|
||||
#include <nccl.h>
|
||||
#endif
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Multi-GPU related
|
||||
#ifdef MULTI_GPU
|
||||
|
||||
#if defined(ENABLE_FP32)
|
||||
const ncclDataType_t ncclFloatX = ncclFloat;
|
||||
#elif defined(ENABLE_FP16)
|
||||
const ncclDataType_t ncclFloatX = ncclHalf;
|
||||
#else // Default to bfloat16
|
||||
const ncclDataType_t ncclFloatX = ncclBfloat16;
|
||||
#endif
|
||||
|
||||
void nccl_check(ncclResult_t status, const char *file, int line) {
|
||||
if (status != ncclSuccess) {
|
||||
printf("[NCCL ERROR] at file %s:%d:\n%s\n", file, line, ncclGetErrorString(status));
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
}
|
||||
#define ncclCheck(err) (nccl_check(err, __FILE__, __LINE__))
|
||||
|
||||
void mpi_check(int status, const char *file, int line) {
|
||||
if (status != MPI_SUCCESS) {
|
||||
char mpi_error[4096];
|
||||
int mpi_error_len = 0;
|
||||
assert(MPI_Error_string(status, &mpi_error[0], &mpi_error_len) == MPI_SUCCESS);
|
||||
printf("[MPI ERROR] at file %s:%d:\n%.*s\n", file, line, mpi_error_len, mpi_error);
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
}
|
||||
#define mpiCheck(err) (mpi_check(err, __FILE__, __LINE__))
|
||||
|
||||
#endif // MULTI_GPU
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// MPI / multi-processing setup
|
||||
|
||||
// Parameters specific to training on multiple GPUs.
|
||||
typedef struct {
|
||||
int process_rank; // Rank of this process among all MPI processes. 0 if no multi-GPU.
|
||||
int num_processes; // Total number of processes. 1 if no multi-GPU.
|
||||
int local_device_idx; // This process GPU index on current machine. 0 if no multi-GPU.
|
||||
|
||||
// Zero Redundancy Optimizer stage - https://fairscale.readthedocs.io/en/stable/deep_dive/oss_sdp_fsdp.html
|
||||
// 0-Disabled
|
||||
// 1-Optimizer State Sharding (OSS)
|
||||
// 2-Optimizer + Gradient State Sharding (SDP)
|
||||
// 3-Optimizer + Gradient + Horizontal Model Sharding (FSDP)
|
||||
int zero_stage;
|
||||
size_t shard_num_parameters;
|
||||
#ifdef MULTI_GPU
|
||||
ncclComm_t nccl_comm; // NCCL communication primitive, used for collective multi-GPU work.
|
||||
cudaStream_t nccl_stream; // CUDA Stream to perform NCCL operations.
|
||||
cudaEvent_t compute_nccl_sync; // Event used to synchronize NCCL with the compute
|
||||
#endif
|
||||
} MultiGpuConfig;
|
||||
|
||||
#ifdef MULTI_GPU
|
||||
// Determine which GPU this process should use.
|
||||
// Processes on the same machines use different GPU indicies. Processes on other machines don't.
|
||||
// Copied from NCCL examples: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html#example-2-one-device-per-process-or-thread
|
||||
int multi_gpu_get_local_device_idx(int process_rank, int num_processes) {
|
||||
char hostname[1024];
|
||||
hostname[1023] = '\0';
|
||||
// All processes on the same machine will share the same hostname.
|
||||
gethostname(hostname, 1023);
|
||||
for (int i=0; i < 1024; i++) {
|
||||
if (hostname[i] == '.') {
|
||||
hostname[i] = '\0';
|
||||
break;
|
||||
}
|
||||
}
|
||||
uint64_t hostname_hash = 5381u;
|
||||
for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5u) + hostname_hash) ^ hostname[c]; }
|
||||
|
||||
// Distribute all hostname hashes to all processes.
|
||||
uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t));
|
||||
all_hostsname_hashes[process_rank] = hostname_hash;
|
||||
mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD));
|
||||
|
||||
// Identify which GPU we need to use.
|
||||
int local_device_idx = 0;
|
||||
for (int current_process = 0; current_process < num_processes; ++current_process) {
|
||||
if (current_process == process_rank) {
|
||||
// Found my gpu, local_device_idx now has my target GPU index.
|
||||
break;
|
||||
}
|
||||
if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) {
|
||||
// This process ID runs on the same machine, but it's not me, skip this GPU
|
||||
local_device_idx++;
|
||||
}
|
||||
}
|
||||
|
||||
free(all_hostsname_hashes);
|
||||
return local_device_idx;
|
||||
}
|
||||
#endif
|
||||
|
||||
MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) {
|
||||
#ifdef MULTI_GPU
|
||||
// Initialize MPI.
|
||||
MultiGpuConfig result;
|
||||
mpiCheck(MPI_Init(argc, argv));
|
||||
mpiCheck(MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank));
|
||||
mpiCheck(MPI_Comm_size(MPI_COMM_WORLD, &result.num_processes));
|
||||
result.local_device_idx = multi_gpu_get_local_device_idx(result.process_rank, result.num_processes);
|
||||
cudaCheck(cudaSetDevice(result.local_device_idx));
|
||||
ncclUniqueId nccl_id;
|
||||
if (result.process_rank == 0) {
|
||||
ncclCheck(ncclGetUniqueId(&nccl_id));
|
||||
}
|
||||
mpiCheck(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
|
||||
ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank));
|
||||
cudaCheck(cudaStreamCreate(&result.nccl_stream));
|
||||
// event without timing for maximum performance
|
||||
cudaCheck(cudaEventCreate(&result.compute_nccl_sync, cudaEventDisableTiming));
|
||||
nvtxNameCudaStreamA(result.nccl_stream, "nccl stream");
|
||||
nvtxNameCudaEventA(result.compute_nccl_sync, "nccl compute sync");
|
||||
return result;
|
||||
#else
|
||||
printf("Multi-GPU support is disabled. Using a single GPU.\n");
|
||||
cudaCheck(cudaSetDevice(0));
|
||||
MultiGpuConfig result;
|
||||
result.process_rank = 0;
|
||||
result.num_processes = 1;
|
||||
result.local_device_idx = 0;
|
||||
return result;
|
||||
#endif
|
||||
}
|
||||
|
||||
void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) {
|
||||
#ifdef MULTI_GPU
|
||||
ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm));
|
||||
cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream));
|
||||
cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync));
|
||||
mpiCheck(MPI_Finalize());
|
||||
#endif
|
||||
}
|
||||
|
||||
void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) {
|
||||
#ifdef MULTI_GPU
|
||||
if (multi_gpu_config->num_processes > 1) {
|
||||
mpiCheck(MPI_Barrier(MPI_COMM_WORLD));
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Offset and size of a tensor shard
|
||||
typedef struct {
|
||||
ptrdiff_t offset;
|
||||
size_t size;
|
||||
} ShardInfo;
|
||||
|
||||
// Get info about sharding for a tensor of elements many numbers
|
||||
ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* multi_gpu_config, int shard_at_stage) {
|
||||
const int nproc = multi_gpu_config->num_processes;
|
||||
if(multi_gpu_config->zero_stage >= shard_at_stage) {
|
||||
if (elements % nproc != 0) {
|
||||
fprintf(stderr, "Number of elements %zu must be a multiple of the number of processes %d\n", elements, nproc);
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
return {(ptrdiff_t) (multi_gpu_config->process_rank * (elements / nproc)), elements / nproc};
|
||||
} else {
|
||||
return {0, elements};
|
||||
}
|
||||
}
|
||||
|
||||
// Block NCCL stream until computations on compute_stream are done, then aggregate multiple pointers in an NCCL group.
|
||||
// This can work either as an all-reduce (i.e., no ZeRo), or a reduce-scatter (ZeRO 1).
|
||||
// The awkward `(&pointers)[N]` syntax ensures we are capturing the parameters as sized arrays, so that it becomes impossible
|
||||
// to call this function if pointers and pointers_sizes do not match.
|
||||
template<int N>
|
||||
void multi_gpu_async_reduce_gradient(
|
||||
floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N],
|
||||
MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) {
|
||||
if (multi_gpu_config->num_processes == 1) {
|
||||
return; // no multi-GPU, just exit.
|
||||
}
|
||||
|
||||
#ifdef MULTI_GPU
|
||||
NVTX_RANGE_FN();
|
||||
// mark an event on the compute stream, and immediately wait on this in the nccl stream
|
||||
// this means that the nccl stream won't start executing before all compute kernels that
|
||||
// have been submitted before this point have finished.
|
||||
// by using an event instead of cudaSyncStream, we avoid having to synchronize the host, and
|
||||
// can enqueue new work to the GPU right away.
|
||||
cudaCheck(cudaEventRecord(multi_gpu_config->compute_nccl_sync, compute_stream));
|
||||
cudaCheck(cudaStreamWaitEvent(multi_gpu_config->nccl_stream, multi_gpu_config->compute_nccl_sync));
|
||||
ncclCheck(ncclGroupStart()); // NCCL group: aggregate all pointers in a single NCCL GPU kernel.
|
||||
for (int i = 0; i < N; ++i) {
|
||||
if(multi_gpu_config->zero_stage == 0) {
|
||||
ncclCheck(ncclAllReduce(
|
||||
pointers[i], pointers[i],
|
||||
pointers_sizes[i],
|
||||
ncclFloatX, ncclAvg,
|
||||
multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream
|
||||
));
|
||||
} else if(multi_gpu_config->zero_stage == 1) {
|
||||
assert(pointers_sizes[i] % multi_gpu_config->num_processes == 0);
|
||||
size_t shard_size = pointers_sizes[i] / multi_gpu_config->num_processes;
|
||||
ptrdiff_t shard_offset = (ptrdiff_t)shard_size * multi_gpu_config->process_rank;
|
||||
ncclCheck(ncclReduceScatter(
|
||||
pointers[i], pointers[i] + shard_offset,
|
||||
shard_size,
|
||||
ncclFloatX, ncclAvg,
|
||||
multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream
|
||||
));
|
||||
}
|
||||
}
|
||||
ncclCheck(ncclGroupEnd());
|
||||
#endif
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
|
|
@ -54,7 +54,7 @@ int main(int argc, char *argv[]) {
|
|||
// do a training step
|
||||
gpt2_forward(&model, x, y, B, T);
|
||||
gpt2_zero_grad(&model);
|
||||
gpt2_backward(&model, x);
|
||||
gpt2_backward(&model, x, true);
|
||||
gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1.f, 1, &multi_gpu_config);
|
||||
cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings
|
||||
|
||||
|
|
|
|||
|
|
@ -130,7 +130,7 @@ for rid, row in kernel_profile_data:
|
|||
# the classifier part, counts only once
|
||||
pass_name = "cls"
|
||||
phase = "bwd"
|
||||
elif "adamw" in kernel or "global_norm" in kernel:
|
||||
elif "adamw" in kernel or "global_norm" in kernel or "copy_and_cast" in kernel:
|
||||
# encoder layer or adam
|
||||
pass_name = "opt"
|
||||
# before the first optimizer run, we create weight copies.
|
||||
|
|
|
|||
|
|
@ -210,7 +210,7 @@ int main(int argc, char *argv[]) {
|
|||
clock_gettime(CLOCK_MONOTONIC, &start);
|
||||
gpt2_forward(&model, x, y, B, T);
|
||||
gpt2_zero_grad(&model);
|
||||
gpt2_backward(&model, x);
|
||||
gpt2_backward(&model, x, true);
|
||||
clock_gettime(CLOCK_MONOTONIC, &end);
|
||||
double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
|
||||
|
||||
|
|
|
|||
463
train_gpt2.cu
463
train_gpt2.cu
|
|
@ -62,157 +62,16 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage.
|
|||
// defines: global_norm_squared
|
||||
#include "llmc/global_norm.cuh"
|
||||
// ----------- Multi-GPU support -----------
|
||||
#ifdef MULTI_GPU
|
||||
#include <mpi.h>
|
||||
#include <nccl.h>
|
||||
#endif
|
||||
#include "llmc/zero.cuh"
|
||||
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// global vars containing information about the GPU this process is running on
|
||||
cudaDeviceProp deviceProp; // fills in common_start()
|
||||
cudaStream_t main_stream;
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Multi-GPU related
|
||||
#ifdef MULTI_GPU
|
||||
|
||||
#if defined(ENABLE_FP32)
|
||||
const ncclDataType_t ncclFloatX = ncclFloat;
|
||||
#elif defined(ENABLE_FP16)
|
||||
const ncclDataType_t ncclFloatX = ncclHalf;
|
||||
#else // Default to bfloat16
|
||||
const ncclDataType_t ncclFloatX = ncclBfloat16;
|
||||
#endif
|
||||
|
||||
void nccl_check(ncclResult_t status, const char *file, int line) {
|
||||
if (status != ncclSuccess) {
|
||||
printf("[NCCL ERROR] at file %s:%d:\n%s\n", file, line, ncclGetErrorString(status));
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
}
|
||||
#define ncclCheck(err) (nccl_check(err, __FILE__, __LINE__))
|
||||
|
||||
void mpi_check(int status, const char *file, int line) {
|
||||
if (status != MPI_SUCCESS) {
|
||||
char mpi_error[4096];
|
||||
int mpi_error_len = 0;
|
||||
assert(MPI_Error_string(status, &mpi_error[0], &mpi_error_len) == MPI_SUCCESS);
|
||||
printf("[MPI ERROR] at file %s:%d:\n%.*s\n", file, line, mpi_error_len, mpi_error);
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
}
|
||||
#define mpiCheck(err) (mpi_check(err, __FILE__, __LINE__))
|
||||
|
||||
#endif // MULTI_GPU
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// MPI / multi-processing setup
|
||||
|
||||
// Parameters specific to training on multiple GPUs.
|
||||
typedef struct {
|
||||
int process_rank; // Rank of this process among all MPI processes. 0 if no multi-GPU.
|
||||
int num_processes; // Total number of processes. 1 if no multi-GPU.
|
||||
int local_device_idx; // This process GPU index on current machine. 0 if no multi-GPU.
|
||||
|
||||
// Zero Redundancy Optimizer stage - https://fairscale.readthedocs.io/en/stable/deep_dive/oss_sdp_fsdp.html
|
||||
// 0-Disabled
|
||||
// 1-Optimizer State Sharding (OSS)
|
||||
// 2-Optimizer + Gradient State Sharding (SDP)
|
||||
// 3-Optimizer + Gradient + Horizontal Model Sharding (FSDP)
|
||||
int zero_stage;
|
||||
size_t shard_num_parameters;
|
||||
size_t shard_offset;
|
||||
#ifdef MULTI_GPU
|
||||
ncclComm_t nccl_comm; // NCCL communication primitive, used for collective multi-GPU work.
|
||||
#endif
|
||||
} MultiGpuConfig;
|
||||
|
||||
// one global variable to hold the multi-GPU configuration for this process
|
||||
MultiGpuConfig multi_gpu_config;
|
||||
|
||||
#ifdef MULTI_GPU
|
||||
// Determine which GPU this process should use.
|
||||
// Processes on the same machines use different GPU indicies. Processes on other machines don't.
|
||||
// Copied from NCCL examples: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html#example-2-one-device-per-process-or-thread
|
||||
int multi_gpu_get_local_device_idx(int process_rank, int num_processes) {
|
||||
char hostname[1024];
|
||||
hostname[1023] = '\0';
|
||||
// All processes on the same machine will share the same hostname.
|
||||
gethostname(hostname, 1023);
|
||||
for (int i=0; i < 1024; i++) {
|
||||
if (hostname[i] == '.') {
|
||||
hostname[i] = '\0';
|
||||
break;
|
||||
}
|
||||
}
|
||||
uint64_t hostname_hash = 5381;
|
||||
for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5) + hostname_hash) ^ hostname[c]; }
|
||||
|
||||
// Distribute all hostname hashes to all processes.
|
||||
uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t));
|
||||
all_hostsname_hashes[process_rank] = hostname_hash;
|
||||
mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD));
|
||||
|
||||
// Identify which GPU we need to use.
|
||||
int local_device_idx = 0;
|
||||
for (int current_process = 0; current_process < num_processes; ++current_process) {
|
||||
if (current_process == process_rank) {
|
||||
// Found my gpu, local_device_idx now has my target GPU index.
|
||||
break;
|
||||
}
|
||||
if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) {
|
||||
// This process ID runs on the same machine, but it's not me, skip this GPU
|
||||
local_device_idx++;
|
||||
}
|
||||
}
|
||||
|
||||
free(all_hostsname_hashes);
|
||||
return local_device_idx;
|
||||
}
|
||||
#endif
|
||||
|
||||
MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) {
|
||||
#ifdef MULTI_GPU
|
||||
// Initialize MPI.
|
||||
MultiGpuConfig result;
|
||||
mpiCheck(MPI_Init(argc, argv));
|
||||
mpiCheck(MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank));
|
||||
mpiCheck(MPI_Comm_size(MPI_COMM_WORLD, &result.num_processes));
|
||||
result.local_device_idx = multi_gpu_get_local_device_idx(result.process_rank, result.num_processes);
|
||||
cudaCheck(cudaSetDevice(result.local_device_idx));
|
||||
ncclUniqueId nccl_id;
|
||||
if (result.process_rank == 0) {
|
||||
ncclCheck(ncclGetUniqueId(&nccl_id));
|
||||
}
|
||||
mpiCheck(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
|
||||
ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank));
|
||||
return result;
|
||||
#else
|
||||
printf("Multi-GPU support is disabled. Using a single GPU.\n");
|
||||
cudaCheck(cudaSetDevice(0));
|
||||
MultiGpuConfig result;
|
||||
result.process_rank = 0;
|
||||
result.num_processes = 1;
|
||||
result.local_device_idx = 0;
|
||||
return result;
|
||||
#endif
|
||||
}
|
||||
|
||||
void multi_gpu_config_free(const MultiGpuConfig* multi_gpu_config) {
|
||||
#ifdef MULTI_GPU
|
||||
ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm));
|
||||
mpiCheck(MPI_Finalize());
|
||||
#endif
|
||||
}
|
||||
|
||||
void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) {
|
||||
#ifdef MULTI_GPU
|
||||
if (multi_gpu_config->num_processes > 1) {
|
||||
mpiCheck(MPI_Barrier(MPI_COMM_WORLD));
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// convenience function that only prints if the rank of process is zero
|
||||
void printf0(const char *format, ...) {
|
||||
if (multi_gpu_config.process_rank == 0) {
|
||||
|
|
@ -224,11 +83,8 @@ void printf0(const char *format, ...) {
|
|||
}
|
||||
|
||||
void set_zero_configs(MultiGpuConfig* multi_gpu_config, int zero_stage, size_t total_parameters) {
|
||||
|
||||
multi_gpu_config->zero_stage = 0;
|
||||
multi_gpu_config->shard_num_parameters = total_parameters;
|
||||
multi_gpu_config->shard_offset = 0;
|
||||
|
||||
// Check the Zero Stage and define sharding parameters
|
||||
if (zero_stage == 0) {
|
||||
printf0("| Zero Optimization is disabled |\n");
|
||||
|
|
@ -242,7 +98,6 @@ void set_zero_configs(MultiGpuConfig* multi_gpu_config, int zero_stage, size_t t
|
|||
printf0("| Zero Stage1 is enabled |\n");
|
||||
multi_gpu_config->zero_stage = 1;
|
||||
multi_gpu_config->shard_num_parameters = total_parameters / multi_gpu_config->num_processes;
|
||||
multi_gpu_config->shard_offset = multi_gpu_config->process_rank * multi_gpu_config->shard_num_parameters;
|
||||
}
|
||||
}
|
||||
else{
|
||||
|
|
@ -337,7 +192,7 @@ void* malloc_and_point_parameters(ParameterTensors* params, size_t* param_elemen
|
|||
return params_memory;
|
||||
}
|
||||
|
||||
#define NUM_ACTIVATION_TENSORS 21
|
||||
constexpr int NUM_ACTIVATION_TENSORS = 23;
|
||||
typedef struct {
|
||||
floatX* encoded; // (B, T, C)
|
||||
floatX* ln1; // (L, B, T, C)
|
||||
|
|
@ -366,6 +221,10 @@ typedef struct {
|
|||
// general scratchpad buffer. Allocation is made large enough to hold (B, T, 3C),
|
||||
// (B, NH, T, T), and (B, T, V) shaped tensors.
|
||||
floatX* output;
|
||||
|
||||
// some additional scratch buffers
|
||||
floatX* scratch_bt4c; // (B, T, 4*C)
|
||||
floatX* scratch_btc; // (B, T, C)
|
||||
} ActivationTensors;
|
||||
|
||||
void fill_in_activation_sizes(size_t* act_sizes, size_t B, size_t T, GPT2Config config, int recompute) {
|
||||
|
|
@ -402,34 +261,9 @@ void fill_in_activation_sizes(size_t* act_sizes, size_t B, size_t T, GPT2Config
|
|||
act_sizes[18] = B * T; // losses
|
||||
act_sizes[19] = L * B * T * 3*C; // qkvr
|
||||
act_sizes[20] = B * T * max(3*C, max(NH*T, Vp)); // output / scratch
|
||||
}
|
||||
|
||||
// Backward pass is conceptually quite different from forward, because we can discard
|
||||
// the activations of a layer as soon as we're done with it. This lets us aggressively
|
||||
// reuse memory, so that we need far fewer tensors for backward state.
|
||||
#ifdef ENABLE_CUDNN
|
||||
#define NUM_BACKWARD_TENSORS 2
|
||||
#else
|
||||
#define NUM_BACKWARD_TENSORS 3
|
||||
#endif
|
||||
|
||||
typedef struct {
|
||||
floatX* bt4c; // (B, T, 4*C)
|
||||
floatX* residual3; // (B, T, C)
|
||||
#ifndef ENABLE_CUDNN
|
||||
floatX* preatt; // (B, NH, T, T)
|
||||
#endif
|
||||
} GradActTensors;
|
||||
|
||||
void fill_in_grad_act_sizes(size_t* act_sizes, size_t B, size_t T, GPT2Config config) {
|
||||
size_t C = config.channels;
|
||||
act_sizes[0] = B * T * 4 * C; // bt4c
|
||||
act_sizes[1] = B * T * C; // residual3
|
||||
|
||||
#ifndef ENABLE_CUDNN
|
||||
size_t NH = config.num_heads;
|
||||
act_sizes[2] = B * NH * T * T; // preatt
|
||||
#endif
|
||||
act_sizes[21] = B * T * 4 * C; // scratch_bt4c
|
||||
act_sizes[22] = B * T * C; // scratch_btc
|
||||
}
|
||||
|
||||
void* malloc_and_point(floatX** targets[], const size_t* act_sizes, size_t n) {
|
||||
|
|
@ -457,21 +291,12 @@ void* malloc_and_point_activations(ActivationTensors* acts, const size_t* act_si
|
|||
&acts->encoded, &acts->ln1, &acts->ln1_mean, &acts->ln1_rstd, &acts->atty,
|
||||
&acts->att, &acts->attproj, &acts->residual2, &acts->ln2, &acts->ln2_mean,
|
||||
&acts->ln2_rstd, &acts->fch, &acts->fch_gelu, &acts->fcproj, &acts->residual3, &acts->lnf,
|
||||
&acts->lnf_mean, &acts->lnf_rstd, &acts->losses, &acts->qkvr, &acts->output
|
||||
&acts->lnf_mean, &acts->lnf_rstd, &acts->losses, &acts->qkvr, &acts->output,
|
||||
&acts->scratch_bt4c, &acts->scratch_btc
|
||||
};
|
||||
return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS);
|
||||
}
|
||||
|
||||
void* malloc_and_point_backward(GradActTensors* acts, const size_t* act_sizes) {
|
||||
floatX** ptrs[] = {
|
||||
&acts->bt4c, &acts->residual3,
|
||||
#ifndef ENABLE_CUDNN
|
||||
&acts->preatt,
|
||||
#endif
|
||||
};
|
||||
return malloc_and_point(ptrs, act_sizes, NUM_BACKWARD_TENSORS);
|
||||
}
|
||||
|
||||
typedef struct {
|
||||
GPT2Config config;
|
||||
// the weights of the model, and their sizes
|
||||
|
|
@ -493,10 +318,6 @@ typedef struct {
|
|||
size_t act_sizes[NUM_ACTIVATION_TENSORS];
|
||||
void* acts_memory;
|
||||
size_t num_activations;
|
||||
// gradients of the activations
|
||||
GradActTensors grads_acts;
|
||||
size_t num_grad_acts;
|
||||
void* grads_acts_memory;
|
||||
// other run state configuration
|
||||
int batch_size; // the batch size (B) of current forward pass
|
||||
int seq_len; // the sequence length (T) of current forward pass
|
||||
|
|
@ -531,7 +352,6 @@ void gpt2_init_common(GPT2 *model) {
|
|||
model->mean_loss = -1.0f; // -1.0f designates no loss, set at end of forward()
|
||||
// memory lazily initialized in backward()
|
||||
model->grads_memory = NULL;
|
||||
model->grads_acts_memory = NULL;
|
||||
model->workload_indices = NULL; // on cpu, for encoder_backward
|
||||
model->bucket_info = NULL; // on cpu, for encoder_backward
|
||||
// memory lazily initialized in update()
|
||||
|
|
@ -901,7 +721,7 @@ void gpt2_zero_grad(GPT2 *model) {
|
|||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
void gpt2_backward(GPT2 *model, int* inputs) {
|
||||
void gpt2_backward(GPT2 *model, int* inputs, bool last_step) {
|
||||
NVTX_RANGE_FN();
|
||||
// double check we forwarded previously, with targets
|
||||
if (model->mean_loss == -1.0f) {
|
||||
|
|
@ -915,17 +735,6 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
// 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);
|
||||
// we're going to be clever for the activations backward pass. we don't need to exactly
|
||||
// mirror the forward pass activations and we will save memory.
|
||||
size_t bw_act_sizes[NUM_BACKWARD_TENSORS];
|
||||
fill_in_grad_act_sizes(bw_act_sizes, model->batch_size, model->seq_len, model->config);
|
||||
// count up and allocate the space
|
||||
model->num_grad_acts = 0;
|
||||
for (size_t i = 0; i < NUM_BACKWARD_TENSORS; i++) {
|
||||
model->num_grad_acts += bw_act_sizes[i];
|
||||
}
|
||||
printf0("allocating %d MiB for activation gradients\n", (int)round(model->num_grad_acts * sizeof(floatX) / (1024 * 1024)));
|
||||
model->grads_acts_memory = malloc_and_point_backward(&model->grads_acts, bw_act_sizes);
|
||||
// init gradients of parameters and activations to zero
|
||||
gpt2_zero_grad(model);
|
||||
// initialise cpu scratch buffers for encoder backward
|
||||
|
|
@ -947,10 +756,10 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
ParameterTensors params = model->params; // for brevity
|
||||
ParameterTensors grads = model->grads;
|
||||
ActivationTensors acts = model->acts;
|
||||
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)));
|
||||
floatX* dresidual = (floatX*)model->acts.scratch_btc; // the main buffer holding the gradient in the backward pass
|
||||
cudaCheck(cudaMemset(dresidual, 0, B * T * C * sizeof(floatX)));
|
||||
|
||||
// re-use the output buffer of the forward pass as a scratchpad during backward pass
|
||||
float* scratchF = (float*)acts.output;
|
||||
|
|
@ -961,11 +770,10 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
// 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, main_stream);
|
||||
matmul_backward(model->acts.scratch_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, main_stream);
|
||||
layernorm_backward(dresidual, grads.lnfw, grads.lnfb, scratchF, model->acts.scratch_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
|
||||
|
|
@ -1015,7 +823,7 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
// notice that there is no l *, because we just have a single copy, and keep
|
||||
// re-using this memory in every Transformer block as we calculate backward pass
|
||||
|
||||
floatX* dl_bt4c = (floatX*)grads_acts.bt4c;
|
||||
floatX* dl_bt4c = (floatX*)model->acts.scratch_bt4c;
|
||||
|
||||
// start the backward pass for this layer
|
||||
if(model->recompute >= 1) {
|
||||
|
|
@ -1042,8 +850,7 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
// 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, main_stream);
|
||||
attention_backward(dl_bt4c, buffer_b, 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, main_stream);
|
||||
|
|
@ -1052,10 +859,38 @@ void gpt2_backward(GPT2 *model, int* inputs) {
|
|||
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, main_stream);
|
||||
|
||||
// Accumulate gradients from this layer in a background stream.
|
||||
if(last_step) {
|
||||
floatX* const pointers[] = {
|
||||
dl_ln1w, dl_ln1b,
|
||||
dl_qkvw, dl_qkvb,
|
||||
dl_attprojw, dl_attprojb,
|
||||
dl_ln2w, dl_ln2b,
|
||||
dl_fcw, dl_fcb,
|
||||
dl_fcprojw, dl_fcprojb
|
||||
};
|
||||
const size_t nelem[] = {
|
||||
C, C,
|
||||
3 * C * C, 3 * C,
|
||||
C * C, C,
|
||||
C, C,
|
||||
4 * C * C, 4 * C,
|
||||
C * 4 * C, C
|
||||
};
|
||||
multi_gpu_async_reduce_gradient(pointers, nelem, &multi_gpu_config, 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), main_stream);
|
||||
|
||||
// Aggregate all gradients that are not part of the transformer blocks
|
||||
if(last_step) {
|
||||
floatX* const pointers[] = {grads.wte, grads.wpe, grads.lnfw, grads.lnfb};
|
||||
const size_t nelem[] = {Vp * C, T * C, C, C};
|
||||
multi_gpu_async_reduce_gradient(pointers, nelem, &multi_gpu_config, main_stream);
|
||||
}
|
||||
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
|
|
@ -1073,32 +908,34 @@ float multi_gpu_cpu_float_sum(float value) {
|
|||
|
||||
// Averages out the loss and gradients across all GPUs. No-op when multi-GPU is disabled.
|
||||
// todo - this version only works if all the parameters are the same size (floatX)
|
||||
void gpt2_multi_gpu_loss_and_grad_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) {
|
||||
void gpt2_multi_gpu_loss_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) {
|
||||
#ifdef MULTI_GPU
|
||||
NVTX_RANGE_FN();
|
||||
// If there's only one process, there is nothing to do
|
||||
if (multi_gpu_config->num_processes == 1) { return; }
|
||||
// Average all losses.
|
||||
model->accumulated_mean_loss = multi_gpu_cpu_float_sum(model->mean_loss) / multi_gpu_config->num_processes;
|
||||
// Now average the gradients
|
||||
if(multi_gpu_config->zero_stage == 0) {
|
||||
// no ZERO == standard DDP: Average all gradients.
|
||||
ncclCheck(ncclAllReduce(model->grads_memory, model->grads_memory,
|
||||
model->num_parameters,
|
||||
ncclFloatX, ncclAvg,
|
||||
multi_gpu_config->nccl_comm, 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, main_stream));
|
||||
}
|
||||
#endif
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
// Gets the offset of a specific tensor for a specific layer in the GPT2 model
|
||||
// layer_id is ignored for weights that are not part of a transformer block
|
||||
ShardInfo gpt2_get_tensor_at_layer(const GPT2 *model, int layer_id, int param_tensor_id) {
|
||||
// first offset our way to the parameter tensor start
|
||||
ptrdiff_t offset = 0;
|
||||
for (int i = 0; i < param_tensor_id; i++) {
|
||||
offset += (ptrdiff_t)model->param_elements[i];
|
||||
}
|
||||
size_t size = model->param_elements[param_tensor_id] ;
|
||||
// if we are in the transformer block, we need to additionally offset by the layer id
|
||||
if(2 <= param_tensor_id && param_tensor_id <= 13) {
|
||||
size /= model->config.num_layers;
|
||||
offset += (ptrdiff_t)(layer_id * size);
|
||||
}
|
||||
return {offset, size};
|
||||
}
|
||||
|
||||
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) {
|
||||
// update the model parameters using the AdamW optimizer
|
||||
// keep in mind that optimizer sharding (ZeRO-1) assigns different parameters to different GPUs
|
||||
|
|
@ -1108,8 +945,6 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
// TODO: revisit and probably refactor this entire function
|
||||
NVTX_RANGE_FN();
|
||||
size_t shard_num_parameters = multi_gpu_config->shard_num_parameters; // num parameters we are responsible for
|
||||
size_t shard_offset = multi_gpu_config->shard_offset; // offset into the full parameter tensor
|
||||
floatX* params_memory = (floatX*)model->params_memory;
|
||||
floatX* grads_memory = (floatX*)model->grads_memory;
|
||||
|
||||
// lazily allocate m,v memory and master weights (usually on the first iteration)
|
||||
|
|
@ -1122,33 +957,46 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
cudaCheck(cudaMemset(model->m_memory, 0, shard_num_parameters * sizeof(float)));
|
||||
cudaCheck(cudaMemset(model->v_memory, 0, shard_num_parameters * sizeof(float)));
|
||||
}
|
||||
|
||||
bool init_master_weights = false;
|
||||
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, 0, main_stream>>>(model->master_weights, params_memory + shard_offset, shard_num_parameters);
|
||||
cudaCheck(cudaGetLastError());
|
||||
init_master_weights = true;
|
||||
}
|
||||
|
||||
// gradient clipping
|
||||
// repurposing this buffer (which isn't needed now) to write grad norm into it
|
||||
float* grad_norm_squared = (float*)model->acts.output;
|
||||
if (multi_gpu_config->zero_stage == 1) {
|
||||
// ^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, 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, main_stream);
|
||||
}
|
||||
// transfer the gradient norm to CPU
|
||||
float grad_norm_squared_cpu = 0.0f;
|
||||
cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost));
|
||||
|
||||
if (multi_gpu_config->zero_stage == 1) {
|
||||
// because of the ncclReduceScatter() in backward,
|
||||
// grads_memory only contains the averaged gradients at the local shards,
|
||||
// so we only calculate the grad norm at the grads_memory belonging to the local shards
|
||||
cudaCheck(cudaMemsetAsync(grad_norm_squared, 0, sizeof(float), main_stream));
|
||||
for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
|
||||
if((i < 2 || i > 13)) {
|
||||
ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i);
|
||||
ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1);
|
||||
ptrdiff_t offset = tensor.offset + shard.offset;
|
||||
global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, false, main_stream);
|
||||
} else {
|
||||
ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i);
|
||||
ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1);
|
||||
ptrdiff_t offset = tensor.offset + shard.offset;
|
||||
global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, tensor.size, model->config.num_layers,
|
||||
false, 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);
|
||||
} 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, true, main_stream);
|
||||
cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost));
|
||||
}
|
||||
|
||||
if(!isfinite(grad_norm_squared_cpu)) {
|
||||
|
|
@ -1162,83 +1010,66 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
|
|||
|
||||
// AdamW update
|
||||
unsigned int seed = random_u32(&model->rng_state);
|
||||
// individually call the adamw_kernel3 on all parameter tensors separately
|
||||
size_t offset = 0;
|
||||
for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
|
||||
size_t num_parameters = model->param_elements[i];
|
||||
// the scope of this GPU's work is the range: [shard_offset, shard_offset + shard_num_parameters)
|
||||
// this parameter's values are in the range: [offset, offset + num_parameters)
|
||||
// so we are responsible for some of its parameters if:
|
||||
// 1) this parameter ends after we begin (i.e. offset + num_parameters > shard_offset)
|
||||
// 2) this parameter begins before we end (i.e. offset < shard_offset + shard_num_parameters)
|
||||
if(offset + num_parameters > shard_offset && offset < shard_offset + shard_num_parameters) {
|
||||
|
||||
// ok this tensor has at least one element inside the range of responsibility of this GPU
|
||||
// let's figure out the exact span we wish to call the AdamW kernel on
|
||||
floatX* params_ptr = NULL;
|
||||
floatX* grad_ptr = NULL;
|
||||
float* m_ptr = NULL;
|
||||
float* v_ptr = NULL;
|
||||
float* master_ptr = NULL;
|
||||
size_t local_params = 0;
|
||||
// does the tensor begin before our responsibility?
|
||||
if(offset <= shard_offset) {
|
||||
// if so, our start point is exactly that of our responsibility, i.e. shard_offset
|
||||
params_ptr = params_memory + shard_offset;
|
||||
grad_ptr = grads_memory + shard_offset;
|
||||
// note that (master_weights, m, v) are already only the "local slice" for this GPU,
|
||||
// and are of size shard_num_parameters, instead of the total number of parameters
|
||||
// so they do not get offset, i.e. we just start at their index 0
|
||||
if (model->master_weights != NULL) { master_ptr = model->master_weights; }
|
||||
m_ptr = model->m_memory;
|
||||
v_ptr = model->v_memory;
|
||||
// the number of parameters we have to update is the minimum of two ranges
|
||||
local_params = min(shard_num_parameters, (offset + num_parameters) - shard_offset);
|
||||
} else {
|
||||
// our start point is the location of this tensor, i.e. offset
|
||||
params_ptr = params_memory + offset;
|
||||
grad_ptr = grads_memory + offset;
|
||||
// this arithmetic gave me a headache but my little doodle example says it's right
|
||||
size_t delta = offset - shard_offset;
|
||||
if (model->master_weights != NULL) { master_ptr = model->master_weights + delta; }
|
||||
m_ptr = model->m_memory + delta;
|
||||
v_ptr = model->v_memory + delta;
|
||||
local_params = min(num_parameters, shard_num_parameters - delta);
|
||||
}
|
||||
// we only want to weight decay the 2D tensors and leave all 1D tensors alone
|
||||
// in particular this also decays the embedding weights, but this is ok:
|
||||
// - the token embeddings are weight shared and participate in the final projection to logits
|
||||
// - 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
|
||||
adamw_update(params_ptr, master_ptr, grad_ptr,
|
||||
m_ptr, v_ptr, local_params, learning_rate,
|
||||
beta1, beta2, t, eps, wd, grad_scale, seed, main_stream);
|
||||
// handle adamw for all the transformer blocks
|
||||
for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
|
||||
int num_layers = model->config.num_layers;
|
||||
if((i < 2 || i > 13)) {
|
||||
num_layers = 1;
|
||||
}
|
||||
|
||||
ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i);
|
||||
ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1);
|
||||
ptrdiff_t local_offset_full = tensor.offset + shard.offset;
|
||||
ptrdiff_t local_offset_partial = tensor.offset / multi_gpu_config->num_processes;
|
||||
|
||||
// we only want to weight decay the 2D tensors and leave all 1D tensors alone
|
||||
// in particular this also decays the embedding weights, but this is ok:
|
||||
// - the token embeddings are weight shared and participate in the final projection to logits
|
||||
// - 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;
|
||||
floatX* param_ptr = (floatX*)model->params_memory + local_offset_full;
|
||||
floatX* grad_ptr = (floatX*)model->grads_memory + local_offset_full;
|
||||
|
||||
ptrdiff_t opt_state_offset = multi_gpu_config->zero_stage < 1 ? local_offset_full : local_offset_partial;
|
||||
float* m_ptr = model->m_memory + opt_state_offset;
|
||||
float* v_ptr = model->v_memory + opt_state_offset;
|
||||
float* master_ptr = NULL;
|
||||
if (model->master_weights != NULL) { master_ptr = model->master_weights + opt_state_offset; }
|
||||
if(init_master_weights) {
|
||||
size_t grid_size = CEIL_DIV(shard_num_parameters, 512);
|
||||
copy_and_cast_kernel<<<dim3(grid_size, num_layers), 512, 0, main_stream>>>(master_ptr, param_ptr, shard.size,
|
||||
shard.size, tensor.size);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
// ok finally call the kernel
|
||||
adamw_update(param_ptr, master_ptr, grad_ptr,
|
||||
m_ptr, v_ptr,
|
||||
shard.size, tensor.size, tensor.size, shard.size, num_layers,
|
||||
learning_rate,
|
||||
beta1, beta2, t, eps, wd, grad_scale, seed, main_stream);
|
||||
cudaCheck(cudaGetLastError());
|
||||
|
||||
if (multi_gpu_config->zero_stage == 1) {
|
||||
#if MULTI_GPU
|
||||
ncclCheck(ncclGroupStart());
|
||||
for(int l = 0; l < num_layers; ++l) {
|
||||
// gather updated shards of model->params_memory from each process
|
||||
ncclCheck(ncclAllGather(param_ptr + l * tensor.size,
|
||||
(floatX*) model->params_memory + tensor.offset + l * tensor.size,
|
||||
shard.size, ncclFloatX,
|
||||
multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream));
|
||||
}
|
||||
ncclCheck(ncclGroupEnd());
|
||||
#endif
|
||||
}
|
||||
// advance the offset pointer to the next parameter tensor
|
||||
offset += num_parameters;
|
||||
}
|
||||
cudaCheck(cudaGetLastError());
|
||||
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
return grad_norm_cpu;
|
||||
}
|
||||
|
||||
void gpt2_multi_gpu_param_gather(GPT2 *model, MultiGpuConfig* multi_gpu_config)
|
||||
{
|
||||
#ifdef MULTI_GPU
|
||||
if (multi_gpu_config->num_processes == 1) { return; } // 1 process => noop
|
||||
if (multi_gpu_config->zero_stage == 1) {
|
||||
// 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, main_stream));
|
||||
}
|
||||
cudaCheck(cudaGetLastError());
|
||||
#endif
|
||||
cudaCheck(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
float gpt2_estimate_mfu(GPT2 *model, int num_tokens, float dt) {
|
||||
/*
|
||||
Estimate model flops utilization (MFU)
|
||||
|
|
@ -1275,7 +1106,6 @@ void gpt2_free(GPT2 *model) {
|
|||
cudaCheck(cudaFree(model->v_memory));
|
||||
cudaCheck(cudaFree(model->master_weights));
|
||||
cudaCheck(cudaFree(model->acts_memory));
|
||||
cudaCheck(cudaFree(model->grads_acts_memory));
|
||||
cudaCheck(cudaFree(model->inputs));
|
||||
cudaCheck(cudaFree(model->targets));
|
||||
cudaCheck(cudaFreeHost(model->cpu_losses));
|
||||
|
|
@ -1809,13 +1639,13 @@ int main(int argc, char *argv[]) {
|
|||
gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T, grad_accum_steps);
|
||||
lossf += model.mean_loss; // the mean_loss was normalized by grad_accum_steps inside gpt2_forward
|
||||
// backward pass. all model params accumulate gradients with += inside this inner loop
|
||||
gpt2_backward(&model, train_loader.inputs);
|
||||
gpt2_backward(&model, train_loader.inputs, micro_step == grad_accum_steps - 1);
|
||||
}
|
||||
// override the mean loss, accounting for the gradient accumulation loop
|
||||
// this is esp important to do here in multigpu update below, where model.mean_loss gets allreduced
|
||||
model.mean_loss = lossf;
|
||||
// average the loss and the gradients between all processes
|
||||
gpt2_multi_gpu_loss_and_grad_reduce(&model, &multi_gpu_config);
|
||||
gpt2_multi_gpu_loss_reduce(&model, &multi_gpu_config);
|
||||
// learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac
|
||||
float step_learning_rate = learning_rate;
|
||||
if (step < warmup_iterations) {
|
||||
|
|
@ -1830,7 +1660,6 @@ int main(int argc, char *argv[]) {
|
|||
}
|
||||
// update the model parameters
|
||||
float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, 1.0f, step+1, &multi_gpu_config);
|
||||
gpt2_multi_gpu_param_gather(&model, &multi_gpu_config);
|
||||
// zero out the gradients for the next iteration
|
||||
gpt2_zero_grad(&model);
|
||||
cudaCheck(cudaEventRecord(end));
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue