Most PyTorch training code starts on one GPU. The moment a model stops fitting, or an epoch starts taking a day, the question becomes how to spread it across several, and in 2026 the PyTorch answer for anything that does not fit is FSDP2. This tutorial explains when you need it, wraps a 7B-parameter model with it, adds activation checkpointing and mixed precision, saves and resumes with distributed checkpointing, and covers the failure modes you will hit on the way.
DDP or FSDP2?
DistributedDataParallel (DDP) replicates the whole model on every GPU and all-reduces gradients. It is simple and fast, and it is the right choice whenever the model, its gradients and optimizer state fit on one GPU with room for activations. For a 7B model in BF16 that is roughly 14 GB of weights plus 14 GB of gradients plus 56 GB of FP32 AdamW state: it does not fit on an 80 GB card, so DDP is out.
FSDP2 (torch.distributed.fsdp.fully_shard) shards parameters, gradients and optimizer state across GPUs, gathering each layer's parameters just in time for its forward and backward. Memory per GPU drops roughly by the number of GPUs; the cost is extra communication. FSDP2 replaced the original FSDP (FullyShardedDataParallel) with a per-parameter DTensor design that is simpler to compose with tensor parallelism, torch.compile and distributed checkpointing. Use it for anything over a few billion parameters, or when DDP runs out of memory. The fully_shard docs are the reference.
Wrapping a 7B model
Install torch (2.x, a recent release) and transformers. The script is launched with torchrun, which sets the rank and world-size environment variables.
# train_fsdp.py
import os
import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
checkpoint_wrapper, apply_activation_checkpointing)
from transformers import AutoModelForCausalLM
MODEL = "meta-llama/Llama-3.1-8B"
def setup():
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
def build_model():
# Materialize on the meta device, then shard, then load weights:
# avoids every rank holding a full copy in host memory.
with torch.device("meta"):
model = AutoModelForCausalLM.from_pretrained(
MODEL, dtype=torch.bfloat16, low_cpu_mem_usage=True)
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16,
reduce_dtype=torch.float32)
# Activation checkpointing on each transformer block.
apply_activation_checkpointing(
model,
checkpoint_wrapper_fn=checkpoint_wrapper,
check_fn=lambda m: m.__class__.__name__.endswith("DecoderLayer"),
)
# Shard each block, then the root. Per-block sharding lets FSDP2
# overlap the all-gather for block n+1 with compute for block n.
for layer in model.model.layers:
fully_shard(layer, mp_policy=mp)
fully_shard(model, mp_policy=mp)
model.to_empty(device="cuda")
return model
The real weights are loaded after sharding: either by calling load_state_dict on a sharded state dict (see checkpointing below), or, for the first run, by loading the pretrained checkpoint through torch.distributed.checkpoint after converting it once. Loading a Hugging Face checkpoint into a meta-initialized FSDP2 model on every rank is the one genuinely fiddly step; the simplest robust pattern is to initialize on rank 0 on CPU, save once with DCP, and have every run load from that.
Three choices in that code matter. param_dtype=bfloat16 with reduce_dtype=float32 keeps communication of gradients in FP32 for numerical safety while computing in BF16. Activation checkpointing recomputes each block's activations in the backward instead of storing them, trading roughly 30% compute for the memory that lets you run sequence length 4096 on an 80 GB card. And sharding per block, not just the root, is what enables compute/communication overlap.
The training loop
def main():
local_rank = setup()
model = build_model()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5, fused=True)
loader = make_loader(rank=dist.get_rank(), world=dist.get_world_size())
for step, batch in enumerate(loader):
ids = batch["input_ids"].cuda(non_blocking=True)
out = model(input_ids=ids, labels=ids)
out.loss.backward()
model.clip_grad_norm_(1.0)
optimizer.step()
optimizer.zero_grad(set_to_none=True)
if step % 50 == 0 and dist.get_rank() == 0:
print(f"step {step} loss {out.loss.item():.4f}")
if step % 500 == 0:
save_checkpoint(model, optimizer, step)
dist.destroy_process_group()
Use a DistributedSampler (or a sharded streaming dataset) so each rank sees a disjoint slice. model.clip_grad_norm_ is the FSDP2-aware version; calling the plain torch.nn.utils function on sharded parameters gives wrong norms. Launch with:
torchrun --nproc_per_node=8 train_fsdp.py # single node
torchrun --nnodes=2 --nproc_per_node=8 \
--rdzv_backend=c10d --rdzv_endpoint=$HEAD:29500 train_fsdp.py # two nodes
Checkpointing and resume
Never gather a sharded model to rank 0 to save it; that defeats the point and runs out of memory at scale. Use torch.distributed.checkpoint (DCP), which writes each rank's shards in parallel and can load into a different world size:
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
get_state_dict, set_state_dict)
def save_checkpoint(model, optimizer, step):
model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save({"model": model_sd, "optim": optim_sd, "step": step},
checkpoint_id=f"ckpt/step-{step}")
def load_checkpoint(model, optimizer, path):
model_sd, optim_sd = get_state_dict(model, optimizer)
state = {"model": model_sd, "optim": optim_sd, "step": 0}
dcp.load(state, checkpoint_id=path)
set_state_dict(model, optimizer, model_state_dict=model_sd,
optim_state_dict=optim_sd)
return state["step"]
Save to shared storage (NFS, FSx, or an S3-backed filesystem) and also save the data-loader position so a resume does not repeat or skip samples. For export to Hugging Face format at the end, DCP's dcp_to_torch_save utility produces a single-file state dict you can load on one CPU host. See the DCP docs.
Debugging OOMs and stragglers
OOM on the first step. Check the order: shard before materializing (to_empty after fully_shard), and confirm activation checkpointing applied (print the model; wrapped blocks show CheckpointWrapper). Reduce micro-batch size and raise gradient accumulation before touching anything else.
OOM after many steps. Usually fragmentation or a growing sequence length. Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and bucket sequences by length.
Slow step time that does not scale. Run torch.profiler on one rank and look at the NCCL kernels. If all-gathers are not overlapping compute, the sharding granularity is wrong (too coarse). If one rank is always waiting, you have a straggler: a bad GPU, a thermally throttled node, or uneven data. nvidia-smi -q -d CLOCK and per-rank step timing find it quickly.
Hangs. Almost always a collective that one rank did not reach: an if rank == 0 around something that calls a collective, or an exception on one rank. Set TORCH_NCCL_ASYNC_ERROR_HANDLING=1 (now the default) so the job dies with a traceback instead of hanging.
When multi-node is premature
Before renting a cluster, exhaust the single node: BF16, activation checkpointing, torch.compile, a fused optimizer, and FSDP2 across 8 GPUs will fine-tune a 70B model on one 8x80 GB node. Multi-node adds interconnect sensitivity, rendezvous failures, and a 20-40% efficiency loss on anything but a well-tuned fabric. It becomes necessary for pretraining-scale runs and for models that do not fit in one node's aggregate memory, and at that point spot instances with frequent DCP checkpoints are the cost play: a preemption costs you the minutes since the last save, not the run.
Our deep learning model development and cloud deployment practices include distributed-training setup and tuning. Contact us if your training runs have outgrown one GPU.