A4 — Multi-GPU with NCCL and ZeRO-1

Files: llmc/zero.cuh, train_gpt2.cu
Needs: 2+ NVIDIA GPUs, NCCL
Builds on: A0–A3
Check with: ./grade A4

---

Overview

Training GPT-2 124M fits on a single GPU. Training GPT-2 1.5B, GPT-3, or Llama-3 does not. A4 implements distributed data-parallel training across multiple GPUs using NCCL and an optional ZeRO-1 optimizer state sharding scheme.

By the end of A4, the same training binary will scale from 1 to N GPUs (N ≤ 8 on a single node, or more across nodes). The loss curve must be identical across GPU counts — run with 1 GPU or 4 GPUs and you get the same answer.

That reproducibility requirement is the hardest part of the assignment. It forces you to think carefully about communication ordering, gradient averaging, and deterministic operations.

This one requires at least 2 NVIDIA GPUs plus NCCL. ./grade A4 reports every check as SKIP otherwise.

To list your work:

grep -rn "TODO (A4)" llmc/ train_gpt2.cu

When these notes and the in-file HINT disagree, the code is authoritative — please report the discrepancy so it can be fixed.

Note that everything you write in zero.cuh lives inside #ifdef MULTI_GPU. The single-GPU fallbacks are provided and already work, which is why A3 ran fine without any of this.

---

Where this fits

