Distributed Training Guide

Choosing a Distributed Strategy

Model fits on 1 GPU?
├── YES → Do you need multiple GPUs for speed?
│         ├── YES → Use DDP (simplest, fastest for <8 GPUs)
│         └── NO  → Single GPU training is fine
└── NO  → How big is the model?
          ├── <10B params → Use FSDP (PyTorch-native, no extra deps)
          ├── 10B–100B    → Use FSDP or DeepSpeed ZeRO Stage 3
          └── >100B       → Use DeepSpeed ZeRO Stage 3 + CPU offload
                            or Megatron-LM tensor parallelism

DDP (DistributedDataParallel)

What it does: Replicates the entire model on each GPU. Each GPU processes a different mini-batch. After backward, gradients are averaged across GPUs via all-reduce.

When to use: Model fits on 1 GPU, you want to scale throughput linearly.

Memory overhead: Full model copy per GPU + gradient communication buffers (~25 MB default).

Launch:

torchrun --nproc_per_node=4 train.py

Code changes:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

# Use DistributedSampler for the dataloader
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, batch_size=per_gpu_batch)

See configs/ddp_config.yaml for a complete configuration.

FSDP (Fully Sharded Data Parallel)

What it does: Shards model parameters, gradients, and optimizer states across GPUs. Each GPU holds only 1/N of the model. Parameters are gathered on-the-fly for each layer's forward/backward.

When to use: Model doesn't fit on 1 GPU (after AMP + gradient checkpointing). Preferred over DeepSpeed when you want a pure PyTorch solution.

Memory per GPU: Roughly total_model_memory / num_gpus + activation memory.

Code changes:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision

mp_policy = MixedPrecision(
    param_dtype=torch.float16,
    reduce_dtype=torch.float16,
    buffer_dtype=torch.float16,
)

model = FSDP(
    model,
    mixed_precision=mp_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
)

See configs/fsdp_config.yaml for a complete configuration.

FSDP Wrapping Policy

The wrapping policy determines which modules become separate FSDP units. This is the single most important FSDP tuning decision.

Rule of thumb: Wrap at the transformer-block level. Each block becomes an independently sharded unit.

from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)
model = FSDP(model, auto_wrap_policy=wrap_policy)

DeepSpeed ZeRO

What it does: Progressive sharding of optimizer states (Stage 1), gradients (Stage 2), and parameters (Stage 3).

When to use: You need fine-grained control over sharding, or you want CPU offloading for truly massive models.

Launch:

deepspeed --num_gpus=8 train.py --deepspeed configs/deepspeed_config.json

See configs/deepspeed_config.json for a Stage 2 configuration.

Performance Comparison

Measured on 4×A100 80GB, training a 7B parameter model:

DDP is always the fastest because it has the least communication overhead. FSDP/ZeRO trade throughput for the ability to train larger models.

Multi-Node Training

For training across multiple machines:

1. Network: Use InfiniBand or 100+ Gbps Ethernet. NCCL performance degrades sharply on slow networks.

2. Launch: Use torchrun with rendezvous:

# On node 0:
torchrun --nnodes=2 --nproc_per_node=8 \
  --rdzv_backend=c10d --rdzv_endpoint=node0:29500 \
  --node_rank=0 train.py

# On node 1:
torchrun --nnodes=2 --nproc_per_node=8 \
  --rdzv_backend=c10d --rdzv_endpoint=node0:29500 \
  --node_rank=1 train.py

3. Data: Each node needs access to the dataset (shared filesystem like NFS/Lustre, or download to local SSD on each node).

4. Batch size: Scale learning rate linearly with the number of GPUs. Use a warmup period when scaling to many GPUs.

Common Pitfalls

1. Forgetting DistributedSampler. Without it, every GPU processes the same data → N× wasted compute, wrong gradients.

2. Logging from all ranks. Only log from rank 0:

if dist.get_rank() == 0:
    logger.info("Epoch %d loss: %.4f", epoch, loss)

3. Saving checkpoints from all ranks. Only save from rank 0 (DDP) or use FSDP's full_state_dict API.

4. Not setting NCCL environment variables. On cloud instances:

export NCCL_SOCKET_IFNAME=eth0    # Or your network interface
export NCCL_IB_DISABLE=1          # Disable InfiniBand if not available

GPU Memory Management Guide

How GPU Memory Works

GPU memory (VRAM) is a fixed resource — you can't add more at runtime. Understanding what consumes it is the first step to optimizing it.

What Lives in GPU Memory During Training

Key insight: For a 1B parameter model in fp32 with Adam optimizer, the *minimum* memory is:

• Parameters: 1B × 4 = 4 GB
• Gradients: 1B × 4 = 4 GB
• Adam states (momentum + variance): 1B × 8 = 8 GB
• Total before activations: 16 GB

