Communication Optimization › Code Implementation
Language AI HandbookAll chapters Communication Optimization › Code ImplementationOne section · Sections in this chapter6 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us build a concrete simulation that demonstrates the quantitative impact of these techniques on training throughput. We will not run actual distributed training (which requires multiple GPUs), but we will simulate the timing characteristics of the synchronization overhead and the improvement from each optimization strategy.
Setup and Baseline Simulation
First, we establish the simulation parameters and baseline (unoptimized) communication behavior.
# Simulation parameters
n_gpus = 64 # Total GPUs in the cluster
n_nodes = 8 # 8 nodes, 8 GPUs each
n_gpus_per_node = n_gpus // n_nodes
model_params = 7e9 # 7B parameter model
bytes_per_param = 4 # float32
# Hardware characteristics
nvlink_bw_gbps = 600 # A100-class NVLink 3.0 aggregate bidirectional bandwidth
infiniband_bw_gbps = 25 # InfiniBand HDR per GPU (GB/s)
pcie_bw_gbps = 32 # PCIe 4.0 x16 (GB/s)
# Compute timings per step
t_fwd_ms = 2000 # Forward pass (ms)
t_bwd_ms = 4000 # Backward pass (ms)
# Ring all-reduce sends ~2*(N-1)/N * total_bytes per GPU
total_grad_bytes = model_params * bytes_per_param
ring_factor = 2 * (n_gpus - 1) / n_gpus # Asymptotically 2.0
# Baseline: flat ring all-reduce over InfiniBand only
t_allreduce_baseline_ms = (
(ring_factor * total_grad_bytes / 1e9) / infiniband_bw_gbps * 1000
)Model: 7.0B parameters Gradient size: 28.0 GB (float32) Ring-reduce data per GPU: 55.1 GB Baseline all-reduce time: 2205 ms (2.2 s) Baseline step time: 8205 ms Communication fraction: 26.9%
The output shows how communication can dominate training time. With a 7B-parameter model on 64 GPUs using a flat ring all-reduce, the gradient tensor is 28 GB in float32, and the ring-reduce algorithm needs to transfer approximately twice this amount per GPU. Even at the full InfiniBand bandwidth, this communication time often equals or exceeds the compute time.
Gradient Compression Impact
Now we simulate how top-K sparsification and quantization reduce communication volume.
# Compression configurations
compression_configs = {
"No compression (FP32)": {
"ratio": 1.0,
"overhead_ms": 0,
"description": "Full gradients, float32",
},
"Top-1% sparse (FP32)": {
"ratio": 0.01,
"overhead_ms": 20, # Sparsification CPU overhead
"description": "1% density, with error feedback",
},
"Top-0.1% sparse (FP32)": {
"ratio": 0.001,
"overhead_ms": 25,
"description": "0.1% density (DGC regime)",
},
"8-bit quantization": {
"ratio": 0.25, # 8-bit vs 32-bit
"overhead_ms": 10,
"description": "Per-layer quantization",
},
"1-bit SGD": {
"ratio": 1 / 32,
"overhead_ms": 15,
"description": "Sign + scalar, 32x compression",
},
"PowerSGD rank-4": {
"ratio": (n_gpus_per_node + n_gpus_per_node) * 4 / (n_gpus_per_node**2),
"overhead_ms": 30,
"description": "Low-rank approx, ~rank 4",
},
}
results_compression = {}
for name, cfg in compression_configs.items():
compressed_bytes = total_grad_bytes * cfg["ratio"]
t_comm_ms = (
(ring_factor * compressed_bytes / 1e9) / infiniband_bw_gbps * 1000
)
t_total_ms = t_fwd_ms + t_bwd_ms + t_comm_ms + cfg["overhead_ms"]
speedup = (t_fwd_ms + t_bwd_ms + t_allreduce_baseline_ms) / t_total_ms
results_compression[name] = {
"compressed_gb": compressed_bytes / 1e9,
"t_comm_ms": t_comm_ms,
"t_total_ms": t_total_ms,
"speedup": speedup,
}Method Comm (GB) Comm (s) Step (s) Speedup --------------------------------------------------------------------------- No compression (FP32) 28.00 2.21 8.21 1.00x Top-1% sparse (FP32) 0.28 0.02 6.04 1.36x Top-0.1% sparse (FP32) 0.03 0.00 6.03 1.36x 8-bit quantization 7.00 0.55 6.56 1.25x 1-bit SGD 0.88 0.07 6.08 1.35x PowerSGD rank-4 28.00 2.21 8.23 1.00x
These numbers show a clear pattern: aggressive compression like DGC's 0.1% density or 1-bit SGD can reduce communication time by 2-3 orders of magnitude, turning a communication-dominated workload into a compute-bound one. The key observation is that even the overhead from the sparsification or quantization operation itself (20-30 ms in our simulation) is negligible compared to the time saved on communication.
Communication Overlap Impact
Next, we model the effect of bucketing and overlap on effective step time.
# Overlap model: gradients computed in reverse layer order during backward pass
# With bucketing, each bucket's communication can overlap with earlier layers' backprop
# Assume backward pass is roughly uniform across layers
n_buckets = 20 # Number of parameter buckets (DDP default: ~25 MB buckets)
# Communication time per bucket (flat ring, no compression)
t_comm_per_bucket_ms = t_allreduce_baseline_ms / n_buckets
# Backward time per bucket's layers
t_bwd_per_bucket_ms = t_bwd_ms / n_buckets
# With overlap: bucket i's communication runs during the backward computation of buckets 0..i-1
# The last bucket has no overlap opportunity; the first n-1 buckets can overlap
overlap_configs = {
"No overlap (sequential)": {"overlap_fraction": 0.0},
"Partial overlap (50%)": {"overlap_fraction": 0.5},
"Full bucketing overlap": {"overlap_fraction": (n_buckets - 1) / n_buckets},
}
results_overlap = {}
for name, cfg in overlap_configs.items():
# Communication that cannot be overlapped
t_comm_exposed_ms = t_allreduce_baseline_ms * (1 - cfg["overlap_fraction"])
t_total_ms = t_fwd_ms + t_bwd_ms + t_comm_exposed_ms
speedup_vs_sequential = (
t_fwd_ms + t_bwd_ms + t_allreduce_baseline_ms
) / t_total_ms
results_overlap[name] = {
"overlap_pct": cfg["overlap_fraction"] * 100,
"exposed_comm_ms": t_comm_exposed_ms,
"t_total_ms": t_total_ms,
"speedup": speedup_vs_sequential,
}Overlap Strategy Overlap % Exposed comm (s) Step time (s) Speedup ------------------------------------------------------------------------------------- No overlap (sequential) 0 2.21 8.21 1.00x Partial overlap (50%) 50 1.10 7.10 1.16x Full bucketing overlap 95 0.11 6.11 1.34x
Full bucketing overlap hides nearly all communication behind backward computation, approaching the ideal case where communication adds no wall-clock overhead. The remaining exposed communication (1/20th of the total, from the last bucket) is unavoidable because there is no remaining backward computation to overlap it with.
Hierarchical All-Reduce Benefit
We now model the bandwidth benefit of the two-level hierarchical all-reduce.
# In hierarchical all-reduce:
# Phase 1: intra-node reduce-scatter via NVLink (full NVLink bandwidth)
# Phase 2: inter-node all-reduce via InfiniBand (only P/n_nodes bytes per node)
# Phase 3: intra-node all-gather via NVLink
# Phase 1: each GPU sends (n_gpus_per_node - 1) / n_gpus_per_node * (P/n_nodes) bytes via NVLink
# Phase 2: each node sends P/n_nodes bytes via InfiniBand (ring over n_nodes nodes)
# Phase 3: symmetric to phase 1
# Intra-node reduce-scatter bytes per GPU
intra_bytes_per_gpu = total_grad_bytes / n_gpus * (n_gpus_per_node - 1)
t_phase1_ms = (intra_bytes_per_gpu / 1e9) / nvlink_bw_gbps * 1000
# Inter-node all-reduce bytes per node (only 1/n_nodes of params)
inter_bytes_per_node = total_grad_bytes / n_nodes
inter_ring_factor = 2 * (n_nodes - 1) / n_nodes
t_phase2_ms = (
(inter_ring_factor * inter_bytes_per_node / 1e9) / infiniband_bw_gbps * 1000
)
# Intra-node all-gather (same as phase 1)
t_phase3_ms = t_phase1_ms
t_hierarchical_ms = t_phase1_ms + t_phase2_ms + t_phase3_ms
speedup_hier = t_allreduce_baseline_ms / t_hierarchical_ms
allreduce_configs = {
"Flat ring (InfiniBand)": t_allreduce_baseline_ms,
"Hierarchical (NVLink + InfiniBand)": t_hierarchical_ms,
}All-reduce time breakdown (hierarchical): Phase 1 intra-node reduce-scatter (NVLink): 5.1 ms Phase 2 inter-node all-reduce (InfiniBand): 245.0 ms Phase 3 intra-node all-gather (NVLink): 5.1 ms Total hierarchical: 255.2 ms Flat ring all-reduce: 2205.0 ms Hierarchical all-reduce: 255.2 ms Speedup from hierarchy: 8.6x
The hierarchical approach is dramatically faster because it routes the vast majority of data movement over NVLink (which is approximately 24x faster than InfiniBand per unit bandwidth) rather than InfiniBand. The flat ring all-reduce unnecessarily sends all 28 GB of gradients over InfiniBand, including the portions that only need to travel between GPUs on the same node. The hierarchical approach sends only 1/8th of the data over InfiniBand, with the remaining intra-node phases completing nearly instantly over NVLink.
Combined Optimization Impact
We compose all three optimizations to estimate the total impact.
# Combined: 1% sparsification + hierarchical + overlap
compression_ratio = 0.01
overhead_compression_ms = 20
# Hierarchical time with compressed gradients
compressed_grad_bytes = total_grad_bytes * compression_ratio
intra_bytes_compressed = compressed_grad_bytes / n_gpus * (n_gpus_per_node - 1)
t_phase1_comp_ms = (intra_bytes_compressed / 1e9) / nvlink_bw_gbps * 1000
inter_bytes_comp_per_node = compressed_grad_bytes / n_nodes
t_phase2_comp_ms = (
(inter_ring_factor * inter_bytes_comp_per_node / 1e9)
/ infiniband_bw_gbps
* 1000
)
t_hier_comp_ms = t_phase1_comp_ms + t_phase2_comp_ms + t_phase1_comp_ms
# With full overlap (expose only 1/n_buckets of hierarchical comm time)
t_exposed_comm_ms = t_hier_comp_ms / n_buckets
t_combined_ms = (
t_fwd_ms + t_bwd_ms + t_exposed_comm_ms + overhead_compression_ms
)
t_baseline_total_ms = t_fwd_ms + t_bwd_ms + t_allreduce_baseline_ms
total_speedup = t_baseline_total_ms / t_combined_ms
comm_reduction = t_allreduce_baseline_ms / (
t_exposed_comm_ms + overhead_compression_ms
)
scenarios = {
"Baseline (flat ring, no overlap)": t_baseline_total_ms,
"+ Hierarchical topology": t_fwd_ms + t_bwd_ms + t_hierarchical_ms,
"+ Compression (1% sparse)": t_fwd_ms
+ t_bwd_ms
+ t_hier_comp_ms
+ overhead_compression_ms,
"+ Overlap (full bucketing)": t_combined_ms,
}Scenario Step time (s) Speedup ----------------------------------------------------------------- Baseline (flat ring, no overlap) 8.21 1.00x + Hierarchical topology 6.26 1.31x + Compression (1% sparse) 6.02 1.36x + Overlap (full bucketing) 6.02 1.36x Communication overhead reduced by 110x Overall training throughput improved by 1.4x
Stacking all three optimizations reveals the multiplicative effect: each technique addresses a different aspect of the communication bottleneck, and together they can reduce the communication-to-computation ratio to where it is no longer a significant training constraint.
