4.2 PyTorch data parallelism from principles to practice

Data parallelism is the most basic and most used parallel strategy: each GPU holds a full model replica, processes a different data shard, then syncs gradients. The evolution and memory/communication principles are in 4.1 Data parallelism in depth; this article picks up from the principles and focuses on engineering practice—from PyTorch communication basics through complete runnable DDP and FSDP code, multi-node multi-GPU launch, and common pitfalls, so you can land data parallelism in a real training script.

📑 Contents

• 1. PyTorch distributed communication basics
• 2. DDP: the standard data-parallel path
• 3. FSDP: memory optimization via parameter sharding
• 4. DDP vs FSDP: how to choose
• 5. torchrun and multi-node training practice
• 6. Common distributed-training pitfalls
• Summary
• Self-check list
• References

---

1. PyTorch distributed communication basics

Before DDP and FSDP, you need the communication they sit on.

1.1 Core concepts

1.2 Choosing a communication backend

PyTorch supports three backends:

📌 Key point: GPU distributed training almost always uses the NCCL backend. NCCL (NVIDIA Collective Communications Library) is deeply tuned for NVLink, NVSwitch, InfiniBand, and other NVIDIA interconnects.

1.3 Collective communication primitives

Data sync in distributed training relies on collective communication. The ones you use most:

💡 Tip: Understanding AllReduce = ReduceScatter + AllGather is key to how DDP and FSDP differ in communication pattern.

1.4 Initializing communication

Every distributed job starts by initializing the process group:

import torch.distributed as dist

# Option 1: init from env vars (most common, with torchrun)
dist.init_process_group(backend="nccl")

# Option 2: pass args explicitly
dist.init_process_group(
    backend="nccl",
    init_method="tcp://192.168.1.1:29500",
    world_size=8,
    rank=0,
)

With torchrun, MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE, LOCAL_RANK, and friends are set for you, so the code only needs dist.init_process_group(backend="nccl").

---

2. DDP: the standard data-parallel path

DDP (DistributedDataParallel) is the most mature, most widely used data-parallel path in PyTorch.

2.1 How DDP works

DDP’s core mechanism in three lines:

1. Init: Broadcast rank 0’s model params to every process so everyone starts aligned
2. Forward + backward: each process computes gradients independently on its data shard
3. Gradient sync: during backward, DDP AllReduces gradients across processes and averages them

⚠️ Note: DDP does not wait for the whole backward to finish before syncing. It uses a Bucket mechanism: parameter gradients are grouped into Buckets in reverse-compute order; as soon as a Bucket’s gradients are ready, AllReduce starts, so compute overlaps communication. That overlap is a key DDP performance win.

2.2 Complete DDP training code

A DDP training example you can run as-is:

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms


def setup():
    """Initialize the distributed environment"""
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank


def cleanup():
    """Tear down the distributed environment"""
    dist.destroy_process_group()


def main():
    local_rank = setup()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 1. Build the model and wrap with DDP
    model = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ).to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # 2. Build the data loader (must use DistributedSampler)
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.view(-1)),
    ])
    dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2)

    # 3. Optimizer and loss
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # 4. Training loop
    model.train()
    for epoch in range(5):
        sampler.set_epoch(epoch)  # Important: different shuffle seed each epoch
        total_loss = 0.0
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(local_rank), target.to(local_rank)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()  # DDP AllReduces automatically here
            optimizer.step()

            total_loss += loss.item()

        if rank == 0:
            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch}: avg_loss = {avg_loss:.4f}")

    cleanup()


if __name__ == "__main__":
    main()

Launch:

# Single-node 4 GPUs
torchrun --nproc_per_node=4 train_ddp.py

2.3 DDP usage essentials

A few easy-to-miss but critical details when using DDP:

1. You must use DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

Without DistributedSampler, every process reads the same data and data parallelism is pointless.

2. Call sampler.set_epoch(epoch) every epoch

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # reshuffle data assignment across epochs
    for batch in dataloader:
        ...

Skipping set_epoch leaves each epoch with the same per-process data assignment.

3. Log and save only on rank 0