Activations (intermediate tensors stored for the backward pass) often dominate — they scale with batch size × sequence length × hidden size × number of layers.

The OOM Debugging Flowchart

When you hit torch.cuda.OutOfMemoryError, follow this decision tree:

1. Is this the FIRST iteration?
   YES → Model + optimizer don't fit on GPU
         → Reduce model size, enable FSDP, or use a bigger GPU
   NO  → Continue to 2

2. Does it happen at a RANDOM iteration?
   YES → Memory fragmentation or variable-length inputs
         → Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
         → Pad/truncate inputs to fixed lengths
   NO  → Continue to 3

3. Does it happen at the SAME iteration each time?
   YES → A specific batch is too large, or memory is growing
         → Check for tensor accumulation (see "Common Memory Leaks")
         → Profile with torch.cuda.memory_snapshot()

Optimization Techniques (Ordered by Effort)

1. Reduce Batch Size (Effort: 30 seconds)

The simplest fix. Halving batch size roughly halves activation memory.

Trade-off: Smaller batches = noisier gradients = potentially worse convergence. Compensate with gradient accumulation:

# Instead of batch_size=32 (OOM):
# Use batch_size=8 with 4 accumulation steps = same effective batch
config = TrainerConfig(gradient_accumulation_steps=4)

2. Enable AMP (Effort: 5 minutes)

Mixed-precision training cuts tensor memory by ~50 %. See guides/mixed-precision.md for details.

config = TrainerConfig(use_amp=True)

3. Clear CUDA Cache (Effort: 1 minute)

The PyTorch caching allocator holds onto freed memory blocks for reuse. Sometimes this causes fragmentation:

from src.memory_manager import clear_cuda_cache
clear_cuda_cache()

4. Gradient Checkpointing (Effort: 15 minutes)

Trade compute for memory: don't store activations during forward pass; recompute them during backward. Saves ~60 % memory, costs ~30 % speed.

# For HuggingFace models:
model.gradient_checkpointing_enable()

# For custom models:
from torch.utils.checkpoint import checkpoint
out = checkpoint(layer, input, use_reentrant=False)

5. Use set_to_none=True (Effort: 1 minute)

optimizer.zero_grad(set_to_none=True)  # Frees gradient memory instead of zeroing it

6. PYTORCH_CUDA_ALLOC_CONF (Effort: 2 minutes)

Environment variable that tunes the CUDA memory allocator:

# Reduce fragmentation by allowing the allocator to expand existing segments
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Limit the maximum size of a single cached block (reduces fragmentation)
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512

Common Memory Leaks

1. Accumulating Tensors in a List

# BAD — losses stays on GPU, grows every iteration
losses = []
for batch in dataloader:
    loss = model(batch)
    losses.append(loss)  # Keeps the computation graph alive!

# GOOD — detach and move to CPU
losses = []
for batch in dataloader:
    loss = model(batch)
    losses.append(loss.item())  # .item() returns a Python float

2. Not Deleting Intermediate Tensors

# BAD — embeddings stays in GPU memory even after you don't need it
embeddings = model.encode(text)
predictions = classifier(embeddings)
# embeddings is still alive here

# GOOD — explicitly free it
embeddings = model.encode(text)
predictions = classifier(embeddings)
del embeddings
torch.cuda.empty_cache()

3. Logging Tensors to TensorBoard

# BAD — logs the full tensor (keeps it on GPU)
writer.add_histogram("weights", model.layer.weight)

# GOOD — clone to CPU first
writer.add_histogram("weights", model.layer.weight.detach().cpu())

Memory Monitoring

Use src/memory_manager.py to track memory during training:

from src.memory_manager import memory_tracking, format_bytes

with memory_tracking(label="training_step") as tracker:
    loss = model(batch)
    loss.backward()
    optimizer.step()

print(f"Step allocated {format_bytes(tracker['allocated_delta'])}")

Use src/gpu_monitor.py for continuous monitoring:

from src.gpu_monitor import query_nvidia_smi, detect_anomalies

states = query_nvidia_smi()
for s in states:
    print(s.summary_line())

anomalies = detect_anomalies(states)
for a in anomalies:
    print(f"[{a.severity}] {a.message}")

Memory Budget Calculator

Quick formula to estimate if your model fits on a given GPU:

Total memory needed (bytes) =
    params × 4                          # Parameters (fp32)
  + params × 4                          # Gradients
  + params × 8                          # Adam optimizer states
  + batch_size × seq_len × hidden × layers × 4  # Activations (rough)
  + 500 MB                              # CUDA overhead

With AMP: multiply the first 3 terms by 0.5
With gradient checkpointing: multiply activations by 0.3
With FSDP: divide params/grads/optimizer by num_gpus