flowchart LR
    A0["A0
tokens, batches, PRNG"] --> A1["A1
forward pass (CPU)"]
    A1 --> A2["A2
backward + AdamW (CPU)"]
    A2 --> A3["A3
port to CUDA (1 GPU)"]
    A3 --> A4["A4
scale to N GPUs"]
    A3 --> A5["A5
schedules, clipping, checkpoints"]
    style A4 fill:#2563eb,stroke:#1e40af,color:#fff

One GPU has a fixed amount of memory and a fixed amount of arithmetic throughput. A4 removes both limits by putting a copy of the model on every GPU, giving each a different slice of the batch, and averaging the gradients before the optimizer step. The hard part is not making it run — it is making it produce exactly the same numbers regardless of how many GPUs you use.

flowchart TB
    batch["global batch
d tokens"] --> s0["GPU 0
micro-batch 0"]
    batch --> s1["GPU 1
micro-batch 1"]
    batch --> sn["GPU N-1
micro-batch N-1"]
    s0 --> g0["grads"]
    s1 --> g1["grads"]
    sn --> gn["grads"]
    g0 --> ar["NCCL AllReduce (ncclAvg)
every GPU ends up with the mean gradient"]
    g1 --> ar
    gn --> ar
    ar --> upd["identical AdamW step on every GPU
→ identical weights"]

Where the speedup actually comes from. Communication is overlapped with computation. The backward pass runs from layer L-1 down to 0, so by the time layer L-2 is being differentiated, layer L-1’s gradients are already crossing the interconnect on a separate CUDA stream. Done right, most of the transfer is free; done wrong (a cudaStreamSynchronize in the loop), you serialise compute and comms and the scaling collapses.

ZeRO-1 in one picture. Optimizer state is 3 extra copies of every parameter (m, v, and the FP32 master weights) — for GPT-2 124M that is more memory than the weights themselves. ZeRO-1 gives each GPU responsibility for 1/N of it:

flowchart LR
    subgraph z0["ZeRO-0: every GPU holds everything"]
        direction TB
        a["GPU0: m,v,master for ALL params"]
        b["GPU1: m,v,master for ALL params"]
    end
    subgraph z1["ZeRO-1: optimizer state sharded"]
        direction TB
        c["GPU0: m,v,master for params[0..P/2)"]
        d["GPU1: m,v,master for params[P/2..P)"]
    end
    z0 -->|"ReduceScatter instead of AllReduce,
+ one AllGather after the update"| z1

The trade is one extra AllGather per step in exchange for an N-fold cut in optimizer memory — which is what lets a model that would not otherwise fit train at all.

Why “same loss on 1 and 4 GPUs” is the real test. It is easy to write multi-GPU code that trains and converges but is subtly wrong: a missed stream sync, a double division by N, a non-deterministic atomicAdd left over from A3. None of those necessarily hurt the loss curve enough to notice. Requiring exact agreement across GPU counts catches all of them, which is why it is graded rather than just encouraged.

---

Background Reading

Source	Read	Why it matters here
The #ifdef MULTI_GPU blocks in llmc/zero.cuh	all of them	The struct fields, the stream and event, and the single-GPU fallbacks that are already written. Your job is the bodies inside these guards.
scripts/multi_node/	all three scripts	Working launch commands for the MPI, shared-filesystem, and TCP init paths. These are the ground truth for how the binary is actually invoked.
NVIDIA, NCCL User Guide	“Collective Operations”, “Group Calls”	AllReduce vs ReduceScatter vs AllGather, and why ncclGroupStart/End matters for launch overhead when you issue 16 collectives in a row.
NVIDIA, Massively Scale Deep Learning with NCCL	all	Why ring-AllReduce moves 2(N-1)/N times the data per GPU regardless of N, which is what makes data parallelism scale.
Li et al. 2020, PyTorch Distributed	§3.2 (gradient bucketing), §4	How DDP overlaps the backward pass with communication. Part 4 is the same design; this paper explains the measurements behind it.
Rajbhandari et al. 2020, ZeRO	§3 and Fig. 1	The memory accounting for stages 1–3. Figure 1 is the clearest single picture of what sharding optimizer state buys you. You implement stage 1 only.
Goyal et al. 2017, Accurate, Large Minibatch SGD	§2.1	Why averaging gradients across N workers is equivalent to one large batch — the assumption behind ncclAvg and behind holding -d fixed when you compare GPU counts.

---

System Model

Data parallelism: Each GPU holds a complete copy of the model weights and processes a different micro-batch. After the backward pass, gradients are averaged across GPUs so every GPU computes the same parameter update.

Gradient communication: NCCL AllReduce (ZeRO-0) or ReduceScatter (ZeRO-1), overlapped with the backward pass on a separate stream.

ZeRO-1: Shards the optimizer state (m, v, master weights) across GPUs. Each GPU updates 1/N of the parameters, then an AllGather synchronizes the updated parameters back to everyone. This reduces optimizer memory by N× at the cost of one extra AllGather per step.

GPU 0: owns params[0 .. P/N-1]     → update them → AllGather
GPU 1: owns params[P/N .. 2P/N-1]  → update them → AllGather
...

---

Part 1 — Process Initialization (llmc/zero.cuh)

multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, init_method)

Returns a MultiGpuConfig. Three init methods:

• "mpi": rank and size come from MPI; rank 0 generates the NCCL unique ID and MPI_Bcasts it.
• "tcp": rank 0 runs a TCP server and sends the ID to each client that connects.
• "fs": rank 0 writes the ID to a file on a shared filesystem; other ranks poll until it appears.

For "mpi":

MPI_Init(NULL, NULL);
MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank);
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) ncclGetUniqueId(&nccl_id);
MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD);

For "tcp" and "fs", rank/size/gpus_per_node come from the command line instead, and local_device_idx = process_rank % gpus_per_node.

Whichever path produced nccl_id, finish with:

cudaSetDevice(result.local_device_idx);
ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank);
cudaStreamCreate(&result.nccl_stream);
cudaEventCreate(&result.compute_nccl_sync, cudaEventDisableTiming);
cudaMallocManaged(&result.unified_buffer, sizeof(float));

The MPI path is guarded by #ifdef USE_MPI — when MPI is not compiled in, init_method == "mpi" must print a clear error and exit rather than silently misbehaving.

multi_gpu_config_free(config) — destroy the communicator, stream, and event; free the unified buffer; MPI_Finalize() under USE_MPI.

multi_gpu_barrier(config) — when num_processes > 1, issue an ncclAllReduce of 1 float on config->unified_buffer (the value is irrelevant), then cudaDeviceSynchronize.

multi_gpu_get_local_device_idx(process_rank, num_processes) — MPI path only. gethostname, truncate at the first ., hash it, MPI_Allgather all hashes, then count how many lower ranks share your hash. That count is your local GPU index.

get_nccl_id_via_fs(result, fs_path) — rank 0 writes the ID to {fs_path}/ncclUniqueId.sync; other ranks sleep and poll fopen until it exists, then read it.

send_nccl_id_to_clients(...) — helper for the TCP path: send the ID to each connected socket and close it.

---

Part 2 — Tensor Sharding (llmc/zero.cuh)

multi_gpu_get_shard_offset(elements, config, shard_at_stage)

If config->zero_stage >= shard_at_stage, return this rank’s shard:

shard_size   = elements / num_processes
shard_offset = process_rank * shard_size

Error out if elements % num_processes != 0. Otherwise (not sharding at this stage) return the full range {0, elements} — that branch is already written for you and is what makes single-GPU runs work.

set_zero_configs(config, zero_stage, total_parameters)

Defaults to zero_stage = 0, shard_num_parameters = total_parameters. For zero_stage == 1, enable sharding only when total_parameters % num_processes == 0; otherwise print the “Can’t equally partition parameters” message and fall back to stage 0. Stages 2 and 3 are not supported and fall back to 0.

---

Part 3 — Gradient Reduction (llmc/zero.cuh)

multi_gpu_async_reduce_gradient(pointers[N], sizes[N], config, compute_stream)

The key function. It runs on a separate NCCL stream so communication overlaps with the backward pass.

// Record an event on the compute stream when this gradient batch is ready
cudaEventRecord(config->compute_nccl_sync, compute_stream);
// Make the NCCL stream wait for that event (non-blocking on the host)
cudaStreamWaitEvent(config->nccl_stream, config->compute_nccl_sync);

ncclGroupStart();
for (int i = 0; i < N; i++) {
    if (zero_stage == 0) {
        ncclAllReduce(pointers[i], pointers[i], sizes[i], ncclFloatX, ncclAvg, comm, nccl_stream);
    } else if (zero_stage == 1) {
        shard_size   = sizes[i] / nproc;
        shard_offset = shard_size * process_rank;
        ncclReduceScatter(pointers[i], pointers[i] + shard_offset, shard_size, ncclFloatX, ncclAvg, comm, nccl_stream);
    }
}
ncclGroupEnd();

The NCCL group batches all collectives into a single kernel launch, minimizing launch overhead.

Using an event rather than cudaStreamSynchronize is what keeps the host free to enqueue the next layer’s backward work immediately.

Why ncclAvg? After the collective each GPU holds the mean gradient across processes, which is what an effective batch of B * N requires. The dloss = 1/(B*T*grad_accum_steps) you set in A3 already accounts for the rest.

---

Part 4 — Multi-GPU Backward and Update (train_gpt2.cu)

gpt2_backward_and_reduce — the TODO (A4) in the function you wrote for A3.

After each layer’s gradients are computed, call multi_gpu_async_reduce_gradient on that layer’s parameter gradient tensors. On the last micro-step only (last_step == true), also reduce the non-layer tensors (wte, wpe, lnfw, lnfb) and the mean loss via ncclAllReduce.

The overlap works because the backward pass runs from layer L-1 down to 0: by the time layer L-2’s backward starts, layer L-1’s reduction is already in flight. This is the same idea as PyTorch DDP’s bucketed all-reduce.

backward(layer L-1) → async_reduce_gradient(layer L-1 params)
backward(layer L-2) → async_reduce_gradient(layer L-2 params)
...  [compute and communicate overlap]
backward(layer 0)   → async_reduce_gradient(layer 0 params)
encoder_backward()  → async_reduce_gradient(wte, wpe, lnf)
ncclAllReduce(mean_loss)
cudaDeviceSynchronize()

gpt2_update — the TODO (A4) in the function you wrote for A3.

When zero_stage == 1, each GPU:

1. Calls the host-side adamw_update launcher only for its shard of each parameter tensor.
2. Calls ncclAllGather to broadcast the updated shard back to all GPUs.

The AllGather is issued per parameter tensor inside one NCCL group to amortize launch overhead:

ncclGroupStart();
for (int l = 0; l < num_layers; l++) {
    ncclAllGather(shard_ptr + l * tensor_size, full_ptr + l * tensor_size,
                  shard_size, ncclFloatX, comm, nccl_stream);
}
ncclGroupEnd();

---

Part 5 — Multi-GPU CPU Sum

multi_gpu_cpu_float_sum(value, config)

A lightweight collective for summing a single float on the host — used for aggregating validation losses and HellaSwag accuracy across ranks. Write the value into config->unified_buffer, ncclAllReduce 1 element with ncclSum, cudaDeviceSynchronize, and read it back. Return value unchanged when num_processes == 1.

---

Testing

Build

make train_gpt2cu

The Makefile prints whether it found OpenMPI and NCCL. If it says ✗ OpenMPI is not found, disabling multi-GPU support then MULTI_GPU is not defined and none of your code is even compiled — fix that first.

Single-node multi-GPU

There is no run.sh; the binary is launched with mpirun:

# 2-GPU DDP (ZeRO-0)
mpirun -np 2 ./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
                            -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin \
                            -b 4 -t 1024 -x 10 -z 0

# 2-GPU ZeRO-1
mpirun -np 2 ./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
                            -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin \
                            -b 4 -t 1024 -x 10 -z 1

Relevant flags (all single-character — --zero-stage will not work):

Flag	Meaning
-z	ZeRO stage (0 or 1)
-d	total batch size in tokens — hold this fixed when comparing GPU counts
-pi	NCCL init method: "mpi" (default), "tcp", "fs"
-pn / -pr / -pg	num processes / process rank / GPUs per node (TCP and FS modes)
-ps	server IP (TCP mode)
-pf	shared filesystem path (FS mode)

For the TCP and FS paths, build with NO_USE_MPI=1 and launch one process per rank. scripts/multi_node/run_gpt2_124M_tcp.sbatch and run_gpt2_124M_fs.sbatch are working examples.

Correctness check

Hold the total batch size -d fixed and vary the GPU count. The step-10 loss must match to 6 decimal places between:

              ./train_gpt2cu -d 32768 -b 4 -t 1024 -x 10   # 1 GPU
mpirun -np 2  ./train_gpt2cu -d 32768 -b 4 -t 1024 -x 10   # 2 GPUs
mpirun -np 4  ./train_gpt2cu -d 32768 -b 4 -t 1024 -x 10   # 4 GPUs

Any divergence indicates a bug in gradient averaging, optimizer state initialization, or a non-deterministic kernel from A3.

Bandwidth check

With 2× A100s, effective gradient bandwidth should reach ≥ 90% of NVLink bisection bandwidth. Compare against nccl-tests:

./build/all_reduce_perf -b 512M -e 512M -f 2 -g 2

---

Common Mistakes

• Not waiting on the compute stream before launching NCCL. Without cudaEventRecord + cudaStreamWaitEvent, the collective can start before the gradients exist. The result is a plausible-looking loss curve that is subtly wrong.
• Dividing by num_processes twice. ncclAvg already divides by N. Dividing dloss by N as well double-counts.
• Forgetting the final cudaDeviceSynchronize. The compute and NCCL streams run independently; you must sync before reading model->mean_loss on the host.
• Comparing runs at fixed -b instead of fixed -d. -b is the per-GPU micro-batch, so holding it fixed changes the total batch size with the GPU count and the losses will legitimately differ. Fix -d.
• ZeRO-1 offset arithmetic. For a per-layer parameter you must offset into the right layer first and then into the shard. Use gpt2_get_tensor_at_layer together with multi_gpu_get_shard_offset.
• Non-determinism from A3. A stray atomicAdd in a gradient-accumulating kernel will not show up on one GPU but breaks reproducibility across GPU counts. If losses diverge only at 4 GPUs, look there before looking at your NCCL code.

---

Extension (Bonus, 10%)

Implement get_nccl_id_via_tcp for multi-node initialization without MPI, and test it across two machines (-pi tcp). Expected: initialization adds < 1% to per-step wall time. scripts/multi_node/run_gpt2_124M_tcp.sbatch shows the launch pattern.

---

Checking your work

./grade A4

Runs the 1-GPU / 2-GPU / 4-GPU comparison at a fixed total batch size (-d) and checks the step-10 losses agree, for both ZeRO-0 and ZeRO-1. If fewer than two GPUs are visible, or the binary was built without NCCL, every line is reported as SKIP rather than failed.

---

Scoring

What ./grade A4 checks, and what each part is worth.

Component	Points
multi_gpu_config_init via the MPI path	15
get_nccl_id_via_fs and the FS init path	10
multi_gpu_get_local_device_idx	5
multi_gpu_config_free, multi_gpu_barrier, multi_gpu_cpu_float_sum	5
multi_gpu_get_shard_offset and set_zero_configs	10
multi_gpu_async_reduce_gradient — ZeRO-0 (AllReduce)	15
multi_gpu_async_reduce_gradient — ZeRO-1 (ReduceScatter)	15
gpt2_update with AllGather for ZeRO-1	15
Reproducible loss across 1/2/4 GPU runs at fixed total batch size	10
Total	100
Bonus: get_nccl_id_via_tcp multi-node init	+10