with contextlib.redirect_stdout(io.StringIO()), \
contextlib.redirect_stderr(io.StringIO()): # download once, quietly
datasets.FashionMNIST('./data', train=True, download=True)
DDP_SCRIPT = r'''
import json, os, sys, time, torch
from torch import nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision import datasets, transforms
from d2l import torch as d2l
rank = int(os.environ["LOCAL_RANK"]); world = int(os.environ["WORLD_SIZE"])
torch.cuda.set_device(rank); dist.init_process_group("nccl")
torch.set_float32_matmul_precision("high"); torch.backends.cudnn.benchmark = True
model = d2l.resnet18(10, 1).to(rank)
model(torch.zeros(1, 1, 64, 64, device=rank)) # materialize Lazy params
model = DDP(model, device_ids=[rank])
opt = torch.optim.SGD(model.parameters(), lr=0.1)
loss = nn.CrossEntropyLoss()
B = int(sys.argv[1]) if len(sys.argv) > 1 else 256 # per-rank batch size
tf = transforms.Compose([transforms.Resize(64), transforms.ToTensor()])
ds = datasets.FashionMNIST("./data", train=True, transform=tf)
sampler = torch.utils.data.distributed.DistributedSampler(ds, world, rank)
loader = torch.utils.data.DataLoader(ds, B, sampler=sampler, num_workers=2)
for epoch in range(2): # epoch 0 warms up
sampler.set_epoch(epoch); n = 0; torch.cuda.synchronize(); t0 = time.time()
for X, y in loader:
X, y = X.to(rank), y.to(rank)
opt.zero_grad(set_to_none=True)
loss(model(X), y).backward(); opt.step(); n += X.shape[0]
torch.cuda.synchronize(); dt = time.time() - t0
here = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(here, f"rank{rank}.json"), "w") as f:
json.dump({"rank": rank, "samples_per_s": n / dt}, f)
dist.destroy_process_group()
'''
DDP_DIR = pathlib.Path('ddp_scratch')
DDP_DIR.mkdir(exist_ok=True)
_ = (DDP_DIR / 'train_ddp.py').write_text(DDP_SCRIPT)Multi-GPU in Practice
What the Hand-Rolled Loop Lacked
Three deficits, all software:
- no overlap — communicate only after the whole backward
- one process — one GIL dispatching k GPUs
- star topology — device 0 as a hub
DDP fixes all three: process per GPU, ring collectives, and buckets that overlap comm with backward.
DDP Overlap
Gradients arrive back-to-front; bucket them and allreduce each as it fills, hiding communication under compute.
DDP, Really Run
Multiple processes from a notebook: write a sidecar script, launch with torchrun, read back per-rank results.
init_process_group + DDP(model) — the loop body is unchanged from single-GPU; the scaffolding is the launcher’s.
Weak Scaling, Measured
Convention first: per-rank batch fixed at 256 ⇒ global batch grows with k — weak scaling.
def ddp_throughput(k, batch_size=256):
"""Launch k-process DDP via torchrun; return aggregate samples/s."""
for stale in DDP_DIR.glob('rank*.json'): # a crashed run must not
stale.unlink() # serve old results
torchrun = str(pathlib.Path(sys.executable).parent / 'torchrun')
subprocess.run([torchrun, '--standalone', f'--nproc-per-node={k}',
str(DDP_DIR / 'train_ddp.py'), str(batch_size)],
check=True)
per_rank = []
for r in range(k):
with open(DDP_DIR / f'rank{r}.json') as f:
per_rank.append(json.load(f)['samples_per_s'])
return sum(per_rank)
ks = [k for k in (1, 2, 4) if k <= d2l.num_gpus()]
tput = [ddp_throughput(k) for k in ks]
for k, t in zip(ks, tput):
print(f'{k} GPU(s): {t:.0f} samples/s, {t / tput[0]:.2f}x, '
f'{100 * t / tput[0] / k:.0f}% weak-scaling efficiency')
d2l.plot(ks, [tput], 'GPUs', 'samples/s')~1.8× at 2 GPUs (88%), ~3.3× at 4 (82%) — sublinear, no cliff. Strong scaling (global batch 512, split thinner) parts company as k grows. NVLink changes only the constant.
Price the Fabric, Then Check the Bill
no_sync() turns gradient sync off ⇒ synced − unsynced steps estimate the communication time.
COMM_SCRIPT = r'''
import contextlib, json, os, time, torch
from torch import nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from d2l import torch as d2l
rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(rank)
dist.init_process_group("nccl", device_id=torch.device("cuda", rank))
torch.set_float32_matmul_precision("high"); torch.backends.cudnn.benchmark = True
model = d2l.resnet18(10, 1).to(rank)
model(torch.zeros(1, 1, 64, 64, device=rank))
model = DDP(model, device_ids=[rank])
loss = nn.CrossEntropyLoss()
X = torch.randn(256, 1, 64, 64, device=rank)
y = torch.randint(0, 10, (256,), device=rank)
def fwd_bwd():
model.zero_grad(set_to_none=True)
loss(model(X), y).backward()
def step_time(sync_grads, steps=30):
ctx = contextlib.nullcontext if sync_grads else model.no_sync
for _ in range(10): # warmup
with ctx(): fwd_bwd()
torch.cuda.synchronize(); dist.barrier(); t0 = time.time()
for _ in range(steps):
with ctx(): fwd_bwd()
torch.cuda.synchronize()
return (time.time() - t0) / steps
synced, silent = step_time(True), step_time(False)
if rank == 0:
here = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(here, "comm.json"), "w") as f:
json.dump({"synced_ms": 1e3 * synced, "no_sync_ms": 1e3 * silent}, f)
dist.destroy_process_group()
'''
_ = (DDP_DIR / 'comm_probe.py').write_text(COMM_SCRIPT)
net = resnet18(10, 1) # the gradient bytes, N
net(torch.zeros(1, 1, 64, 64)) # materialize Lazy params
n_bytes = 4 * sum(p.numel() for p in net.parameters())
beta = 4.5e9 # NCCL allreduce, effective bytes/device/s on this box (13.5)
if d2l.num_gpus() >= 2:
torchrun = str(pathlib.Path(sys.executable).parent / 'torchrun')
subprocess.run([torchrun, '--standalone', '--nproc-per-node=2',
str(DDP_DIR / 'comm_probe.py')], check=True)
with open(DDP_DIR / 'comm.json') as f:
comm = json.load(f)
print(f"k=2: predicted comm 2N/beta = {2 * n_bytes / beta * 1e3:.0f} "
f"ms/step; measured {comm['synced_ms'] - comm['no_sync_ms']:.0f} ms "
f"(synced {comm['synced_ms']:.0f}, no_sync {comm['no_sync_ms']:.0f})")The cost model’s 2N/\beta lands within tens of percent of the measurement. Prediction agreeing with measurement is the result.
Configured vs. Default: Factors, Not Percent
§13.5 diagnosed the fallback transport; one documented switch re-routes it over the DMA copy engines:
BENCH_SCRIPT = r'''
import json, os, time, torch
import torch.distributed as dist
rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(rank)
dist.init_process_group("nccl", device_id=torch.device("cuda", rank))
x = torch.randn(16 * 1024 * 1024, device=rank) # 64 MB of fp32
for _ in range(5): # warmup
dist.all_reduce(x)
torch.cuda.synchronize(); t0 = time.time()
for _ in range(15):
dist.all_reduce(x)
torch.cuda.synchronize(); dt = (time.time() - t0) / 15
if rank == 0:
here = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(here, "bench.json"), "w") as f:
json.dump({"gbs": 2 * x.numel() * 4 / dt / 1e9}, f)
dist.destroy_process_group()
'''
_ = (DDP_DIR / 'allreduce_bench.py').write_text(BENCH_SCRIPT)
def bare_allreduce_gbs(**env):
(DDP_DIR / 'bench.json').unlink(missing_ok=True)
torchrun = str(pathlib.Path(sys.executable).parent / 'torchrun')
subprocess.run([torchrun, '--standalone', '--nproc-per-node=2',
str(DDP_DIR / 'allreduce_bench.py')],
check=True, env={**os.environ, **env})
with open(DDP_DIR / 'bench.json') as f:
return json.load(f)['gbs']
if d2l.num_gpus() >= 2:
slow = bare_allreduce_gbs()
# Works around an NCCL transport bottleneck on this P2P-less box
# (13.5's diagnosis) -- a platform-specific fix; validate on yours.
fast = bare_allreduce_gbs(NCCL_SHM_USE_CUDA_MEMCPY='1')
print(f'bare allreduce, effective bytes/device/s: '
f'{slow:.1f} GB/s at the default, {fast:.1f} GB/s configured '
f'({fast / slow:.1f}x)')~5× from configuration alone — yet the same mode deadlocks DDP’s overlapped training on this box, so the notebooks keep the defaults. A workaround is validated per platform and per workload — measure the workload, not just the wire.
Sharding the Redundant: FSDP
DDP replicates the whole 16P-byte state k times. FSDP shards it: allreduce = reduce-scatter + all-gather, kept separate — gather a layer just-in-time, free it after. The §13.5 identity, cashed in.
JAX: Annotate the Layout
One process. Describe the layout; jit the unchanged step; XLA writes the collectives.
imgs, labels = d2l.FashionMNIST().train # raw arrays; resize once
X_all = np.asarray(jax.image.resize(
jnp.float32(imgs[:, :, :, None]) / 255, (len(imgs), 64, 64, 1),
'bilinear'))
y_all = np.asarray(labels, np.int32)
k = min(4, jax.local_device_count())
mesh = make_mesh(k)
model = ResNet18(rngs=nnx.Rngs(0))
opt = nnx.Optimizer(model, optax.sgd(0.1), wrt=nnx.Param)
replicate(model, mesh); replicate(opt, mesh)
Xs, ys = shard_batch(X_all[:256], y_all[:256], mesh)
loss = train_step(model, opt, Xs, ys) # one data-parallel step
print(f'one sharded step: loss {float(loss):.2f}')
print('the batch, split along its leading axis:')
jax.debug.visualize_array_sharding(Xs.reshape(len(Xs), -1))
print('a weight after the step, still replicated:')
jax.debug.visualize_array_sharding(model.net.layers[-1].kernel[...])one sharded step: loss 2.47
the batch, split along its leading axis:
GPU 0
...
GPU 0,1,2,3
Batch sharded, weights replicated — and still replicated after the step: the compiler’s allreduce kept every copy identical.
The Receipt
Nobody wrote a collective — so find it in the compiled program:
hlo = train_step.lower(model, opt, Xs, ys).compile().as_text()
ops = [line.strip() for line in hlo.splitlines()
if ' all-reduce(' in line or ' all-reduce-start(' in line]
sizes = [int(re.search(r'\[(\d+)\]', op).group(1)) for op in ops]
print(f'{len(ops)} all-reduce ops in the compiled step; the largest '
f'sums {max(sizes) / 1e6:.1f}M floats:')
print(ops[sizes.index(max(sizes))][:88] + ' ...')36 all-reduce ops in the compiled step; the largest sums 6.4M floats:
%all-reduce-start.34 = f32[6447680]{0} all-reduce-start(%wrapped_concatenate.34), channe ...Two big allreduces carry the gradients (XLA bucketed them); the small ones are batch-norm statistics — global under jit. Measured k-sweep + a k=1-vs-k=2 equality check close the loop. Change the PartitionSpec, not the code, to move between data, tensor, and FSDP sharding.
Explicit vs. Declarative
| PyTorch | JAX | |
|---|---|---|
| processes | one per GPU | one, all GPUs |
| collectives | explicit | XLA inserts |
| switch strategy | different API | different PartitionSpec |
Neither is better; knowing both is knowing the design space. Past one node → 3D parallelism, the Language Models part.