if dist.get_rank() == 0:
    print(f"Loss: {loss.item()}")
    torch.save(model.module.state_dict(), "checkpoint.pt")  # note: model.module

⚠️ Note: save with model.module.state_dict(), not model.state_dict(). DDP prefixes parameter names with module.; saving the wrapper directly breaks keys on load.

2.4 DDP’s limit

DDP is simple and fast, but has one fundamental limit: every GPU holds a full model replica. That means:

• Params + gradients + optimizer state must fit in one GPU’s memory
• For 7B+ models, a single 80 GB GPU is often not enough

That is the problem FSDP solves.

---

3. FSDP: memory optimization via parameter sharding

FSDP (FullyShardedDataParallel) is PyTorch’s native take on DeepSpeed ZeRO—the core idea is not to store a full model on every GPU, but to shard params, gradients, and optimizer state.

3.1 Understanding FSDP from ZeRO

ZeRO (Zero Redundancy Optimizer)‘s insight: in DDP every GPU redundantly stores full optimizer state, gradients, and params—a huge waste.

ZeRO removes redundancy in three stages:

Example: 7B model, 8 GPUs (ZeRO paper’s 16 bytes/param estimate: FP16 params 2B + FP16 grads 2B + FP32 optimizer state 12B):

• DDP: ~112 GB per GPU (does not fit an 80 GB H100)
• FSDP (FULL_SHARD): ~$\frac{112}{8}$ = ~14 GB per GPU (fits easily)

3.2 FSDP’s workflow

FSDP’s flow is richer than DDP: params are “assembled on demand” and “discarded when done”:

1. Initial state: each GPU only stores the param shards it owns
2. Forward: for a layer, AllGather all shards into full params → compute → free non-local shards
3. Backward: AllGather full params again → compute grads → ReduceScatter grads back to each GPU → free full params
4. Update: each GPU only updates its own shards’ params and optimizer state

📌 Key point: FSDP trades communication for memory. Every layer’s forward and backward need one AllGather + one ReduceScatter, so communication volume is higher than DDP, but memory drops a lot.

3.3 FSDP sharding strategies

PyTorch FSDP offers four sharding strategies with a lot of flexibility:

💡 Tip: HYBRID_SHARD is the recommended multi-node strategy. Inside a node it runs FSDP (high-bandwidth NVLink); across nodes it runs data parallel (fits lower InfiniBand bandwidth)—the best balance of communication efficiency and memory savings.

3.4 Complete FSDP training code

import os
import functools
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms


def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank


def cleanup():
    dist.destroy_process_group()


class SimpleMLP(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=2048, num_layers=8, output_dim=10):
        super().__init__()
        layers = []
        for i in range(num_layers):
            in_d = input_dim if i == 0 else hidden_dim
            out_d = output_dim if i == num_layers - 1 else hidden_dim
            layers.append(nn.Linear(in_d, out_d))
            if i < num_layers - 1:
                layers.append(nn.ReLU())
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


def main():
    local_rank = setup()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 1. Create the model
    model = SimpleMLP().to(local_rank)

    # 2. Mixed-precision config
    bf16_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )

    # 3. Auto-wrap policy (submodules over the threshold are sharded separately)
    wrap_policy = functools.partial(
        size_based_auto_wrap_policy,
        min_num_params=1_000_000,
    )

    # 4. Wrap with FSDP
    model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        mixed_precision=bf16_policy,
        auto_wrap_policy=wrap_policy,
        device_id=local_rank,
    )

    # 5. Data loading
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.view(-1)),
    ])
    dataset = datasets.MNIST("./data", train=True, download=(rank == 0), transform=transform)
    dist.barrier()  # wait for rank 0 to finish download
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=2)

    # 6. Optimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # 7. Training loop
    model.train()
    for epoch in range(5):
        sampler.set_epoch(epoch)
        total_loss = 0.0
        for data, target in dataloader:
            data, target = data.to(local_rank), target.to(local_rank)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        if rank == 0:
            print(f"Epoch {epoch}: avg_loss = {total_loss / len(dataloader):.4f}")

    cleanup()


