LLM Lab Notes 01: Distributed Training 101 — From init_process_group to all_reduce

Series note: This series is based on experiments actually run by the author on GPU clusters (llm_lab/), breaking down the core mechanisms of LLM training infrastructure line by line. Note 01 starts at the most unassuming — yet decisive — starting point of distributed training: process group initialization.

Code files: llm_lab/01_ddp/ddp_test.py, ddp_test2.py

1. Experiment Goal

The first question anyone asks when touching multi-GPU training is usually: how do programs on multiple GPUs "know" each other?

This experiment has three goals:

1. Get NCCL process group initialization working with the shortest possible code;
2. Understand what rank and local_rank really are, and why set_device(local_rank) is mandatory;
3. Verify with all_reduce that "collective communication = multi-GPU data synchronization."

2. Code Walkthrough

Full code (ddp_test.py, 18 lines):

import os
import torch
import torch.distributed as dist

dist.init_process_group("nccl")          # ① init process group

rank = dist.get_rank()                   # ② global id
local_rank = int(os.environ["LOCAL_RANK"])   # ③ node-local id

torch.cuda.set_device(local_rank)        # ④ bind current GPU

x = torch.ones(1).cuda() * rank          # ⑤ each GPU holds a distinct value

dist.all_reduce(x)                       # ⑥ cross-GPU sum

print(f"Rank {rank}: {x.item()}")        # ⑦ result

dist.destroy_process_group()             # ⑧ cleanup

① dist.init_process_group("nccl")

This single line is the foundation of every multi-GPU PyTorch program. It does three things:

• Handshakes with all participating processes and establishes the process group;
• Selects the communication backend — nccl (NVIDIA Collective Communications Library) is the standard backend for multi-GPU training;
• Reads MASTER_ADDR / MASTER_PORT / WORLD_SIZE / RANK from the environment and completes global topology alignment.

A common confusion: "why can I run without passing any arguments?" — because the arguments come from environment variables. When launched with torchrun, they are injected automatically:

torchrun --nproc_per_node=4 ddp_test.py

②③ rank vs local_rank

• rank: the process's global ID, unique across the whole training job (0, 1, 2, 3, …), numbered consecutively across nodes. It decides "who you are."
• local_rank: the process's ID within the current node. It restarts from 0 on every node, and is also the index of the physical GPU you should use.

⚠️ Key distinction: rank is your identity in the communication world; local_rank is your identity in the hardware world. In multi-node setups, the process with rank=8 and the one with rank=0 both run on local_rank=0 of their own nodes.

④ torch.cuda.set_device(local_rank)

Why is this required? Because PyTorch defaults to cuda:0. Without explicit switching, all processes on an 8-GPU node would pile onto GPU 0 — initialization passes, and memory blows up immediately.

⑤⑥⑦ ones(1) * rank and all_reduce

Each GPU constructs a tensor containing only its own rank: GPU 0 gets [0], GPU 1 gets [1], GPU 2 gets [2], GPU 3 gets [3].

dist.all_reduce(x) is the most basic collective primitive: every process contributes its x, a full reduction is performed, and the result is broadcast to everyone. The default reduction operator is SUM.

3. How It Works: Collectives and the "Starting Line"

The semantics of all_reduce can be summarized in one sentence:

Every process feeds in one piece of data, and every process receives the same, aggregated result.

This is the mathematical foundation of data-parallel gradient synchronization: each GPU computes a different gradient; all_reduce sums (or averages) them and broadcasts the result, so all GPUs hold identical gradients — which is exactly what DistributedDataParallel executes behind the scenes every training step.

Another important fact: during init_process_group, NCCL establishes a communication ring/tree topology and runs bandwidth probes, which determines how fast all subsequent collectives can run. Initialization itself is overhead — which is why launching tens of thousands of processes on a 10K-GPU cluster is a systems-engineering problem in its own right.

4. Results

With 4 GPUs, each feeding in its rank [0, 1, 2, 3], after all_reduce(SUM) every GPU should print:

Rank 0: 6.0
Rank 1: 6.0
Rank 2: 6.0
Rank 3: 6.0

(0 + 1 + 2 + 3 = 6; for N GPUs it is N(N−1)/2. This is a deterministic mathematical result of the collective, usable as a correctness check.)

How to tell the experiment succeeded: all four lines must show the same value. If any GPU differs, the process group is misaligned or a link is broken — and in a real cluster, that is the first scene of communication troubleshooting.

5. Engineering Takeaways

1. NCCL initialization is the starting line of 10K-GPU training: failed/timed-out init_process_group usually means a wrong MASTER_ADDR, a firewalled MASTER_PORT, or unreachable RDMA between nodes. When many jobs start at once, the handshake storm itself can overwhelm the scheduler — common practice on large clusters is staggered startup.

2. Misusing local_rank is the classic pitfall: forgetting set_device(local_rank) causes memory contention; using rank as the device ID goes out of bounds in multi-node setups. The canonical form is always torch.cuda.set_device(local_rank).

3. Run the 18-line minimal experiment before optimizing: in production troubleshooting, I start with nccl-tests and this minimal script to rule out "the process group is fundamentally broken," then move into bandwidth and topology analysis — otherwise you are tuning against nothing.

4. Don't skip destroy_process_group: leftover NCCL resources from repeated experiments can interfere with consecutive runs on the same node.

Next: with the process group up, wrap a real model in DDP — see LLM Lab Notes 02: DDP Minimal Training.