From e0abfae2bc731d4c093830266d4e701f71c19ced Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Sat, 1 Jun 2024 19:06:42 +0530 Subject: [PATCH 01/22] Benchmark modal script revamped - profiling and external libraries --- dev/cuda/benchmark_on_modal.py | 72 +++++++++++++++++++++++++--------- 1 file changed, 54 insertions(+), 18 deletions(-) diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py index 7a055ec..49173ff 100644 --- a/dev/cuda/benchmark_on_modal.py +++ b/dev/cuda/benchmark_on_modal.py @@ -2,18 +2,34 @@ 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" -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. -""" +OR if you want to use cuDNN and openMPI etc. +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 \ + ./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" + +-> 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 lot of errors 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 @@ -27,19 +43,15 @@ GPU_NAME_TO_MODAL_CLASS_MAP = { 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}") + # Custom image built with cuda tools and nsight system along with axolotl + Image.from_registry("totallyvyom/cuda-env:latest") .run_commands( - "git clone https://github.com/OpenAccess-AI-Collective/axolotl /root/axolotl", - "cd /root/axolotl && git checkout v0.4.0", + "cd /root/axolotl && git checkout v0.4.0" ) .pip_install("huggingface_hub==0.20.3", "hf-transfer==0.1.5") .env( @@ -49,34 +61,58 @@ 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 update", + "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" + ) ) 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, 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" 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.qdrep /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 \ No newline at end of file From fbc1e4983cbd9b75b27fa962a923909f1b0bd1c0 Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Mon, 3 Jun 2024 17:00:10 +0530 Subject: [PATCH 02/22] Benchmark modal script revamped - profiling and external libraries --- dev/cuda/benchmark_on_modal.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py index 49173ff..7ac38cd 100644 --- a/dev/cuda/benchmark_on_modal.py +++ b/dev/cuda/benchmark_on_modal.py @@ -1,28 +1,23 @@ """ Script for running benchmarks on the Modal platform. This is useful for folks who do not have access to expensive GPUs locally. - 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 and openMPI etc. 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 \ ./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" - -> 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 lot of errors 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 """ @@ -33,13 +28,11 @@ 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") From 9b540fa8595445e229ecb6429fbaf84deca5eba2 Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Mon, 3 Jun 2024 18:59:57 +0530 Subject: [PATCH 03/22] Makefile updated to match script --- Makefile | 1 + 1 file changed, 1 insertion(+) diff --git a/Makefile b/Makefile index c8b555a..502974c 100644 --- a/Makefile +++ b/Makefile @@ -14,6 +14,7 @@ CUDA_OUTPUT_FILE = -o $@ # NVCC flags # -t=0 is short for --threads, 0 = number of CPUs on the machine NVCC_FLAGS = -O3 -t=0 --use_fast_math +NVCC_FLAGS += --std=c++17 NVCC_LDFLAGS = -lcublas -lcublasLt NVCC_INCLUDES = NVCC_LDLIBS = From 4564a24c692a168f60fe5e4d4feee5581fc11f8c Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Tue, 4 Jun 2024 23:39:36 +0530 Subject: [PATCH 04/22] Modal script updated for cudnn and gcc version match --- dev/cuda/benchmark_on_modal.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py index 7ac38cd..034af49 100644 --- a/dev/cuda/benchmark_on_modal.py +++ b/dev/cuda/benchmark_on_modal.py @@ -40,11 +40,14 @@ GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, size=str(GPU_ME APP_NAME = "llm.c benchmark run" +AXOLOTL_REGISTRY_SHA = ( + "d5b941ba2293534c01c23202c8fc459fd2a169871fa5e6c45cb00f363d474b6a" +) axolotl_image = ( - # Custom image built with cuda tools and nsight system along with axolotl - Image.from_registry("totallyvyom/cuda-env:latest") + Image.from_registry("winglian/axolotl:main-20240604-py3.10-cu118-2.1.2") .run_commands( - "cd /root/axolotl && git checkout v0.4.0" + "git clone https://github.com/OpenAccess-AI-Collective/axolotl /root/axolotl", + "cd /root/axolotl && git checkout v0.4.0", ) .pip_install("huggingface_hub==0.20.3", "hf-transfer==0.1.5") .env( @@ -55,6 +58,11 @@ axolotl_image = ( ) ) .run_commands( + "wget -q https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-nsight-systems-12-5_12.5.0-1_amd64.deb", + "apt-get update", + # "dpkg -i cuda-nsight-12-5_12.5.39-1_amd64.deb", + "apt-get install -y cuda-nsight-systems-12-5", + "apt-get install -f -y", "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", @@ -63,10 +71,13 @@ axolotl_image = ( .run_commands( "apt-get update", "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( + "apt-get install -y cuda-toolkit-12-5", + ) + ) stub = modal.App(APP_NAME) @@ -100,7 +111,8 @@ def run_benchmark(compile_command: str, run_command: str): timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") execute_command("mkdir report1_" + timestamp) - execute_command("mv /root/report1.qdrep /root/report1_" + timestamp + "/") + execute_command("mv /root/report1.nsys-rep /root/report1_" + timestamp + "/") + execute_command("mv /root/report1.sqlite /root/report1_" + timestamp + "/") execute_command("mv /root/report1_" + timestamp + "/" + " /cuda-env/") return None From 1e8efe94c063fd8ff3dd115e5cfe36ad20c46596 Mon Sep 17 00:00:00 2001 From: Vyom Sharma <93402393+vyom1611@users.noreply.github.com> Date: Wed, 5 Jun 2024 12:02:09 +0530 Subject: [PATCH 05/22] Update Makefile - revert back --- Makefile | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 948c2ab..eddbca8 100644 --- a/Makefile +++ b/Makefile @@ -14,7 +14,6 @@ CUDA_OUTPUT_FILE = -o $@ # NVCC flags # -t=0 is short for --threads, 0 = number of CPUs on the machine NVCC_FLAGS = -O3 -t=0 --use_fast_math -NVCC_FLAGS += --std=c++17 NVCC_LDFLAGS = -lcublas -lcublasLt NVCC_INCLUDES = NVCC_LDLIBS = @@ -268,4 +267,4 @@ profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN) clean: $(REMOVE_FILES) $(TARGETS) - $(REMOVE_BUILD_OBJECT_FILES) \ No newline at end of file + $(REMOVE_BUILD_OBJECT_FILES) From 58c59064ec431a79600f1a0e4f7da21cb035ea54 Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Fri, 7 Jun 2024 14:21:13 +0530 Subject: [PATCH 06/22] Most updated working modal script --- dev/cuda/benchmark_on_modal.py | 60 ++++++++++++++++------------------ 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py index 034af49..b0db3fb 100644 --- a/dev/cuda/benchmark_on_modal.py +++ b/dev/cuda/benchmark_on_modal.py @@ -5,20 +5,26 @@ 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 and openMPI etc. +OR if you want to use cuDNN etc. + + 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" + --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 \ + --run-command "--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 lot of errors on the command line but it + +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 @@ -40,15 +46,8 @@ GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, size=str(GPU_ME APP_NAME = "llm.c benchmark run" -AXOLOTL_REGISTRY_SHA = ( - "d5b941ba2293534c01c23202c8fc459fd2a169871fa5e6c45cb00f363d474b6a" -) -axolotl_image = ( - Image.from_registry("winglian/axolotl:main-20240604-py3.10-cu118-2.1.2") - .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( @@ -58,26 +57,25 @@ axolotl_image = ( ) ) .run_commands( - "wget -q https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-nsight-systems-12-5_12.5.0-1_amd64.deb", - "apt-get update", - # "dpkg -i cuda-nsight-12-5_12.5.39-1_amd64.deb", - "apt-get install -y cuda-nsight-systems-12-5", - "apt-get install -f -y", - "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", - ) + "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 update", "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( - "apt-get install -y cuda-toolkit-12-5", + "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) @@ -89,14 +87,14 @@ def execute_command(command: str): @stub.function( gpu=GPU_CONFIG, - image=axolotl_image, + image=image, allow_concurrent_inputs=4, container_idle_timeout=900, 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" you would get your report + # 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} @@ -112,7 +110,7 @@ def run_benchmark(compile_command: str, run_command: str): execute_command("mkdir report1_" + timestamp) execute_command("mv /root/report1.nsys-rep /root/report1_" + timestamp + "/") - execute_command("mv /root/report1.sqlite /root/report1_" + timestamp + "/") + execute_command("mv /root/report1.qdstrm /root/report1_" + timestamp + "/") execute_command("mv /root/report1_" + timestamp + "/" + " /cuda-env/") return None From ecb13ea74df0c3ac9eb4905752422765b7abb485 Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Fri, 7 Jun 2024 16:14:59 +0530 Subject: [PATCH 07/22] missing parameter in comments added --- dev/cuda/benchmark_on_modal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py index b0db3fb..907a831 100644 --- a/dev/cuda/benchmark_on_modal.py +++ b/dev/cuda/benchmark_on_modal.py @@ -17,7 +17,7 @@ GPU_MEM=80 modal run dev/cuda/benchmark_on_modal.py \ 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 "--cuda-graph-trace=graph --python-backtrace=cuda --cuda-memory-usage=true \ + --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" From 50a9411114857a19d5361afc2c375dff9ac6e17f Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 00:45:16 +0300 Subject: [PATCH 08/22] overlap backward with nccl for non-ZeRO runs --- train_gpt2.cu | 80 +++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 72 insertions(+), 8 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 286afbb..1e59113 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -123,7 +123,9 @@ typedef struct { 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. + 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; @@ -186,6 +188,10 @@ MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { } 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)); + cudaCheck(cudaEventCreate(&result.compute_nccl_sync)); + 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"); @@ -201,6 +207,8 @@ MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { void multi_gpu_config_free(const 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 } @@ -901,7 +909,37 @@ void gpt2_zero_grad(GPT2 *model) { cudaCheck(cudaDeviceSynchronize()); } -void gpt2_backward(GPT2 *model, int* inputs) { +// Block NCCL stream until computations on compute_stream are done, then aggregate multiple pointers in an NCCL group. +template +void multi_gpu_async_all_reduce_pointers_group( + floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N], + MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) { +#ifdef MULTI_GPU + NVTX_RANGE_FN(); + if (multi_gpu_config->num_processes == 1) { + return; // no multi-GPU, just exit. + } + // 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) { + ncclCheck(ncclAllReduce( + pointers[i], pointers[i], + pointers_sizes[i], + ncclFloatX, ncclAvg, + multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream + )); + } + ncclCheck(ncclGroupEnd()); +#endif +} + +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) { @@ -1052,10 +1090,39 @@ 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* 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 + }; + 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_all_reduce_pointers_group(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* pointers[] = {grads.wte, grads.wpe, grads.lnfw, grads.lnfb}; + size_t nelem[] = {Vp * C, T * C, C, C}; + multi_gpu_async_all_reduce_pointers_group(pointers, nelem, &multi_gpu_config, main_stream); + } + cudaCheck(cudaDeviceSynchronize()); } @@ -1082,12 +1149,9 @@ void gpt2_multi_gpu_loss_and_grad_reduce(GPT2* model, MultiGpuConfig* multi_gpu_ 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)); + // gradients get averaged during backward } else if (multi_gpu_config->zero_stage == 1) { + assert(false && "This code is WIP; ZeRO-1 is currently broken"); // 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, @@ -1809,7 +1873,7 @@ 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 From 6d2c72cb4b673b841aff40598dd874bd2c8a7e72 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 01:17:34 +0300 Subject: [PATCH 09/22] rearranged adam code --- test_gpt2.cu | 2 +- train_gpt2.cu | 97 +++++++++++++++++++++++++++------------------------ 2 files changed, 52 insertions(+), 47 deletions(-) diff --git a/test_gpt2.cu b/test_gpt2.cu index d1d5c64..c24f4c4 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -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; diff --git a/train_gpt2.cu b/train_gpt2.cu index 1e59113..476cbda 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -221,6 +221,25 @@ void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) { #endif } +struct ShardInfo { + ptrdiff_t offset; + size_t size; +}; +ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* multi_gpu_config, int shard_at_stage) { + if(multi_gpu_config->zero_stage >= shard_at_stage) { + if (elements % multi_gpu_config->num_processes != 0) { + fprintf(stderr, "Number of elements %zu must be a multiple of the number of processes %d\n", elements, + multi_gpu_config->num_processes); + exit(EXIT_FAILURE); + } + return {(ptrdiff_t) (multi_gpu_config->process_rank * (elements / multi_gpu_config->num_processes)), + elements / multi_gpu_config->num_processes + }; + } else { + return {0, elements}; + } +} + // convenience function that only prints if the rank of process is zero void printf0(const char *format, ...) { if (multi_gpu_config.process_rank == 0) { @@ -1226,63 +1245,49 @@ 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); + // handle adamw for all the transformer blocks + for(int l = 0; l < model->config.num_layers; ++l) { + ptrdiff_t tensor_offset = 0; + for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + // skip over preceding layers + if(i > 0) { + tensor_offset += model->param_elements[i - 1]; } + ptrdiff_t local_offset; + size_t num_parameters; + if(2 <= i && i <= 12) { + size_t tensor_nelem = model->param_elements[i] / model->config.num_layers; + ShardInfo shard = multi_gpu_get_shard_offset(tensor_nelem, multi_gpu_config, 1); + local_offset = tensor_offset + l * tensor_nelem + shard.offset; + num_parameters = shard.size; + } else { + if(l != 0) continue; + ShardInfo shard = multi_gpu_get_shard_offset(model->param_elements[i], multi_gpu_config, 1); + local_offset = tensor_offset + shard.offset; + num_parameters = shard.size; + } + // 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; + floatX* grad_ptr = (floatX*)model->grads_memory + local_offset; + float* m_ptr = model->m_memory + local_offset; + float* v_ptr = model->v_memory + local_offset; + float* master_ptr = NULL; + if (model->master_weights != NULL) { master_ptr = model->master_weights + local_offset; } + // ok finally call the kernel - adamw_update(params_ptr, master_ptr, grad_ptr, - m_ptr, v_ptr, local_params, learning_rate, + adamw_update(param_ptr, master_ptr, grad_ptr, + m_ptr, v_ptr, + num_parameters, learning_rate, beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); + cudaCheck(cudaGetLastError()); } - // advance the offset pointer to the next parameter tensor - offset += num_parameters; } - cudaCheck(cudaGetLastError()); cudaCheck(cudaDeviceSynchronize()); return grad_norm_cpu; From fdcb30256df807f6866670ec7fb73623fd709e2f Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 02:16:14 +0300 Subject: [PATCH 10/22] first working ZeRO-1 --- train_gpt2.cu | 69 +++++++++++++++++++++++++++------------------------ 1 file changed, 37 insertions(+), 32 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 476cbda..bda894a 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -226,15 +226,13 @@ struct ShardInfo { size_t size; }; ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* multi_gpu_config, int shard_at_stage) { + int nproc = multi_gpu_config->num_processes; if(multi_gpu_config->zero_stage >= shard_at_stage) { - if (elements % multi_gpu_config->num_processes != 0) { - fprintf(stderr, "Number of elements %zu must be a multiple of the number of processes %d\n", elements, - multi_gpu_config->num_processes); + 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 / multi_gpu_config->num_processes)), - elements / multi_gpu_config->num_processes - }; + return {(ptrdiff_t) (multi_gpu_config->process_rank * (elements / nproc)), elements / nproc}; } else { return {0, elements}; } @@ -1159,25 +1157,13 @@ 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) { - // gradients get averaged during backward - } else if (multi_gpu_config->zero_stage == 1) { - assert(false && "This code is WIP; ZeRO-1 is currently broken"); - // 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()); } @@ -1248,23 +1234,32 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // handle adamw for all the transformer blocks for(int l = 0; l < model->config.num_layers; ++l) { - ptrdiff_t tensor_offset = 0; + ptrdiff_t full_tensor_offset = 0; + ptrdiff_t partial_tensor_offset = 0; for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { // skip over preceding layers if(i > 0) { - tensor_offset += model->param_elements[i - 1]; + full_tensor_offset += model->param_elements[i - 1]; + partial_tensor_offset += model->param_elements[i - 1] / multi_gpu_config->num_processes; } - ptrdiff_t local_offset; + + ptrdiff_t global_offset; + ptrdiff_t local_offset_full; + ptrdiff_t local_offset_partial; size_t num_parameters; - if(2 <= i && i <= 12) { + if(2 <= i && i <= 13) { size_t tensor_nelem = model->param_elements[i] / model->config.num_layers; ShardInfo shard = multi_gpu_get_shard_offset(tensor_nelem, multi_gpu_config, 1); - local_offset = tensor_offset + l * tensor_nelem + shard.offset; + global_offset = full_tensor_offset + l * tensor_nelem; + local_offset_full = full_tensor_offset + l * tensor_nelem + shard.offset; + local_offset_partial = partial_tensor_offset + l * shard.size; num_parameters = shard.size; } else { if(l != 0) continue; ShardInfo shard = multi_gpu_get_shard_offset(model->param_elements[i], multi_gpu_config, 1); - local_offset = tensor_offset + shard.offset; + global_offset = full_tensor_offset; + local_offset_full = full_tensor_offset + shard.offset; + local_offset_partial = partial_tensor_offset; num_parameters = shard.size; } @@ -1273,12 +1268,14 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // - 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; - floatX* grad_ptr = (floatX*)model->grads_memory + local_offset; - float* m_ptr = model->m_memory + local_offset; - float* v_ptr = model->v_memory + local_offset; + 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 + local_offset; } + if (model->master_weights != NULL) { master_ptr = model->master_weights + opt_state_offset; } // ok finally call the kernel adamw_update(param_ptr, master_ptr, grad_ptr, @@ -1286,6 +1283,14 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl num_parameters, learning_rate, beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); cudaCheck(cudaGetLastError()); + + if (multi_gpu_config->zero_stage == 1) { + // gather updated shards of model->params_memory from each process + ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + global_offset, + num_parameters, ncclFloatX, + multi_gpu_config->nccl_comm, main_stream)); + } + cudaCheck(cudaGetLastError()); } } @@ -1884,7 +1889,7 @@ int main(int argc, char *argv[]) { // 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) { @@ -1899,7 +1904,7 @@ 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); + //gpt2_multi_gpu_param_gather(&model, &multi_gpu_config); // zero out the gradients for the next iteration gpt2_zero_grad(&model); cudaCheck(cudaEventRecord(end)); From ecdb0ca86e965d9181e518faa8d494c66865615b Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 12:55:17 +0300 Subject: [PATCH 11/22] refactored code --- llmc/global_norm.cuh | 7 +- llmc/zero.cuh | 215 +++++++++++++++++++++++++++++ train_gpt2.cu | 312 ++++++++----------------------------------- 3 files changed, 275 insertions(+), 259 deletions(-) create mode 100644 llmc/zero.cuh diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index f360cb3..713da7e 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -2,6 +2,7 @@ Global norm, used in gradient clipping */ #include +#include // llmc internal imports #include "cuda_common.h" #include "cuda_utils.cuh" @@ -34,7 +35,7 @@ __global__ void global_norm_squared_kernel(float* out, const T* data, size_t cou // kernel launcher template -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, 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,7 +45,9 @@ 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)); + if(reset) { + cudaCheck(cudaMemsetAsync(out, 0, sizeof(float), stream)); + } global_norm_squared_kernel<<>>(out, values, count); cudaCheck(cudaGetLastError()); } diff --git a/llmc/zero.cuh b/llmc/zero.cuh new file mode 100644 index 0000000..76ead6b --- /dev/null +++ b/llmc/zero.cuh @@ -0,0 +1,215 @@ +/* +Utilities for ZeRO sharding +*/ + +#ifndef LLMC_ZERO_CUH +#define LLMC_ZERO_CUH + +#include +#include +#include +#include +#include + +#ifdef MULTI_GPU +#include +#include +#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)); + cudaCheck(cudaEventCreate(&result.compute_nccl_sync)); + 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. +template +void multi_gpu_async_all_reduce_pointers_group( + floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N], + MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) { +#ifdef MULTI_GPU + NVTX_RANGE_FN(); + if (multi_gpu_config->num_processes == 1) { + return; // no multi-GPU, just exit. + } + // 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) { + ncclCheck(ncclAllReduce( + pointers[i], pointers[i], + pointers_sizes[i], + ncclFloatX, ncclAvg, + multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream + )); + } + ncclCheck(ncclGroupEnd()); +#endif +} + +#endif + diff --git a/train_gpt2.cu b/train_gpt2.cu index bda894a..6e09252 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -62,198 +62,27 @@ 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 -#include -#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. - 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; - // 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)); - cudaCheck(cudaStreamCreate(&result.nccl_stream)); - cudaCheck(cudaEventCreate(&result.compute_nccl_sync)); - 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(const 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 -} - -struct ShardInfo { - ptrdiff_t offset; - size_t size; -}; -ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* multi_gpu_config, int shard_at_stage) { - 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}; - } -} - // convenience function that only prints if the rank of process is zero -void printf0(const char *format, ...) { +template +void printf0(const char *format, Args... args) { if (multi_gpu_config.process_rank == 0) { - va_list args; - va_start(args, format); - vprintf(format, args); - va_end(args); + printf(format, args...); } } 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"); @@ -267,7 +96,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{ @@ -926,36 +754,6 @@ void gpt2_zero_grad(GPT2 *model) { cudaCheck(cudaDeviceSynchronize()); } -// Block NCCL stream until computations on compute_stream are done, then aggregate multiple pointers in an NCCL group. -template -void multi_gpu_async_all_reduce_pointers_group( - floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N], - MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) { -#ifdef MULTI_GPU - NVTX_RANGE_FN(); - if (multi_gpu_config->num_processes == 1) { - return; // no multi-GPU, just exit. - } - // 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) { - ncclCheck(ncclAllReduce( - pointers[i], pointers[i], - pointers_sizes[i], - ncclFloatX, ncclAvg, - multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream - )); - } - ncclCheck(ncclGroupEnd()); -#endif -} - void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { NVTX_RANGE_FN(); // double check we forwarded previously, with targets @@ -1168,6 +966,24 @@ void gpt2_multi_gpu_loss_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { 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 gtp2_get_tensor_at_layer(const GPT2 *model, int layer_id, int param_tensor_id) { + 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(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 @@ -1177,8 +993,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) @@ -1191,12 +1005,12 @@ 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<<>>(model->master_weights, params_memory + shard_offset, shard_num_parameters); - cudaCheck(cudaGetLastError()); + init_master_weights = true; } // gradient clipping @@ -1206,11 +1020,24 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // ^1 because of the ncclReduceScatter() in gpt2_multi_gpu_loss_and_grad_reduce, // grads_memory only contains the averaged gradients at the local shard // so we only calculate the grad norm at the grads_memory belonging to the local shard - global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters, main_stream); + cudaCheck(cudaMemsetAsync(grad_norm_squared, 0, sizeof(float), main_stream)); + for(int l = 0; l < model->config.num_layers; ++l) { + for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + if((i < 2 || i > 13) && l != 0) { + continue; + } + // if(2 <= i && i <= 13) + + ShardInfo tensor = gtp2_get_tensor_at_layer(model, l, i); + ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); + ptrdiff_t local_offset_full = tensor.offset + shard.offset; + global_norm_squared(grad_norm_squared, grads_memory + local_offset_full, shard.size, false, 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); + global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters, true, main_stream); } // transfer the gradient norm to CPU float grad_norm_squared_cpu = 0.0f; @@ -1234,34 +1061,15 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // handle adamw for all the transformer blocks for(int l = 0; l < model->config.num_layers; ++l) { - ptrdiff_t full_tensor_offset = 0; - ptrdiff_t partial_tensor_offset = 0; for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { - // skip over preceding layers - if(i > 0) { - full_tensor_offset += model->param_elements[i - 1]; - partial_tensor_offset += model->param_elements[i - 1] / multi_gpu_config->num_processes; + if((i < 2 || i > 13) && l != 0) { + continue; } - ptrdiff_t global_offset; - ptrdiff_t local_offset_full; - ptrdiff_t local_offset_partial; - size_t num_parameters; - if(2 <= i && i <= 13) { - size_t tensor_nelem = model->param_elements[i] / model->config.num_layers; - ShardInfo shard = multi_gpu_get_shard_offset(tensor_nelem, multi_gpu_config, 1); - global_offset = full_tensor_offset + l * tensor_nelem; - local_offset_full = full_tensor_offset + l * tensor_nelem + shard.offset; - local_offset_partial = partial_tensor_offset + l * shard.size; - num_parameters = shard.size; - } else { - if(l != 0) continue; - ShardInfo shard = multi_gpu_get_shard_offset(model->param_elements[i], multi_gpu_config, 1); - global_offset = full_tensor_offset; - local_offset_full = full_tensor_offset + shard.offset; - local_offset_partial = partial_tensor_offset; - num_parameters = shard.size; - } + ShardInfo tensor = gtp2_get_tensor_at_layer(model, l, 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: @@ -1276,18 +1084,23 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl 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<<>>(master_ptr, param_ptr, shard.size); + cudaCheck(cudaGetLastError()); + } // ok finally call the kernel adamw_update(param_ptr, master_ptr, grad_ptr, m_ptr, v_ptr, - num_parameters, learning_rate, + shard.size, learning_rate, beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); cudaCheck(cudaGetLastError()); if (multi_gpu_config->zero_stage == 1) { // gather updated shards of model->params_memory from each process - ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + global_offset, - num_parameters, ncclFloatX, + ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + tensor.offset, + shard.size, ncclFloatX, multi_gpu_config->nccl_comm, main_stream)); } cudaCheck(cudaGetLastError()); @@ -1298,21 +1111,6 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl 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) From ecca976132713052d8b6ce3e1084575afdb099ce Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 15:51:24 +0300 Subject: [PATCH 12/22] re-enable reduce-scatter operation --- llmc/zero.cuh | 32 +++++++++++++++++++++--------- train_gpt2.cu | 54 +++++++++++++++++++++++++-------------------------- 2 files changed, 49 insertions(+), 37 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 76ead6b..6f4c560 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -182,15 +182,18 @@ ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* mult } // 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) depending on the scatter argument template -void multi_gpu_async_all_reduce_pointers_group( +void multi_gpu_async_reduce_pointer_group( floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N], + bool scatter, MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) { -#ifdef MULTI_GPU - NVTX_RANGE_FN(); 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. @@ -200,12 +203,23 @@ void multi_gpu_async_all_reduce_pointers_group( 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) { - ncclCheck(ncclAllReduce( - pointers[i], pointers[i], - pointers_sizes[i], - ncclFloatX, ncclAvg, - multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream - )); + if(!scatter) { + ncclCheck(ncclAllReduce( + pointers[i], pointers[i], + pointers_sizes[i], + ncclFloatX, ncclAvg, + multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream + )); + } else { + 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 diff --git a/train_gpt2.cu b/train_gpt2.cu index 6e09252..95bdac3 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -908,7 +908,7 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { // Accumulate gradients from this layer in a background stream. if(last_step) { - floatX* pointers[] = { + floatX* const pointers[] = { dl_ln1w, dl_ln1b, dl_qkvw, dl_qkvb, dl_attprojw, dl_attprojb, @@ -916,7 +916,7 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { dl_fcw, dl_fcb, dl_fcprojw, dl_fcprojb }; - size_t nelem[] = { + const size_t nelem[] = { C, C, 3 * C * C, 3 * C, C * C, C, @@ -924,8 +924,9 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { 4 * C * C, 4 * C, C * 4 * C, C }; - multi_gpu_async_all_reduce_pointers_group(pointers, nelem, - &multi_gpu_config, main_stream); + multi_gpu_async_reduce_pointer_group(pointers, nelem, + multi_gpu_config.zero_stage == 1, + &multi_gpu_config, main_stream); } } encoder_backward(grads.wte, grads.wpe, scratchX, model->workload_indices, model->bucket_info, @@ -933,9 +934,9 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { // Aggregate all gradients that are not part of the transformer blocks if(last_step) { - floatX* pointers[] = {grads.wte, grads.wpe, grads.lnfw, grads.lnfb}; - size_t nelem[] = {Vp * C, T * C, C, C}; - multi_gpu_async_all_reduce_pointers_group(pointers, nelem, &multi_gpu_config, main_stream); + 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_pointer_group(pointers, nelem, multi_gpu_config.zero_stage == 1, &multi_gpu_config, main_stream); } cudaCheck(cudaDeviceSynchronize()); @@ -1016,35 +1017,32 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // gradient clipping // repurposing this buffer (which isn't needed now) to write grad norm into it float* grad_norm_squared = (float*)model->acts.output; + float grad_norm_squared_cpu = 0.0f; + 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 - cudaCheck(cudaMemsetAsync(grad_norm_squared, 0, sizeof(float), main_stream)); - for(int l = 0; l < model->config.num_layers; ++l) { - for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + // 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) * NUM_PARAMETER_TENSORS, main_stream)); + for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + for(int l = 0; l < model->config.num_layers; ++l) { if((i < 2 || i > 13) && l != 0) { continue; } - // if(2 <= i && i <= 13) - ShardInfo tensor = gtp2_get_tensor_at_layer(model, l, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); - ptrdiff_t local_offset_full = tensor.offset + shard.offset; - global_norm_squared(grad_norm_squared, grads_memory + local_offset_full, shard.size, false, main_stream); + ptrdiff_t offset = tensor.offset + shard.offset; + global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, false, 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, true, 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) { + 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, true, main_stream); + cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); } if(!isfinite(grad_norm_squared_cpu)) { @@ -1101,7 +1099,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // gather updated shards of model->params_memory from each process ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + tensor.offset, shard.size, ncclFloatX, - multi_gpu_config->nccl_comm, main_stream)); + multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); } cudaCheck(cudaGetLastError()); } @@ -1681,7 +1679,7 @@ 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, micro_step != grad_accum_steps - 1); + 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 From 8cba9ba0c0f385a68767026976d5520e2d651b03 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 18:07:33 +0300 Subject: [PATCH 13/22] cleanup --- llmc/zero.cuh | 14 ++++++++------ train_gpt2.cu | 8 ++------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 6f4c560..62b6009 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -129,7 +129,8 @@ MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { 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)); - cudaCheck(cudaEventCreate(&result.compute_nccl_sync)); + // 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; @@ -182,11 +183,12 @@ ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* mult } // 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) depending on the scatter argument +// 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 -void multi_gpu_async_reduce_pointer_group( +void multi_gpu_async_reduce_gradient( floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N], - bool scatter, MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) { if (multi_gpu_config->num_processes == 1) { return; // no multi-GPU, just exit. @@ -203,14 +205,14 @@ void multi_gpu_async_reduce_pointer_group( 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(!scatter) { + 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 { + } else if(multi_gpu_config->zero_stage == 1) { 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( diff --git a/train_gpt2.cu b/train_gpt2.cu index 95bdac3..37fb9d4 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -924,9 +924,7 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { 4 * C * C, 4 * C, C * 4 * C, C }; - multi_gpu_async_reduce_pointer_group(pointers, nelem, - multi_gpu_config.zero_stage == 1, - &multi_gpu_config, main_stream); + 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, @@ -936,7 +934,7 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { 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_pointer_group(pointers, nelem, multi_gpu_config.zero_stage == 1, &multi_gpu_config, main_stream); + multi_gpu_async_reduce_gradient(pointers, nelem, &multi_gpu_config, main_stream); } cudaCheck(cudaDeviceSynchronize()); @@ -1101,7 +1099,6 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl shard.size, ncclFloatX, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); } - cudaCheck(cudaGetLastError()); } } @@ -1700,7 +1697,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)); From 368644959ac0d9dc0180e9ef8724ba369e17dbfb Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 18:14:56 +0300 Subject: [PATCH 14/22] fix non-multi-gpu case --- profile_gpt2.cu | 2 +- train_gpt2.cu | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index 0cb2d77..669e5ef 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -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 diff --git a/train_gpt2.cu b/train_gpt2.cu index 37fb9d4..b7bef1b 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1094,10 +1094,12 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl cudaCheck(cudaGetLastError()); if (multi_gpu_config->zero_stage == 1) { +#if MULTI_GPU // gather updated shards of model->params_memory from each process ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + tensor.offset, shard.size, ncclFloatX, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); +#endif } } } From 8818a751572c4ca1e13521a8d5f8d94e420d748b Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 10 Jun 2024 11:15:26 +0300 Subject: [PATCH 15/22] fixup to profiling script --- profile_gpt2cu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index de2edfd..bcf097c 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -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. From 9d9da3b209a262b492a826036d2ea4f3531e5e91 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 9 Jun 2024 22:19:07 +0300 Subject: [PATCH 16/22] fewer global norm kernel calls --- llmc/global_norm.cuh | 18 +++++++++++++----- train_gpt2.cu | 19 +++++++++++-------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 713da7e..92866e8 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -2,6 +2,7 @@ Global norm, used in gradient clipping */ #include +#include #include // llmc internal imports #include "cuda_common.h" @@ -11,7 +12,7 @@ Global norm, used in gradient clipping // CUDA kernels template -__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 @@ -24,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(accumulator); + // block-level reduce + return blockReduce(accumulator); +} + +template +__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); } @@ -35,7 +41,7 @@ __global__ void global_norm_squared_kernel(float* out, const T* data, size_t cou // kernel launcher template -void global_norm_squared(float* out, const T* values, size_t count, bool reset, 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 @@ -48,7 +54,9 @@ void global_norm_squared(float* out, const T* values, size_t count, bool reset, if(reset) { cudaCheck(cudaMemsetAsync(out, 0, sizeof(float), stream)); } - global_norm_squared_kernel<<>>(out, values, count); + const int gx = CEIL_DIV(grid_size, num_slices); + const int gy = num_slices; + global_norm_squared_kernel<<>>(out, values, count, stride); cudaCheck(cudaGetLastError()); } diff --git a/train_gpt2.cu b/train_gpt2.cu index b7bef1b..3d3ff33 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1021,16 +1021,19 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // 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) * NUM_PARAMETER_TENSORS, main_stream)); + cudaCheck(cudaMemsetAsync(grad_norm_squared, 0, sizeof(float), main_stream)); for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { - for(int l = 0; l < model->config.num_layers; ++l) { - if((i < 2 || i > 13) && l != 0) { - continue; - } - ShardInfo tensor = gtp2_get_tensor_at_layer(model, l, i); + if((i < 2 || i > 13)) { + ShardInfo tensor = gtp2_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, false, main_stream); + global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, false, main_stream); + } else { + ShardInfo tensor = gtp2_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)); @@ -1039,7 +1042,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl } 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, true, main_stream); + 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)); } From 99965a1e3e5ff896e92d4b6fcacf913d6d3ab240 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 10 Jun 2024 18:49:16 +0300 Subject: [PATCH 17/22] fix typo --- train_gpt2.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 3d3ff33..f5db9d9 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -967,7 +967,7 @@ void gpt2_multi_gpu_loss_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { // 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 gtp2_get_tensor_at_layer(const GPT2 *model, int layer_id, int param_tensor_id) { +ShardInfo gpt2_get_tensor_at_layer(const GPT2 *model, int layer_id, int param_tensor_id) { ptrdiff_t offset = 0; for (int i = 0; i < param_tensor_id; i++) { offset += (ptrdiff_t)model->param_elements[i]; @@ -1024,12 +1024,12 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl 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 = gtp2_get_tensor_at_layer(model, 0, i); + 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 = gtp2_get_tensor_at_layer(model, 0, i); + 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, @@ -1065,7 +1065,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl continue; } - ShardInfo tensor = gtp2_get_tensor_at_layer(model, l, i); + ShardInfo tensor = gpt2_get_tensor_at_layer(model, l, 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; From f071efd94fd2784664d55b15430fad2907865efe Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 10 Jun 2024 18:49:16 +0300 Subject: [PATCH 18/22] stride-based adam and copy-and-cast --- llmc/adamw.cuh | 29 ++++++++++++----- llmc/cuda_utils.cuh | 4 +-- train_gpt2.cu | 77 +++++++++++++++++++++++---------------------- 3 files changed, 63 insertions(+), 47 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index eb244a6..94d4642 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -16,9 +16,9 @@ __device__ float lerp(float start, float end, float weight) { } template -__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 +__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 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<<>>(params_memory, master_params_memory, grads_memory, - m_memory, v_memory, num_parameters, + adamw_kernel3<<>>(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()); diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 14ee2b2..1bca60a 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -102,11 +102,11 @@ __device__ float cast_value(__nv_bfloat16 val) { } template -__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(src[idx]); + dst[idx + stride_dst * blockIdx.y] = cast_value(src[idx + stride_src * blockIdx.y]); } } diff --git a/train_gpt2.cu b/train_gpt2.cu index f5db9d9..53010a1 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1059,51 +1059,52 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl unsigned int seed = random_u32(&model->rng_state); // handle adamw for all the transformer blocks - for(int l = 0; l < model->config.num_layers; ++l) { - for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { - if((i < 2 || i > 13) && l != 0) { - continue; - } + 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, l, 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; + 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; + // 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<<>>(master_ptr, param_ptr, shard.size); - cudaCheck(cudaGetLastError()); - } - - // ok finally call the kernel - adamw_update(param_ptr, master_ptr, grad_ptr, - m_ptr, v_ptr, - shard.size, learning_rate, - beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); + 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<<>>(master_ptr, param_ptr, shard.size, + shard.size, tensor.size); cudaCheck(cudaGetLastError()); + } - if (multi_gpu_config->zero_stage == 1) { + // 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 - // gather updated shards of model->params_memory from each process - ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + tensor.offset, - shard.size, ncclFloatX, - multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); + // gather updated shards of model->params_memory from each process + ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + tensor.offset, + shard.size, ncclFloatX, + multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); #endif - } } } From 2a6797b1cd5edc2cb1839f23eca4add78d8dca27 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Thu, 13 Jun 2024 13:58:14 +0200 Subject: [PATCH 19/22] Add debugging tip to README --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index b6c6128..d3dea87 100644 --- a/README.md +++ b/README.md @@ -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: From 3fc8d6192fbe09054b0bc997bb3003fb5db0db86 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 10 Jun 2024 18:49:16 +0300 Subject: [PATCH 20/22] strided all-gather --- train_gpt2.cu | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 53010a1..e2e192b 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1100,10 +1100,15 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl if (multi_gpu_config->zero_stage == 1) { #if MULTI_GPU - // gather updated shards of model->params_memory from each process - ncclCheck(ncclAllGather(param_ptr, (floatX*)model->params_memory + tensor.offset, - shard.size, ncclFloatX, - multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); + 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 } } From 43a903e1b4a90f5b839261ebb31eea04f0117aac Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 13 Jun 2024 23:18:22 +0000 Subject: [PATCH 21/22] revert old printf0, minor touchups and comments, add assert --- llmc/zero.cuh | 1 + train_gpt2.cu | 15 ++++++++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 62b6009..160dae7 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -213,6 +213,7 @@ void multi_gpu_async_reduce_gradient( 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( diff --git a/train_gpt2.cu b/train_gpt2.cu index e2e192b..1d27408 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -73,10 +73,12 @@ cudaStream_t main_stream; MultiGpuConfig multi_gpu_config; // convenience function that only prints if the rank of process is zero -template -void printf0(const char *format, Args... args) { +void printf0(const char *format, ...) { if (multi_gpu_config.process_rank == 0) { - printf(format, args...); + va_list args; + va_start(args, format); + vprintf(format, args); + va_end(args); } } @@ -968,18 +970,17 @@ void gpt2_multi_gpu_loss_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { // 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}; } @@ -1103,7 +1104,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl 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, + 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)); From cbe82e783a569581abe6ca7a61438812e42f112a Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 14 Jun 2024 00:54:58 +0000 Subject: [PATCH 22/22] consolidate memory by deleting the grad activation struct, move some of the last pieces of memory to the forward pass --- llmc/attention.cuh | 34 +++++++++++--------- train_gpt2.cu | 79 +++++++++------------------------------------- 2 files changed, 34 insertions(+), 79 deletions(-) diff --git a/llmc/attention.cuh b/llmc/attention.cuh index 6d4f0e2..efad322 100644 --- a/llmc/attention.cuh +++ b/llmc/attention.cuh @@ -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(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<<>>(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<<>>(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<<>>(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, diff --git a/train_gpt2.cu b/train_gpt2.cu index 1d27408..49e94c3 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -192,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) @@ -221,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) { @@ -257,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) { @@ -312,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 @@ -348,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 @@ -386,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() @@ -770,17 +735,6 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { // 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 @@ -802,10 +756,10 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { 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; @@ -816,11 +770,10 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { // 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 @@ -870,7 +823,7 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { // 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) { @@ -897,8 +850,7 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { // 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); @@ -1154,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));