if __name__ == "__main__":
    main()

3.5 Saving and loading FSDP models

Saving an FSDP model is harder than DDP because params are sharded. Prefer PyTorch Distributed Checkpoint (DCP), with helpers from torch.distributed.checkpoint.state_dict to get/write sharded state correctly:

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_state_dict,
    set_state_dict,
)

# Save (all ranks participate): helpers correctly extract FSDP-sharded model / optimizer state
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save(
    {"model": model_sd, "optim": optim_sd},
    checkpoint_id="./checkpoints/epoch_5",
)

# Load (all ranks participate): take current sharded placeholders, fill in place via load, then write back
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.load(
    {"model": model_sd, "optim": optim_sd},
    checkpoint_id="./checkpoints/epoch_5",
)
set_state_dict(model, optimizer, model_state_dict=model_sd, optim_state_dict=optim_sd)

⚠️ Note: do not save an FSDP model with torch.save(model.state_dict()) directly. FSDP’s state_dict behavior depends on state_dict_type; a naive save may only dump local shards. DCP + get_state_dict / set_state_dict handles sharded save/restore correctly and supports resharding (save on 8 GPUs, restore on 4).

3.6 FSDP2: the next-gen API

PyTorch is pushing a new FSDP API—torch.distributed.fsdp.fully_shard (usually called FSDP2). Unlike FSDP1’s module-wrapper style, FSDP2 uses composable per-parameter sharding:

from torch.distributed.fsdp import fully_shard

# FSDP2 style: call fully_shard on modules directly
for layer in model.layers:
    fully_shard(layer)
fully_shard(model)

FSDP2 upsides:

• Combines better with other parallel strategies such as tensor parallelism (TP)
• Built on the DTensor (Distributed Tensor) abstraction for a more unified API
• Finer control over sharding grain

💡 Tip: as of PyTorch 2.x, FSDP2 is usable but still iterating fast. New projects can try FSDP2; production still tends to prefer the more stable FSDP1.

---

4. DDP vs FSDP: how to choose

Selection tips:

• ✅ Params + optimizer state fit one GPU → use DDP
• ✅ Full model does not fit one GPU → use FSDP
• ✅ Multi-node + large model → use FSDP + HYBRID_SHARD
• ✅ Huge models (10B+) → consider FSDP + TP + PP (Megatron-LM or DeepSpeed)

---

5. torchrun and multi-node training practice

5.1 What torchrun is

torchrun is PyTorch’s recommended distributed launcher (replacing the deprecated torch.distributed.launch). Its core capabilities:

• Auto-set env vars: RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT, etc.
• Elastic training: nodes can join/leave dynamically
• Fault tolerance: worker failures can trigger automatic restart

5.2 Single-node multi-GPU launch

# Single-node 4 GPUs
torchrun --nproc_per_node=4 train.py --batch_size 64

# Pin master port (avoid port conflicts)
torchrun --nproc_per_node=4 --master_port=29501 train.py

5.3 Multi-node multi-GPU launch

Multi-node training starts torchrun on each machine with the same rendezvous info:

# Machine 0 (master)
torchrun \
    --nnodes=2 \
    --nproc_per_node=4 \
    --node_rank=0 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

# Machine 1
torchrun \
    --nnodes=2 \
    --nproc_per_node=4 \
    --node_rank=1 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

5.4 Elastic training

torchrun supports elastic training—node count can move inside a range:

# At least 1 machine, at most 4
torchrun \
    --nnodes=1:4 \
    --nproc_per_node=4 \
    --rdzv_id=my_training_job \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

⚠️ Note: elastic training needs the script to save and restore checkpoints correctly. When nodes join or leave, torchrun restarts all workers; the script must resume from the latest checkpoint.

5.5 Reading distributed info in training code

With torchrun, read distributed info from env vars:

import os

local_rank = int(os.environ["LOCAL_RANK"])     # GPU id on this machine
rank = int(os.environ["RANK"])                  # global process id
world_size = int(os.environ["WORLD_SIZE"])      # total processes
local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])  # processes on this machine

---

