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

Merge pull request #632 from gordicaleksa/multi_node_my

Browse source 
MPI/TCP/FS for NCCL-init
This commit is contained in:
Andrej 2024-06-24 14:36:30 -07:00 • committed by GitHub
commit 69b50adcf5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 681 additions and 70 deletions
Download patch file Download diff file Expand all files Collapse all files

32
Makefile
View file

@ -188,27 +188,37 @@ else
endif
endif
# Check if OpenMPI and NCCL are available, include them if so, for multi-GPU training
# Check if NCCL is available, include if so, for multi-GPU training
ifeq ($(NO_MULTI_GPU), 1)
$(info → Multi-GPU (OpenMPI + NCCL) is manually disabled)
$(info → Multi-GPU (NCCL) is manually disabled)
else
ifneq ($(OS), Windows_NT)
# Detect if running on macOS or Linux
ifeq ($(SHELL_UNAME), Darwin)
$(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping OpenMPI + NCCL support)
else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists)
$(info ✓ OpenMPI found, OK to train with multiple GPUs)
NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include
NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
NVCC_LDLIBS += -lmpi -lnccl
NVCC_FLAGS += -DMULTI_GPU
$(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping NCCL support)
else ifeq ($(shell dpkg -l | grep -q nccl && echo "exists"), exists)
$(info ✓ NCCL found, OK to train with multiple GPUs)
NVCC_LDLIBS += -lnccl
else
$(info ✗ OpenMPI is not found, disabling multi-GPU support)
$(info ---> On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`)
$(info ✗ NCCL is not found, disabling multi-GPU support)
$(info ---> On Linux you can try install NCCL with `sudo apt install libnccl2 libnccl-dev`)
endif
endif
endif
ifeq ($(NO_USE_MPI), 1)
$(info → MPI is manually disabled)
else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists)
$(info ✓ MPI enabled)
NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include
NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
NVCC_LDLIBS += -lmpi
NVCC_FLAGS += -DUSE_MPI
NVCC_FLAGS += -DMULTI_GPU
else
$(info ✗ MPI not found)
endif
# Precision settings, default to bf16 but ability to override
PRECISION ?= BF16
VALID_PRECISIONS := FP32 FP16 BF16

21
README.md
View file

@ -134,7 +134,9 @@ sudo apt-get -y install libcudnn9-dev-cuda-12
On top of this you need the [cuDNN frontend](https://github.com/NVIDIA/cudnn-frontend/tree/main), but this is just header files. Simply clone the repo to your disk. The Makefile currently looks for it in either your home directory or the current directory. If you have put it elsewhere, add `CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include` to the `make` command-line.
**multi-GPU training using MPI and NCCL**. Make sure you install MPI and NCCL, e.g. on Linux:
## multi-GPU training
Make sure you install MPI and NCCL, e.g. on Linux:
```bash
sudo apt install openmpi-bin openmpi-doc libopenmpi-dev
@ -149,6 +151,23 @@ make train_gpt2cu
mpirun -np <number of GPUs> ./train_gpt2cu
```
or simply run one of our scripts under `./scripts/`.
## multi-node training
Make sure you've installed `NCCL` following instructions from [multi-GPU](#multi-gpu-training) section.
There are 3 ways we currently support that allow you to run multi-node training:
1) Use OpenMPI to exchange nccl id and initialize NCCL. See e.g. `./scripts/multi_node/run_gpt2_124M_mpi.sh` script for details.
2) Use shared file system to init NCCL. See `./scripts/multi_node/run_gpt2_124M_fs.sbatch` script for details.
3) Use TCP sockets to init NCCL. See `./scripts/multi_node/run_gpt2_124M_tcp.sbatch` script for details.
Note:
* If you're running in a slurm environment and your slurm doesn't support PMIx (which we assume will be a common situation given that `slurm-wlm` dropped PMIx support) you will have to use FS (2) or TCP (3) approach. To test whether your slurm supports PMIx run: `srun --mpi=list` and see whether you get `pmix` in the output.
* If you don't have slurm set up, you can kick off a multi-node run using `mpirun` - MPI (1).
None of these 3 methods is superior, we just offer you options so that you can run in your specific environment.
## experiments / sweeps
Just as an example process to sweep learning rates on a machine with 4 GPUs on TinyStories. Run a shell script `sweep.sh` (after you of course `chmod u+x sweep.sh`):

411
llmc/zero.cuh
View file

@ -5,6 +5,12 @@ Utilities for ZeRO sharding
#ifndef LLMC_ZERO_CUH
#define LLMC_ZERO_CUH
#ifdef _WIN32
#include <winsock2.h>
#else
#include <arpa/inet.h>
#endif
#include <cuda_runtime_api.h>
#include <stdint.h>
#include <stdlib.h>
@ -12,8 +18,10 @@ Utilities for ZeRO sharding
#include <stddef.h>
#ifdef MULTI_GPU
#include <mpi.h>
#include <nccl.h>
#ifdef USE_MPI
#include <mpi.h>
#endif
#endif
// ----------------------------------------------------------------------------
@ -36,6 +44,7 @@ void nccl_check(ncclResult_t status, const char *file, int line) {
}
#define ncclCheck(err) (nccl_check(err, __FILE__, __LINE__))
#ifdef USE_MPI
void mpi_check(int status, const char *file, int line) {
if (status != MPI_SUCCESS) {
char mpi_error[4096];
@ -46,15 +55,14 @@ void mpi_check(int status, const char *file, int line) {
}
}
#define mpiCheck(err) (mpi_check(err, __FILE__, __LINE__))
#endif
#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 process_rank; // Rank of this process among all 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.
@ -69,70 +77,381 @@ typedef struct {
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
float* unified_buffer;
#endif
} MultiGpuConfig;
#ifdef MULTI_GPU
#ifdef _WIN32
void send_nccl_id_to_clients_windows(ncclUniqueId *nccl_id, SOCKET client_sockets[], int num_clients) {
for (int i = 0; i < num_clients; ++i) {
if (send(client_sockets[i], (const char *)nccl_id, sizeof(*nccl_id), 0) == SOCKET_ERROR) {
printf("Failed to send nccl_id");
WSACleanup();
exit(EXIT_FAILURE);
}
closesocket(client_sockets[i]);
}
}
#else
void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int num_clients) {
for (int i = 0; i < num_clients; ++i) {
if (send(client_sockets[i], nccl_id, sizeof(*nccl_id), 0) == -1) {
printf("Failed to send nccl_id");
exit(EXIT_FAILURE);
}
close(client_sockets[i]);
}
}
#endif
#ifdef _WIN32
// Same as get_nccl_id_via_tcp but for Windows
ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* server_ip) {
ncclUniqueId nccl_id;
int SERVER_PORT = 12345; // hardcoded an arbitrary port number between 1024 and 49151 (registered ports)
WSADATA wsaData;
if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
printf("WSAStartup failed");
exit(EXIT_FAILURE);
}
if (result->process_rank == 0) {
ncclCheck(ncclGetUniqueId(&nccl_id));
int MAX_CLIENTS = result->num_processes - 1;
SOCKET client_sockets[MAX_CLIENTS];
int num_clients = 0;
SOCKET server_socket, new_socket;
struct sockaddr_in address;
int addrlen = sizeof(address);
// Step 1) create a server TCP socket
if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) {
printf("Socket failed");
WSACleanup();
exit(EXIT_FAILURE);
}
// Step 2) set the server address and port
address.sin_family = AF_INET; // IPv4
address.sin_addr.s_addr = inet_addr(server_ip);
address.sin_port = htons(SERVER_PORT);
// Step 3) bind the socket to the address and port
if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) == SOCKET_ERROR) {
printf("Bind failed");
closesocket(server_socket);
WSACleanup();
exit(EXIT_FAILURE);
}
// Step 4) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server
if (listen(server_socket, MAX_CLIENTS) == SOCKET_ERROR) {
printf("Listen failed");
closesocket(server_socket);
WSACleanup();
exit(EXIT_FAILURE);
}
// Step 5) accept connections from clients
printf("Waiting for clients to connect...\n");
while (num_clients < MAX_CLIENTS) {
if ((new_socket = accept(server_socket, (struct sockaddr *)&address, &addrlen)) == INVALID_SOCKET) {
printf("Accept failed");
closesocket(server_socket);
WSACleanup();
exit(EXIT_FAILURE);
}
client_sockets[num_clients++] = new_socket;
printf("Client %d connected\n", num_clients);
}
// Step 6) send the NCCL ID to all clients
send_nccl_id_to_clients_windows(&nccl_id, client_sockets, num_clients);
printf("NCCL ID sent to all clients\n");
closesocket(server_socket);
} else {
int num_connection_attempts = 5;
int time_to_sleep = 2;
SOCKET client_socket;
struct sockaddr_in serv_addr;
// Step 1) create a client TCP socket
if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) {
printf("Socket creation error");
WSACleanup();
exit(EXIT_FAILURE);
}
// Step 2) set the server address and port
serv_addr.sin_family = AF_INET;
serv_addr.sin_port = htons(SERVER_PORT);
if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) {
printf("Invalid address or address not supported");
closesocket(client_socket);
WSACleanup();
exit(EXIT_FAILURE);
}
// Step 3) Try to connect to the server - retry up to `num_connection_attempts` times if the connection fails
while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) == SOCKET_ERROR) {
printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep);
if (--num_connection_attempts == 0) {
printf("Failed to connect to the server\n");
closesocket(client_socket);
WSACleanup();
exit(EXIT_FAILURE);
}
Sleep(time_to_sleep * 1000);
}
// Step 4) receive the NCCL ID from the server
if (recv(client_socket, (char *)&nccl_id, sizeof(nccl_id), 0) <= 0) {
printf("Failed to receive nccl_id");
closesocket(client_socket);
WSACleanup();
exit(EXIT_FAILURE);
}
printf("Received NCCL ID\n");
closesocket(client_socket);
}
WSACleanup();
return nccl_id;
}
#else
ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) {
ncclUniqueId nccl_id;
int SERVER_PORT = 12345; // hardcoded an arbitrary port number between 1024 and 49151 (registered ports)
if (result->process_rank == 0) {
ncclCheck(ncclGetUniqueId(&nccl_id));
int MAX_CLIENTS = result->num_processes - 1;
int client_sockets[MAX_CLIENTS];
int num_clients = 0;
int server_socket, new_socket;
struct sockaddr_in address;
int addrlen = sizeof(address);
int opt = 1;
// Step 1) create a server TCP socket
if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
printf("Socket failed");
exit(EXIT_FAILURE);
}
// Step 2) set socket options
// SOL_SOCKET - means that option is configured at socket level
// SO_REUSEADDR - allows to bind to an address which is in a TIME_WAIT state (already used by another socket) - useful when restarting the server
// SO_REUSEPORT - allows to bind to the same port multiple times
if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)) < 0) {
printf("Setsockopt failed");
exit(EXIT_FAILURE);
}
// Step 3) set the server address and port
address.sin_family = AF_INET; // IPv4
address.sin_addr.s_addr = inet_addr(server_ip); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet
address.sin_port = htons(SERVER_PORT);
// Step 4) bind the socket to the address and port
if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) < 0) {
printf("Bind failed");
exit(EXIT_FAILURE);
}
// Step 5) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server
if (listen(server_socket, MAX_CLIENTS) < 0) {
printf("Listen failed");
exit(EXIT_FAILURE);
}
// Step 6) accept connections from clients
printf("Waiting for clients to connect...\n");
while (num_clients < MAX_CLIENTS) {
if ((new_socket = accept(server_socket, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) {
printf("Accept failed");
exit(EXIT_FAILURE);
}
client_sockets[num_clients++] = new_socket;
printf("Client %d connected\n", num_clients);
}
// Step 7) send the NCCL ID to all clients
send_nccl_id_to_clients(&nccl_id, client_sockets, num_clients);
printf("NCCL ID sent to all clients\n");
close(server_socket);
} else {
int num_connection_attempts = 5;
int time_to_sleep = 2;
int client_socket;
struct sockaddr_in serv_addr;
// Step 1) create a client TCP socket
if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
printf("Socket creation error");
exit(EXIT_FAILURE);
}
// Step 2) set the server address and port
serv_addr.sin_family = AF_INET;
serv_addr.sin_port = htons(SERVER_PORT);
if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) {
printf("Invalid address or address not supported");
exit(EXIT_FAILURE);
}
// Step 3) Try to connect to the server - retry up to `num_connection_attempts` times if the connection fails
while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) {
printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep);
if (--num_connection_attempts == 0) {
printf("Failed to connect to the server\n");
exit(EXIT_FAILURE);
}
sleep(time_to_sleep);
}
// Step 4) receive the NCCL ID from the server
if (recv(client_socket, &nccl_id, sizeof(nccl_id), 0) <= 0) {
printf("Failed to receive nccl_id");
exit(EXIT_FAILURE);
}
printf("Received NCCL ID\n");
close(client_socket);
}
return nccl_id;
}
#endif
ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) {
// Works assuming that the filesystem is shared among all processes
ncclUniqueId nccl_id;
FILE* idFile;
static char filename[1024];
snprintf(filename, sizeof(filename), "%s/ncclUniqueId.sync", fs_path);
if (result->process_rank != 0) { // client processse should wait for the server to write to the file
// This is a naive and not 100% robust way to synchronize the processes but it should work almost always
sleep(2);
}
if (result->process_rank == 0) {
ncclCheck(ncclGetUniqueId(&nccl_id));
idFile = fopen(filename, "wb");
assert(idFile != NULL);
fwrite(&nccl_id, sizeof(nccl_id), 1, idFile);
fclose(idFile);
} else {
// Other ranks wait until the file is available and read the unique ID
do {
sleep(1); // 1 second
idFile = fopen(filename, "rb");
if (idFile != NULL) break;
} while (idFile == NULL);
freadCheck(&nccl_id, sizeof(nccl_id), 1, idFile);
fclose(idFile);
}
return nccl_id;
}
#ifdef USE_MPI
// 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;
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]; }
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));
// 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++;
}
}
// 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;
free(all_hostsname_hashes);
return local_device_idx;
}
#endif
MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) {
#endif
MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, int gpus_per_node, char* server_ip, char* fs_path, char* init_method) {
#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));
// Get nccl_id using MPI, TCP, or FS (file system synchronization) methods
// On newer slurm versions (slurm-wlm package) PMIx is disabled so we can not use MPI for NCCL init in multi node setup
if (strcmp(init_method, "mpi") == 0) {
#ifdef USE_MPI
mpiCheck(MPI_Init(NULL, NULL));
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);
if (result.process_rank == 0) {
ncclCheck(ncclGetUniqueId(&nccl_id));
}
mpiCheck(MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
#else
printf("MPI support is disabled. Please enable MPI support to use MPI-based NCCL-init method.\n");
exit(EXIT_FAILURE);
#endif
} else {
result.process_rank = process_rank;
result.num_processes = num_processes;
result.local_device_idx = process_rank % gpus_per_node;
if (strcmp(init_method, "tcp") == 0) {
#ifdef _WIN32
nccl_id = get_nccl_id_via_tcp_windows(&result, server_ip);
#else
nccl_id = get_nccl_id_via_tcp(&result, server_ip);
#endif
} else if (strcmp(init_method, "fs") == 0) {
nccl_id = get_nccl_id_via_fs(&result, fs_path);
} else {
printf("Invalid NCCL-init method\n");
exit(EXIT_FAILURE);
}
}
mpiCheck(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
cudaCheck(cudaSetDevice(result.local_device_idx));
ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank));
cudaCheck(cudaStreamCreate(&result.nccl_stream));
// event without timing for maximum performance
cudaCheck(cudaEventCreate(&result.compute_nccl_sync, cudaEventDisableTiming));
nvtxNameCudaStreamA(result.nccl_stream, "nccl stream");
nvtxNameCudaEventA(result.compute_nccl_sync, "nccl compute sync");
cudaCheck(cudaMallocManaged(&result.unified_buffer, sizeof(float)));
return result;
#else
printf("Multi-GPU support is disabled. Using a single GPU.\n");
@ -150,15 +469,19 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) {
ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm));
cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream));
cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync));
cudaCheck(cudaFree(multi_gpu_config->unified_buffer));
#ifdef USE_MPI
mpiCheck(MPI_Finalize());
#endif
#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));
ncclCheck(ncclAllReduce(multi_gpu_config->unified_buffer, multi_gpu_config->unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream));
}
cudaCheck(cudaDeviceSynchronize());
#endif
}

8
profile_gpt2.cu
View file

@ -28,7 +28,13 @@ the profile.ncu-rep from a cloud box to local to pretty view.
#include "train_gpt2.cu"
int main(int argc, char *argv[]) {
multi_gpu_config = multi_gpu_config_init(&argc, &argv);
char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi"
int num_processes = -1; // doesn't matter when using MPI
int process_rank = -1; // doesn't matter when using MPI
int gpus_per_node = -1; // doesn't matter when using MPI
char server_ip[256] = ""; // doesn't matter when using MPI
char fs_path[256] = ""; // doesn't matter when using MPI
multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method);
common_start(true, true);
// build the GPT-2 model from a checkpoint

85
scripts/multi_node/run_gpt2_124M_fs.sbatch Executable file
View file

@ -0,0 +1,85 @@
#!/bin/bash
#SBATCH --job-name=llmc-multinode # job name
#SBATCH --output=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.log # output file
#SBATCH --error=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.err # error file
#SBATCH --partition=llmc # Specify the GPU partition
#SBATCH --ntasks=16 # total number of processes to launch on all nodes
#SBATCH --nodes=2 # total number of nodes
#SBATCH --ntasks-per-node=8 # assuming each node has 8 gpus
#SBATCH --gres=gpu:8 # request 8 gpus from each node
# NOTE: change the above slurm arguments to match your system!
# Run with `sbatch <path_to_this_script.sh>`
make train_gpt2cu USE_CUDNN=1 NO_USE_MPI=1
# NOTE: change the following to match your system
binary_path="/home/ubuntu/llm.c/train_gpt2cu"
out_dir="/ephemeral/data/fineweb/log_gpt2_124M_multi"
train_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_train_*.bin'
val_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_val_*.bin'
sync_fs_path=$out_dir # needs to be a shared filesystem path that all nodes can access
# In case the file system is shared this is a no-op.
# Otherwise, we need to copy the binary to all nodes.
current_user=$USER
hosts=$(scontrol show hostnames $SLURM_JOB_NODELIST) # get the hostnames of the allocated nodes
current_host=$(hostname)
for host in $hosts; do
if [ $host == $current_host ]; then
continue
fi
echo "copying $binary_path to $current_user@$host"
scp -r $binary_path $current_user@$host:$binary_path
done
# Use this for NCCL debugging if you run into issues
# export NCCL_DEBUG=INFO
# export NCCL_DEBUG_SUBSYS=ALL
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
# Optimization flags
export NCCL_NET_GDR_LEVEL=2 # use GPUDirect RDMA - allows for direct memory access between GPUs across different nodes by bypassing the CPU
export NCCL_IB_DISABLE=0 # use InfiniBand if available
# NOTE: change the following environment variables to match your system - or comment them out if you don't need them
export NCCL_SOCKET_IFNAME=ens17
export OMPI_MCA_btl_tcp_if_include=ens17
export NCCL_P2P_LEVEL=PXB
if [ -z "$SLURM_JOB_ID" ]; then
echo "Make sure you're running in a SLURM environment. Did you forget to run with sbatch? Aborting."
exit 1
else
DATESTRING=`date "+%Y-%m-%dT%H:%M:%S"`
echo "Running in a SLURM environment (job ID: $SLURM_JOB_ID, user: $current_user)"
echo "Running on hosts: $(echo $(scontrol show hostname))"
echo "$DATESTRING"
fi
srun -l -u bash -c "
$binary_path \
-i '$train_data_path' \
-j '$val_data_path' \
-o $out_dir \
-v 250 -s 20000 -g 144 \
-h 1 \
-b 64 -t 1024 \
-d 2097152 \
-r 0 \
-z 1 \
-c 0.1 \
-l 0.0006 \
-q 0.0 \
-u 700 \
-n 5000 \
-y 1 \
-e d12 \
-pn \$SLURM_NTASKS \
-pr \$SLURM_PROCID \
-pg \$SLURM_NTASKS_PER_NODE \
-pf $sync_fs_path \
-pi "fs" \
"
echo "$DATESTRING"

49
scripts/multi_node/run_gpt2_124M_mpi.sh Executable file
View file

@ -0,0 +1,49 @@
make train_gpt2cu USE_CUDNN=1
# NOTE: change the following to match your system
binary_path="/home/ubuntu/llm.c/train_gpt2cu"
out_dir="/ephemeral/data/fineweb/log_gpt2_124M_multi"
train_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_train_*.bin'
val_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_val_*.bin'
# You can find these names either in `/etc/hosts`` file or in the terminal (user@host:~$).
host1="h100-node-1-0" # master and worker node
host2="h100-node-1-1" # worker node
# In case the file system is shared this is a no-op.
# Otherwise, we need to copy the binary to all nodes.
scp -r $binary_path $USER@$host2:$binary_path
# Use this for NCCL debugging if you run into issues
# export NCCL_DEBUG=INFO
# export NCCL_DEBUG_SUBSYS=ALL
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
# Optimization flags
export NCCL_NET_GDR_LEVEL=2 # use GPUDirect RDMA - allows for direct memory access between GPUs across different nodes by bypassing the CPU
export NCCL_IB_DISABLE=0 # use InfiniBand if available
# NOTE: change the following environment variables to match your system - or comment them out if you don't need them
export NCCL_SOCKET_IFNAME=ens17
export OMPI_MCA_btl_tcp_if_include=ens17
export NCCL_P2P_LEVEL=PXB
mpirun -np 16 --host $host1:8,$host2:8 \
$binary_path \
-i "$train_data_path" \
-j "$val_data_path" \
-o $out_dir \
-v 250 -s 20000 -g 144 \
-h 1 \
-b 64 -t 1024 \
-d 2097152 \
-r 0 \
-z 1 \
-c 0.1 \
-l 0.0006 \
-q 0.1 \
-u 700 \
-n 1000 \
-y 0 \
-e d12 \
-pi "mpi" \

86
scripts/multi_node/run_gpt2_124M_tcp.sbatch Executable file
View file

@ -0,0 +1,86 @@
#!/bin/bash
#SBATCH --job-name=llmc-multinode # job name
#SBATCH --output=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.log # output file
#SBATCH --error=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.err # error file
#SBATCH --partition=llmc # Specify the GPU partition
#SBATCH --ntasks=16 # total number of processes to launch on all nodes
#SBATCH --nodes=2 # total number of nodes
#SBATCH --ntasks-per-node=8 # assuming each node has 8 gpus
#SBATCH --gres=gpu:8 # request 8 gpus from each node
# NOTE: change the above slurm arguments to match your system!
# Run with `sbatch <path_to_this_script.sh>`
make train_gpt2cu USE_CUDNN=1 NO_USE_MPI=1
# NOTE: change the following to match your system
binary_path="/home/ubuntu/llm.c/train_gpt2cu"
out_dir="/ephemeral/data/fineweb/log_gpt2_124M_multi"
train_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_train_*.bin'
val_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_val_*.bin'
# NOTE: change the server_ip to the IP address of the machine that is running process zero
server_ip="10.0.1.220"
# In case the file system is shared this is a no-op.
# Otherwise, we need to copy the binary to all nodes.
current_user=$USER
hosts=$(scontrol show hostnames $SLURM_JOB_NODELIST) # get the hostnames of the allocated nodes
current_host=$(hostname)
for host in $hosts; do
if [ $host == $current_host ]; then
continue
fi
echo "copying $binary_path to $current_user@$host"
scp -r $binary_path $current_user@$host:$binary_path
done
# Use this for NCCL debugging if you run into issues
# export NCCL_DEBUG=INFO
# export NCCL_DEBUG_SUBSYS=ALL
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
# Optimization flags
export NCCL_NET_GDR_LEVEL=2 # use GPUDirect RDMA - allows for direct memory access between GPUs across different nodes by bypassing the CPU
export NCCL_IB_DISABLE=0 # use InfiniBand if available
# NOTE: change the following environment variables to match your system - or comment them out if you don't need them
export NCCL_SOCKET_IFNAME=ens17
export OMPI_MCA_btl_tcp_if_include=ens17
export NCCL_P2P_LEVEL=PXB
if [ -z "$SLURM_JOB_ID" ]; then
echo "Make sure you're running in a SLURM environment. Did you forget to run with sbatch? Aborting."
exit 1
else
DATESTRING=`date "+%Y-%m-%dT%H:%M:%S"`
echo "Running in a SLURM environment (job ID: $SLURM_JOB_ID, user: $current_user)"
echo "Running on hosts: $(echo $(scontrol show hostname))"
echo "$DATESTRING"
fi
srun -l -u bash -c "
$binary_path \
-i '$train_data_path' \
-j '$val_data_path' \
-o $out_dir \
-v 250 -s 20000 -g 144 \
-h 1 \
-b 64 -t 1024 \
-d 2097152 \
-r 0 \
-z 1 \
-c 0.1 \
-l 0.0006 \
-q 0.0 \
-u 700 \
-n 5000 \
-y 1 \
-e d12 \
-pn \$SLURM_NTASKS \
-pr \$SLURM_PROCID \
-pg \$SLURM_NTASKS_PER_NODE \
-ps $server_ip \
-pi "tcp" \
"
echo "$DATESTRING"

8
test_gpt2.cu
View file

@ -89,7 +89,13 @@ float* float_cpu_malloc_and_point_parameters(FloatParameterTensors* params, size
}
int main(int argc, char *argv[]) {
multi_gpu_config = multi_gpu_config_init(&argc, &argv);
char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi"
int num_processes = -1; // doesn't matter when using MPI
int process_rank = -1; // doesn't matter when using MPI
int gpus_per_node = -1; // doesn't matter when using MPI
char server_ip[256] = ""; // doesn't matter when using MPI
char fs_path[256] = ""; // doesn't matter when using MPI
multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method);
common_start(false, true);
// set the right paths

51
train_gpt2.cu
View file

@ -2,6 +2,10 @@
GPT-2 Transformer Neural Net training loop. See README.md for usage.
*/
#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#endif
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
@ -893,12 +897,15 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) {
}
// Compute sum of a single CPU value across all GPU processes. No-op when multi-GPU is disabled.
float multi_gpu_cpu_float_sum(float value) {
float multi_gpu_cpu_float_sum(float value, MultiGpuConfig* multi_gpu_config) {
#ifdef MULTI_GPU
// note MPI doesn't support all reduce with mean, only sum
float result;
mpiCheck(MPI_Allreduce(&value, &result, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD));
return result;
if (multi_gpu_config->num_processes == 1) return value;
float* unified_buffer = multi_gpu_config->unified_buffer;
*unified_buffer = value;
ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream));
cudaCheck(cudaDeviceSynchronize());
return *unified_buffer;
#else
return value;
#endif
@ -912,7 +919,7 @@ void gpt2_multi_gpu_loss_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) {
// 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;
model->accumulated_mean_loss = multi_gpu_cpu_float_sum(model->mean_loss, multi_gpu_config) / multi_gpu_config->num_processes;
#endif
cudaCheck(cudaDeviceSynchronize());
}
@ -990,7 +997,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl
global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream);
cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost));
// further sum the (partial) squared norm across all GPUs (see comment ^1 above)
grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu);
grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu, multi_gpu_config);
} 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
@ -1315,14 +1322,19 @@ void error_usage() {
// memory management
fprintf(stderr, " -z <int> zero_stage, Zero Optimization Stage, 0,1,2,3 (default = 0)\n");
fprintf(stderr, " -r <int> recompute: less memory but less speed. (default = 1), 0|1|2 = none,gelu,gelu+ln\n");
// multi-node settings
fprintf(stderr, " -pn <int> num_processes (default = 1)\n");
fprintf(stderr, " -pr <int> process_rank (default = 0)\n");
fprintf(stderr, " -pg <int> gpus_per_node (default = 8)\n");
fprintf(stderr, " -pm <string> nccl_init_method: tcp,fs,mpi (default = mpi)\n");
fprintf(stderr, " -ps <string> server_ip - used only when nccl_init_method is tcp (default = -1)\n");
fprintf(stderr, " -pp <string> fs_path - used only when nccl_init_method is fs (default = /tmp)\n");
exit(EXIT_FAILURE);
}
// ----------------------------------------------------------------------------
// main training loop
int main(int argc, char *argv[]) {
multi_gpu_config = multi_gpu_config_init(&argc, &argv);
// read in the (optional) command line arguments
const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin";
const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin";
@ -1349,10 +1361,17 @@ int main(int argc, char *argv[]) {
int recompute = 1; // recompute during backward setting, 0 = none, 1 = recompute gelu
int zero_stage = 0; // Zero Optimization Stage for Multi-GPU training
int hellaswag_eval = 0;
// multi-node settings
int num_processes = 1; // this should be set by the slurm environment
int process_rank = 0; // this should be set by the slurm environment
int gpus_per_node = 8; // this should be set by the slurm environment
char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi"
char server_ip[256] = ""; // used if init_method set to "tcp" -> set to your server ip address
char fs_path[256] = ""; // used if init_method set to "fs" -> set to a shared filesystem path
for (int i = 1; i < argc; i+=2) {
if (i + 1 >= argc) { error_usage(); } // must have arg after flag
if (argv[i][0] != '-') { error_usage(); } // must start with dash
if (strlen(argv[i]) != 2) { error_usage(); } // must be -x (one dash, one letter)
if (!(strlen(argv[i]) == 2 || strlen(argv[i]) == 3)) { error_usage(); } // must be -x[y] (one dash, one or two letters)
// read in the args
if (argv[i][1] == 'i') { train_data_pattern = argv[i+1]; }
else if (argv[i][1] == 'j') { val_data_pattern = argv[i+1]; }
@ -1379,8 +1398,16 @@ int main(int argc, char *argv[]) {
else if (argv[i][1] == 'r') { recompute = atoi(argv[i+1]); }
else if (argv[i][1] == 'h') { hellaswag_eval = atoi(argv[i+1]); }
else if (argv[i][1] == 'k') { lr_scheduler_type = argv[i+1]; }
else if (argv[i][1] == 'p' && argv[i][2] == 'i') { strcpy(nccl_init_method, argv[i+1]); }
else if (argv[i][1] == 'p' && argv[i][2] == 'f') { strcpy(fs_path, argv[i+1]); }
else if (argv[i][1] == 'p' && argv[i][2] == 's') { strcpy(server_ip, argv[i+1]); }
else if (argv[i][1] == 'p' && argv[i][2] == 'n') { num_processes = atoi(argv[i+1]); }
else if (argv[i][1] == 'p' && argv[i][2] == 'r') { process_rank = atoi(argv[i+1]); }
else if (argv[i][1] == 'p' && argv[i][2] == 'g') { gpus_per_node = atoi(argv[i+1]); }
else { error_usage(); }
}
multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method);
// should do a bit more error checking here
assert(warmup_iterations >= 0);
if (output_log_dir != NULL) {
@ -1579,7 +1606,7 @@ int main(int argc, char *argv[]) {
val_loss += model.mean_loss;
}
val_loss /= val_num_batches;
val_loss = multi_gpu_cpu_float_sum(val_loss) / multi_gpu_config.num_processes;
val_loss = multi_gpu_cpu_float_sum(val_loss, &multi_gpu_config) / multi_gpu_config.num_processes;
printf0("val loss %f\n", val_loss);
logger_log_val(&logger, step, val_loss);
}
@ -1598,7 +1625,7 @@ int main(int argc, char *argv[]) {
eval_acc_norm += (float)correct;
}
// careful because not all ranks may have the exact same allocation of number of examples
eval_acc_norm = multi_gpu_cpu_float_sum(eval_acc_norm);
eval_acc_norm = multi_gpu_cpu_float_sum(eval_acc_norm, &multi_gpu_config);
printf0("HellaSwag: %d/%d = %f\n", (int)eval_acc_norm, eval_loader.num_examples, eval_acc_norm / eval_loader.num_examples);
logger_log_eval(&logger, step, eval_acc_norm / eval_loader.num_examples);
}