NCCL Primitives Deep Dive
When you’re training a 70B parameter LLM across 128 GPUs, every single GPU needs to coordinate with others. Whether it’s gathering weights, summing gradients, or routing tokens to experts in an MoE architecture, efficient communication is the bottleneck that determines if your training is rapid or not.
In this post, we’ll dive deep into 3 critical NCCL primitives that power modern distributed training:
- All-Gather → Weight gathering across GPUs (tensor parallelism)
- Ring All-Reduce → Gradient aggregation (data parallelism)
- All-to-All → MoE token/output passing (expert parallelism)
We’ll analyze the bandwidth requirements, walk through step-by-step examples, and show pseudo code to build intuition.
Why These Primitives Matter
Modern LLM training uses different parallelism strategies:
- Tensor Parallelism: Shard model weights across GPUs → needs All-Gather
- Data Parallelism: Each GPU has full model, processes different batches → needs All-Reduce for gradients
- Expert Parallelism (MoE): Experts on different GPUs, tokens routed dynamically → needs All-to-All
Each strategy has different communication patterns, and using the wrong primitive is a recipe for disaster.
1. All-Gather: Weight Gathering Across GPUs
The Use Case
In tensor parallelism or FSDP Zero 3 variant, model weights are sharded across multiple GPUs. For example, a large linear layer might be split so each GPU holds 1/N of the weight matrix. During the forward pass, each GPU needs the complete weight matrix to compute its output.
The Problem: How do we efficiently gather all shards on every GPU?
Naive Approach: The Bottleneck
The naive solution: Have each GPU send its shard to every other GPU.
# Pseudo-code for naive all-gather
for each GPU i:
for each other GPU j (j != i):
send shard_i to GPU_j
Communication per GPU: Each GPU sends M/N data to N-1 other GPUs = M(N-1)/N data sent.
The Bottleneck: What if all GPUs send simultaneously → network congestion! Also, each GPU must handle N-1 separate send operations.
Ring All-Gather: The Elegant Solution
The ring pattern in distributed systems is an architecture where nodes are organized in a logical or physical circle, with each node connecting to exactly two neighbors.
In a ring topology, GPUs form a logical ring: GPU_0 → GPU_1 → GPU_2 → ... → GPU_{N-1} → GPU_0
Algorithm (for N GPUs, each with shard of size M/N):
Step 0: Each GPU has its own shard
Step 1: Pass shard to right neighbor (receive from left)
Step 2: Pass newly received shard to right neighbor
...
Step N-1: Complete after N-1 steps
After N-1 steps, every GPU has accumulated all N shards.
Communication Analysis:
- Each step: Send
M/Ndata - Total steps:
N-1 - Every GPU: Data sent and received per step is
M/N - Total data sent per GPU:
M(N-1)/N - Bandwidth utilization: Every link used simultaneously, no bottleneck!
Step-by-Step Example (4 GPUs)
Initial state:
GPU_0: [A, _, _, _]
GPU_1: [_, B, _, _]
GPU_2: [_, _, C, _]
GPU_3: [_, _, _, D]
Step 1: GPU_i sends to GPU_{(i+1)%4}
GPU_0: [A, _, _, D]
GPU_1: [A, B, _, _]
GPU_2: [_, B, C, _]
GPU_3: [_, _, C, D]
Step 2: GPU_i sends the new shard again to GPU_{(i+1)%4}
GPU_0: [A, _, C, D]
GPU_1: [A, B, _, D]
GPU_2: [A, B, C, _]
GPU_3: [_, B, C, D]
Step 3: Repeat
GPU_0: [A, B, C, D]
GPU_1: [A, B, C, D]
GPU_2: [A, B, C, D]
GPU_3: [A, B, C, D]
Compared to the naive approach, the total data sent per GPU remains the same M(N-1)/N
Pseudo Code
import numpy as np
def ring_all_gather(shard, rank, num_gpu):
"""
Ring All-Gather pseud code, for each GPU rank
Args:
shard: Local shard of size (M/N,) on this GPU
rank: Current GPU rank (0 to world_size-1)
num_gpu: Total number of GPUs (N)
Returns:
gathered: Complete data of size (M,) with all shards
"""
# Initialize buffer to hold all shards (M/N)
shard_size = len(shard) #[4x1]
gathered = np.zeros(shard_size * num_gpu) #[16x1]
# Place local shard in correct position
gathered[rank * shard_size:(rank + 1) * shard_size] = shard
# Ring all-gather: N-1 steps
for step in range(num_gpu - 1):
# Which shard to send in this step?
send_idx = (rank - step) % num_gpu
# Simulate send/receive with neighbors
send_shard = gathered[send_idx * shard_size : (send_idx+1) * shard_size]
# A conceptual function; Send to the right neighbour, receive from left neighbour
received_shard = receive_from_left_neighbour(send_shard)
# Place received shard
recv_idx = (rank - step - 1) % num_gpu
gathered[recv_idx * shard_size:(recv_idx + 1) * shard_size] = received_shard
return gathered
# Example usage
if __name__ == "__main__":
# Simulate 4 GPUs, each with a shard
num_gpu = 4
data_size = 16
shard_size = data_size // num_gpu
# Each GPU has different shard
shards = [
np.array([1, 1, 1, 1]), # GPU 0 [4x1]
np.array([2, 2, 2, 2]), # GPU 1 [4x1]
np.array([3, 3, 3, 3]), # GPU 2 [4x1]
np.array([4, 4, 4, 4]), # GPU 3 [4x1]
]
# Each GPU runs all-gather, calls it async.
for rank in range(num_gpu):
result = ring_all_gather(shards[rank], rank, num_gpu)
print(f"GPU {rank} result: {result}")
Why Ring All-Gather Outperforms Naive
Both approaches send/receive the same total data: M(N-1)/N per GPU.
The difference is communication pattern:
- Naive: Each GPU → (N-1) GPUs simultaneously
- Network interface contention
- May not align with physical topology
- Links underutilized or congested
- Ring: Each GPU → 1 neighbor at a time
- Fully pipelined, all links active
- Maps to physical NVLink topology
- Optimal bandwidth utilization
2. Ring All-Reduce: Gradient Aggregation
The Use Case
In data parallelism (or FSDP Zero1/2/3), each GPU holds a complete copy of the model (or model shards) and processes a different batch of data. After the backward pass, each GPU computes gradients for all parameters. To update the model, we need the sum of all gradients on every GPU.
The Problem: How do we efficiently sum N gradient tensors of size M and distribute the result to all GPUs?
The Parameter Server Bottleneck
At first glance, it seems straightforward: just send everything to a central server, sum it up, and send it back. But if you try that with a 70B parameter model, your central server immediately becomes the bottleneck!
The Math of the Bottleneck:
In a naive setting with a central parameter server:
- Server receives:
M × Ndata (gradients from N GPUs) - Server sends:
M × Ndata (summed gradients to N GPUs)
As you add more GPUs, the server’s bandwidth requirement scales linearly. The system chokes!
This is where Ring All-Reduce comes into play.
Ring All-Reduce: The Two-Phase Solution
Ring All-Reduce combines two primitives:
Phase 1: Reduce-Scatter
- Break model gradients into
Nchunks - After
N-1steps, each GPU holds the complete sum for exactly1/Nof the parameter gradients
Phase 2: All-Gather
- Pass these completed shard sums around the ring
- After
N-1steps, each GPU has the full summed gradient
Total Communication: 2M(N-1)/N ≈ 2M per GPU (almost independent of N!)
Step-by-Step Reduce-Scatter
Ring All reduce uses a reduce scatter primitive which does exactly as it sounds -> Each GPU reduces its corresponding shard and sends to the neighbour. After N-1 rounds, each GPU has the completed sum of its shard.
Initial state (4 GPUs, gradients split into 4 chunks):
GPU_0: [A0, B0, C0, D0]
GPU_1: [A1, B1, C1, D1]
GPU_2: [A2, B2, C2, D2]
GPU_3: [A3, B3, C3, D3]
Step 1: GPU_i sends its shard to GPU_{(i+1)%4}, and GPU_{(i+1)%4} accumulates the received shard.
GPU_0: [A0, B0, C0, D0+D3]
GPU_1: [A0+A1, B1, C1, D1]
GPU_2: [A2, B1+B2, C2, D2]
GPU_3: [A3, B3, C2+C3, D3]
Step 2: Repeats
GPU_0: [A0, B0, C0+C2+C3, D0+D3]
GPU_1: [A0+A1, B1, C1, D1+D0+D3]
GPU_2: [A2+A0+A1, B1+B2, C2, D2]
GPU_3: [A3, B3+B1+B2, C2+C3, D3]
Step 3: Repeats
GPU_0: [A0, B0+B3+B1+B2, C0+C2+C3, D0+D3]
GPU_1: [A0+A1, B1, C1+C0+C2+C3, D1+D0+D3]
GPU_2: [A2+A0+A1, B1+B2, C2, D2+D1+D0+D3]
GPU_3: [A3+A2+A0+A1, B3+B1+B2, C2+C3, D3]
Final: After reduce-scatter (Note: The other partials can be ignored / will be overwritten by the all-gather in the next step)
GPU_0: [_, B0+B1+B2+B3, _, _]
GPU_1: [_, _, C1+C0+C2+C3, _]
GPU_2: [_, _, _, D0+D1+D2+D3]
GPU_3: [A0+A1+A2+A3, _, _, _]
Step-by-Step All-Gather
Refer to the previous section by using a ring all-gather approach.
After all-gather, every GPU has:
[A0+A1+A2+A3, B0+B1+B2+B3, C0+C1+C2+C3, D0+D1+D2+D3]
Note: A keen eyed observer would realize, during the reduce scatter, the initial ownership of the shard has been shifted. For example,
- Initial ownership: GPU 0 starts with A0, GPU 1 with B1, etc.
- Final ownership (after reduce-scatter): Chunks rotate, so GPU 0 ends up with sum(B), not sum(A)
Pseudo Code
import numpy as np
def reduce_scatter(chunks, rank, num_gpu):
"""
Reduce-Scatter: Each GPU gets sum of one chunk
Args:
chunks: List of N chunks, each of size M/N
rank: Current GPU rank
num_gpu: Total number of GPUs
Returns:
buffer: returns the full reduced buffer
"""
chunk_size = len(chunks[0]) #[4]
buffer = np.concatenate(chunks) #[16x1]
for step in range(num_gpu-1):
step_idx = (rank - step) % num_gpu
send_shard = buffer[step_idx*chunk_size: (step_idx+1)*chunk_size]
# Simulated function to reduce scatter: Send to right, Receive from left.
received_shard_sum = reduce_scatter_from_left(send_shard)
received_idx = (rank - step - 1) % num_gpu
buffer[received_idx*chunk_size : (received_idx+1)*chunk_size] += received_shard_sum
return buffer #[16x1]
def ring_all_reduce(gradient, rank, num_gpu):
"""
Ring All-Reduce: Sum gradients across all GPUs
Args:
gradient: Local gradient tensor of size M
rank: Current GPU rank
num_gpu: Total number of GPUs
Returns:
summed_gradient: Sum of all gradients, size M
"""
# Step 1: Split gradient into chunks
chunk_size = len(gradient) // num_gpu #[16/4 = 4]
chunks = [gradient[i*chunk_size:(i+1)*chunk_size] for i in range(num_gpu)]
# Step 2: Reduce-scatter phase
reduced_chunks = reduce_scatter(chunks, rank, num_gpu)
# Step 3: All-gather phase
summed_gradient = ring_all_gather(reduced_chunks, rank, num_gpu)
return summed_gradient
# Example usage
if __name__ == "__main__":
# Simulate 4 GPUs, each with gradients
num_gpu = 4
gradient_size = 16
gradients = [
np.ones(gradient_size) * 1, # GPU 0 [16x1]
np.ones(gradient_size) * 2, # GPU 1 [16x1]
np.ones(gradient_size) * 3, # GPU 2 [16x1]
np.ones(gradient_size) * 4, # GPU 3 [16x1]
]
expected_sum = np.ones(gradient_size) * 10 # Sum = 1+2+3+4
for rank in range(num_gpu):
result = ring_all_reduce(gradients[rank], rank, num_gpu)
# print(f"GPU {rank} summed gradient: {result}")
3. All-to-All: MoE Token Routing
The Use Case
In Mixture of Experts (MoE) architectures, you have E experts distributed across N GPUs. During the forward pass:
- A routing network decides which expert should process each token
- Tokens need to be grouped to the GPU hosting their assigned expert
- After expert computation, outputs need to be grouped back
The Problem: Each GPU needs to send different amounts of data to each other GPU (dynamic routing). This is a gpu specific exchange.
Why All-Gather Won’t Work
All-Gather assumes every GPU contributes the same amount of data. In MoE, routing is dynamic:
- GPU 0 might send 10 tokens to GPU 1, 5 to GPU 2, 0 to GPU 3
- GPU 1 might send 3 tokens to GPU 0, 8 to GPU 2, 4 to GPU 3
We need a primitive that supports arbitrary dynamic exchange.
All-to-All: Personalized Communication
In an All-to-All primitive:
- Each GPU has a separate information for every other GPU
- All GPUs simultaneously exchange their specific messages
Algorithm:
for each GPU i:
for each GPU j:
send message_i_to_j to GPU_j
receive message_j_to_i from GPU_j
Communication: In the worst case, if each GPU sends M/N tokens to each other GPU, total sent = M(N-1)/N per GPU.
MoE Specifics:
- Forward pass: All-to-All to route tokens to experts
- Backward pass: All-to-All to route gradients back to original GPUs
Basic View All-to-All exchanges data between all GPUs in a group.
Think of data as a matrix where rows = source GPU, columns = destination GPU
Before All-to-All:
Dest: GPU0 GPU1 GPU2 GPU3
Source GPU0: A0 A1 A2 A3
GPU1: B0 B1 B2 B3
GPU2: C0 C1 C2 C3
GPU3: D0 D1 D2 D3
After All-to-All (transpose):
Dest: GPU0 GPU1 GPU2 GPU3
Source GPU0: A0 B0 C0 D0
GPU1: A1 B1 C1 D1
GPU2: A2 B2 C2 D2
GPU3: A3 B3 C3 D3
Step-by-Step Example (4 GPUs, 8 Experts)
Scenario: 16 tokens, 4 GPUs, 8 experts (2 experts per GPU)
Data Parallelism + Expert Parallelism
GPU 0: Experts [E0, E1]
GPU 1: Experts [E2, E3]
GPU 2: Experts [E4, E5]
GPU 3: Experts [E6, E7]
Tokens on GPU 0: [T0, T1, T2, T3]
Routing decisions:
T0 → E2 (GPU 1)
T1 → E5 (GPU 2)
T2 → E0 (GPU 0, stays local)
T3 → E7 (GPU 3)
All-to-All exchange:
GPU 0 sends: [T0] to GPU 1, [T1] to GPU 2, [T3] to GPU 3
GPU 0 receives: tokens routed to E0 and E1 from other GPUs
All-to-All redistributes tokens from different sequences to the GPUs hosting their assigned experts. This makes sense because:
- GPU 0 has sequence 0’s tokens but needs to send them to other experts
- GPU 0 receives tokens from OTHER sequences that need its local experts E0, E1
Expert Parallelism Flow
Forward Pass:
Step 1: Routing (local on each GPU)
# Each GPU routes its tokens
router_logits = router(tokens) # Local
top_k_experts = topk(router_logits, k=2)
Step 2: Dispatch (All-to-All)
# Group tokens by destination expert
# All-to-All sends tokens to expert's GPU
dispatched_tokens = all_to_all(grouped_tokens)
Step 3: Expert Computation (local)
# Each GPU processes its experts' tokens
for expert in my_experts:
expert_output = expert(dispatched_tokens[expert_id])
Step 4: Combine (All-to-All)
# Send results back to originating GPUs
combined_output = all_to_all(expert_outputs)
Communication cost: 2 All-to-Alls per MoE layer
Simplified Pseudo Code
def all_to_all_moe(tokens, routing, rank, num_gpus, num_experts):
"""
All-to-All for MoE token routing
Args:
tokens: List of token IDs on this GPU
routing: Expert assignment for each token
rank: Current GPU rank
num_gpus: Total number of GPUs
num_experts: Total number of experts
Returns:
Tokens this GPU will process after All-to-All
"""
experts_per_gpu = num_experts // num_gpus
# Group tokens by destination GPU based on expert assignment
send_to_gpu = {i: [] for i in range(num_gpus)}
for token_id, expert_id in zip(tokens, routing):
target_gpu = expert_id // experts_per_gpu
send_to_gpu[target_gpu].append((token_id, expert_id))
# All to ALL will happen here to send and receive the tokens
# For this case, return what this GPU keeps locally + tokens received from all to all for this gpu
return send_to_gpu[rank]
# Example: GPU 0 in MoE with Data Parallelism
if __name__ == "__main__":
num_gpus = 4
num_experts = 8
# GPU 0 has tokens from Sequence 0
rank = 0
tokens = [0, 1, 2, 3] # Token IDs: T0, T1, T2, T3
routing = [2, 5, 0, 7] # T0→E2, T1→E5, T2→E0, T3→E7
# Expert distribution across GPUs:
# GPU 0: E0, E1
# GPU 1: E2, E3
# GPU 2: E4, E5
# GPU 3: E6, E7
local_tokens = all_to_all_moe(tokens, routing, rank, num_gpus, num_experts)
# GPU 0 keeps: Token 2 for Expert 0
# GPU 0 receives: Tokens from other GPUs routed to E0, E1
Bandwidth Analysis & Comparison
| Primitive | Use Case | Data Sent Per GPU | Bandwidth Scaling |
|---|---|---|---|
| All-Gather | Weight gathering (TP) | M(N-1)/N | (~M) |
| All-Reduce | Gradient aggregation (DP) | 2M(N-1)/N | (~2M) |
| All-to-All | Token routing (MoE) | Variable, up to M(N-1)/N | Depends on routing |
Key Insight: All three primitives scale near-constant with the number of GPUs!
All-Gather as a Special Case of All-to-All
All-Gather is actually a special case of All-to-All where each GPU sends the same data to everyone.
Example with 4 GPUs:
Initial state:
GPU 0: [1]
GPU 1: [2]
GPU 2: [3]
GPU 3: [4]
All-Gather behavior:
GPU 0 sends [1] to GPU 1, 2, 3
GPU 1 sends [2] to GPU 0, 2, 3
GPU 2 sends [3] to GPU 0, 1, 3
GPU 3 sends [4] to GPU 0, 1, 2
Result (all GPUs):
[1, 2, 3, 4]
Same result using All-to-All:
GPU 0 All-to-All send:
To GPU 0: [1] (keep local)
To GPU 1: [1] (same shard)
To GPU 2: [1] (same shard)
To GPU 3: [1] (same shard)
GPU 1 All-to-All send:
To GPU 0: [2]
To GPU 1: [2] (keep local)
To GPU 2: [2]
To GPU 3: [2]
... (GPU 2, 3 similarly)
Result (all GPUs):
[1, 2, 3, 4] ✓ Same as All-Gather!
The Difference:
- All-Gather: Everyone sends their shard to everyone (broadcast pattern)
- All-to-All (MoE): Each GPU sends different, selective data based on routing
- GPU 0 might send [T0] to GPU 1, [T1] to GPU 2, [T2] to GPU 3
- More efficient for sparse, dynamic patterns!
Conclusion
The Big Picture:
- All-Gather: Collect sharded data without reduction
- All-Reduce: Aggregate (sum) and distribute
- All-to-All: Personalized, dynamic exchange
All three use clever topologies (rings) to avoid bottlenecks and utilize full network bandwidth.
References
-
HuggingFace Ultra-Scale Playbook
https://huggingface.co/spaces/weege007/ultrascale-playbook -
Machine Learning Systems for Dummies - Abhishek Maiti
https://abhishekmaiti.com/mlsys-for-dummies/ -
NCCL Overview - NVIDIA
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/overview.html
This post is part of my ongoing deep dive into systems-level ML engineering on GPUs. Some diagrams and examples were developed with AI assistance.