6. Common distributed-training pitfalls

6.1 Deadlock (Hang)

The most common failure is “stuck forever.” Usual causes:

6.2 Useful debug env vars

# Turn on NCCL debug logs
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

# Raise collective timeout (NCCL backend defaults to 10 minutes); set in code:
# dist.init_process_group(..., timeout=datetime.timedelta(seconds=3600))

# Watchdog heartbeat timeout: how long the monitor waits on a stuck collective before killing the process (PyTorch 2.x)
# Note: this controls "stuck detection/abort", not extending a single collective's timeout
export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=600

# PyTorch distributed logging
export TORCH_DISTRIBUTED_DEBUG=DETAIL

# Pin visible GPUs (avoid GPU id mismatches)
export CUDA_VISIBLE_DEVICES=0,1,2,3

6.3 Random-seed consistency

Randomness in distributed training needs care:

import torch
import random
import numpy as np

def set_seed(seed, rank):
    # Model init needs a shared seed (so DDP ranks start with the same params)
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    # Augmentation etc. need different seeds (avoid identical data across ranks)
    torch.manual_seed(seed + rank)

📌 Key point: DDP needs identical initial params across ranks (DDP Broadcasts from rank 0 at construct time), but augmentation, Dropout, and similar should differ across ranks—otherwise “data parallel” loses diversity.

6.4 Learning-rate scaling

Effective batch size under data parallelism = per-GPU batch size × world_size. When world_size grows, you usually scale LR in proportion:

$$lr_{\text{effective}} = lr_{\text{base}} \times \text{world\_size}$$

Linear scaling can destabilize training when world_size is huge. In practice people often pair it with LR warmup:

# Linear warmup + linear scaling
base_lr = 1e-4
effective_lr = base_lr * world_size
warmup_steps = 1000

for step in range(total_steps):
    if step < warmup_steps:
        lr = effective_lr * (step / warmup_steps)
    else:
        lr = effective_lr  # or add cosine decay later
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

💡 Tip: linear LR scaling is a rule of thumb, not absolute. At large scale (hundreds of GPUs), people sometimes scale by $\sqrt{\text{world\_size}}$ (the LAMB/LARS idea) for stability.

---

📝 Summary

This article picks up 4.1’s principles and works through PyTorch data-parallel engineering:

• Communication basics: World Size / Rank / Local Rank, the NCCL backend, AllReduce / AllGather / ReduceScatter and other collectives
• DDP: the most mature data-parallel path for models that fit one GPU, with complete runnable code. Core is bucketed AllReduce gradient sync, plus DistributedSampler + set_epoch and model-save details
• FSDP: parameter sharding for models that do not fit one GPU, with complete code and four sharding strategies. Core is on-demand AllGather + ReduceScatter; prefer DCP for saves
• torchrun: the recommended launcher for single-node, multi-node, and elastic training
• Common issues: hang debugging, seed management, LR scaling

With these practice skills, next steps are 3D parallelism (TP + PP + DP), DeepSpeed ZeRO, Megatron-LM, and other advanced distributed-training stacks.

---

🎯 Self-check list

• Explain the core difference between DDP and FSDP (full replica vs param sharding)
• Sketch the AllReduce gradient-sync data flow in DDP
• State what each of ZeRO’s three stages shards, and which FSDP strategy matches
• Convert a single-GPU training script to multi-GPU with PyTorch DDP
• Configure FSDP sharding strategy, mixed precision, and auto-wrap policy correctly
• Launch single-node multi-GPU and multi-node multi-GPU training with torchrun
• Explain what DistributedSampler does and why set_epoch is required
• List 3 common hang causes and how to debug them
• Compute effective batch size under data parallelism and scale LR correctly
• Save and restore an FSDP model with Distributed Checkpoint (DCP)

---

📚 References

• PyTorch Distributed Communication Package — torch.distributed
• PyTorch DistributedDataParallel — Getting Started
• PyTorch FSDP — FullyShardedDataParallel
• PyTorch torchrun (Elastic Launch)
• PyTorch Distributed Checkpoint
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
• NCCL Documentation