Anatomy of a DDP Training Script

We'll build a DDP pre-training script step by step, evolving from a single-GPU baseline to a full multi-node script using Qwen-2.5-0.5B as the model.

Step 1: Single-GPU Baseline

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import DataLoader, Dataset

class SimpleTextDataset(Dataset):
    """Minimal dataset: pre-tokenized random sequences for demonstration."""
    def __init__(self, tokenizer, num_samples=10000, seq_length=512):
        self.input_ids = torch.randint(
            0, tokenizer.vocab_size, (num_samples, seq_length)
        )
    
    def __len__(self):
        return len(self.input_ids)
    
    def __getitem__(self, idx):
        input_ids = self.input_ids[idx]
        return {"input_ids": input_ids, "labels": input_ids.clone()}

# --- Setup ---
model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
model = model.cuda()

dataset = SimpleTextDataset(tokenizer)
loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4, pin_memory=True)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# --- Training Loop ---
model.train()
for epoch in range(3):
    for batch in loader:
        input_ids = batch["input_ids"].cuda()
        labels = batch["labels"].cuda()

        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch} complete")

Step 2: Initialize the Process Group

The process group is how PyTorch discovers all participating GPUs. We use torchrun to launch, which sets environment variables (RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT) automatically.

import os
import torch
import torch.distributed as dist
from datetime import timedelta

def setup_distributed():
    """Initialize the distributed process group."""
    dist.init_process_group(
        backend="nccl",
        timeout=timedelta(minutes=30),  # Increase from 10-min default for large clusters
    )
    
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    
    torch.cuda.set_device(local_rank)
    
    # Reproducibility (we'll use rank-dependent seeding in the full script — Step 5)
    torch.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    
    return rank, local_rank, world_size

def cleanup():
    """Destroy the process group."""
    dist.destroy_process_group()

NCCL Timeout: The default NCCL timeout is 10 minutes (PyTorch 2.x). For large clusters (32+ nodes) or when debugging, increase it. If a rank hangs longer than the timeout, PyTorch raises an error instead of hanging indefinitely.

Step 3: Wrap the Model with DDP

from torch.nn.parallel import DistributedDataParallel as DDP

# After model is on the correct GPU:
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

After wrapping, model.forward() works the same as before. DDP hooks into .backward() to fire AllReduce on gradients transparently.

Tip: Access the underlying model with model.module (e.g., for saving checkpoints).

find_unused_parameters: If your model has conditionally unused parameters (common with some HuggingFace models that have optional heads), DDP will hang silently waiting for gradients that never arrive. Fix: DDP(model, device_ids=[local_rank], find_unused_parameters=True). Only enable this if needed — it adds overhead.

Step 4: Use DistributedSampler for Data

from torch.utils.data import DataLoader, DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)

loader = DataLoader(
    dataset,
    batch_size=4,
    sampler=sampler,
    num_workers=4,
    pin_memory=True,
    drop_last=True,
    persistent_workers=True,
)

# In the training loop, set the epoch for proper shuffling:
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # <-- critical for shuffling
    for batch in loader:
        ...

Learning Rate Scaling

When you go from 1 GPU to N GPUs, your effective batch size increases N×. The linear scaling rule adjusts the learning rate accordingly:

base_lr = 1e-4            # LR you'd use on a single GPU
scaled_lr = base_lr * world_size  # Linear scaling

# Alternative: square root scaling (more conservative, often works better)
scaled_lr = base_lr * (world_size ** 0.5)

Rule of thumb: Start with linear scaling. If training is unstable, try square root scaling or reduce the LR manually.

Step 5: Full Working Script

Save this as train_ddp.py:

#!/usr/bin/env python3
"""
PyTorch DDP pre-training script using Qwen2.5-0.5B.
Launch with: torchrun --nproc_per_node=NUM_GPUS train_ddp.py
"""

import os
import signal
from contextlib import nullcontext
import torch
import torch.distributed as dist
from datetime import timedelta
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset, DistributedSampler
from transformers import AutoModelForCausalLM, AutoTokenizer


# ─── Configuration ────────────────────────────────────────────────────────────

BATCH_SIZE_PER_GPU = 4
ACCUMULATION_STEPS = 2
NUM_EPOCHS = 3
BASE_LR = 1e-4
MAX_GRAD_NORM = 1.0          # Gradient clipping
CHECKPOINT_DIR = "/fsx/checkpoints/qwen-ddp"
CHECKPOINT_INTERVAL = 500    # Steps between checkpoints
MODEL_NAME = "Qwen/Qwen2.5-0.5B"


# ─── Dataset ──────────────────────────────────────────────────────────────────

class SimpleTextDataset(Dataset):
    """Random token sequences for demonstration purposes."""
    def __init__(self, vocab_size: int, num_samples: int = 10000, seq_length: int = 512):
        self.input_ids = torch.randint(0, vocab_size, (num_samples, seq_length))
    
    def __len__(self):
        return len(self.input_ids)
    
    def __getitem__(self, idx):
        input_ids = self.input_ids[idx]
        return {"input_ids": input_ids, "labels": input_ids.clone()}


# ─── Distributed Setup ────────────────────────────────────────────────────────

def setup():
    dist.init_process_group(
        backend="nccl",
        timeout=timedelta(minutes=30),  # Default is 10 min; increase for large clusters
    )
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    torch.cuda.set_device(local_rank)
    
    # Reproducibility
    torch.manual_seed(42 + rank)
    torch.cuda.manual_seed_all(42 + rank)
    
    return rank, local_rank, world_size


def cleanup():
    dist.destroy_process_group()


# ─── Checkpointing ───────────────────────────────────────────────────────────

def save_checkpoint(model, optimizer, epoch, step, loss, checkpoint_dir):
    """Save model + optimizer state. Only rank 0 writes."""
    dist.barrier()
    rank = int(os.environ["RANK"])
    
    if rank == 0:
        os.makedirs(checkpoint_dir, exist_ok=True)
        checkpoint = {
            "epoch": epoch,
            "step": step,
            "model_state_dict": model.module.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": loss,
        }
        tmp_path = os.path.join(checkpoint_dir, "checkpoint_tmp.pt")
        final_path = os.path.join(checkpoint_dir, f"checkpoint_epoch{epoch}_step{step}.pt")
        latest_path = os.path.join(checkpoint_dir, "checkpoint_latest.pt")
        
        torch.save(checkpoint, tmp_path)
        os.rename(tmp_path, final_path)
        
        if os.path.islink(latest_path) or os.path.exists(latest_path):
            os.remove(latest_path)
        os.symlink(final_path, latest_path)
        print(f"[Rank 0] Checkpoint saved: {final_path}")
    
    dist.barrier()


def load_checkpoint(model, optimizer, checkpoint_dir):
    """Load checkpoint on all ranks. Returns (epoch, step) to resume from."""
    latest_path = os.path.join(checkpoint_dir, "checkpoint_latest.pt")
    
    if not os.path.exists(latest_path):
        return 0, 0
    
    local_rank = int(os.environ["LOCAL_RANK"])
    checkpoint = torch.load(latest_path, map_location=f"cuda:{local_rank}", weights_only=False)
    
    model.module.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    
    rank = int(os.environ["RANK"])
    if rank == 0:
        print(f"[Rank 0] Resumed from epoch {checkpoint['epoch']}, step {checkpoint['step']}")
    
    dist.barrier()
    return checkpoint["epoch"], checkpoint["step"]


# ─── Main ─────────────────────────────────────────────────────────────────────

def main():
    rank, local_rank, world_size = setup()
    
    # ─── SIGTERM handler for graceful preemption (HyperPod) ────────────────
    should_stop = False
    def sigterm_handler(signum, frame):
        nonlocal should_stop
        should_stop = True
        if rank == 0:
            print("[Rank 0] SIGTERM received — saving checkpoint and exiting gracefully.")
    signal.signal(signal.SIGTERM, sigterm_handler)
    
    def log(msg):
        if rank == 0:
            print(msg)
    
    # ─── Effective batch size ─────────────────────────────────────────────
    effective_batch_size = BATCH_SIZE_PER_GPU * world_size * ACCUMULATION_STEPS
    scaled_lr = BASE_LR * world_size  # Linear scaling rule
    
    log(f"DDP initialized: world_size={world_size}")
    log(f"Effective batch size: {BATCH_SIZE_PER_GPU} × {world_size} × {ACCUMULATION_STEPS} = {effective_batch_size}")
    log(f"Learning rate: {BASE_LR} × {world_size} = {scaled_lr}")
    
    # ─── Model ────────────────────────────────────────────────────────────
    log(f"Loading model: {MODEL_NAME}")
    
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    model = model.to(local_rank)
    model = DDP(model, device_ids=[local_rank])
    
    # ─── Data ─────────────────────────────────────────────────────────────
    dataset = SimpleTextDataset(
        vocab_size=tokenizer.vocab_size,
        num_samples=10000,
        seq_length=512,
    )
    
    sampler = DistributedSampler(
        dataset, num_replicas=world_size, rank=rank, shuffle=True
    )
    
    loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE_PER_GPU,
        sampler=sampler,
        num_workers=4,
        pin_memory=True,
        drop_last=True,
        persistent_workers=True,
    )
    
    # ─── Optimizer ────────────────────────────────────────────────────────
    optimizer = torch.optim.AdamW(model.parameters(), lr=scaled_lr, weight_decay=0.01)
    
    # ─── Resume from checkpoint ───────────────────────────────────────────
    start_epoch, start_step = load_checkpoint(model, optimizer, CHECKPOINT_DIR)
    
    # ─── Training Loop ────────────────────────────────────────────────────
    model.train()
    
    for epoch in range(start_epoch, NUM_EPOCHS):
        sampler.set_epoch(epoch)
        total_loss = 0.0
        num_steps = 0
        optim_step = 0
        
        for step, batch in enumerate(loader):
            # Skip steps already completed when resuming
            # Note: DataLoader still fetches skipped batches — acceptable for this
            # demo dataset. Production: use DataLoader state_dict (PyTorch 2.4+).
            if epoch == start_epoch and step <= start_step and start_step > 0:
                continue
            
            # Check for preemption signal
            if should_stop:
                save_checkpoint(model, optimizer, epoch, step, total_loss / max(num_steps, 1), CHECKPOINT_DIR)
                cleanup()
                return
            
            input_ids = batch["input_ids"].to(local_rank)
            labels = batch["labels"].to(local_rank)
            
            # ─── Gradient accumulation with no_sync ───────────────────────
            # DDP fires AllReduce on every .backward(). For accumulation,
            # skip the sync on intermediate steps to avoid wasted communication.
            is_accumulation_step = (step + 1) % ACCUMULATION_STEPS != 0
            context = model.no_sync() if is_accumulation_step else nullcontext()
            
            with context:
                outputs = model(input_ids=input_ids, labels=labels)
                loss = outputs.loss / ACCUMULATION_STEPS
                loss.backward()
            
            total_loss += loss.item() * ACCUMULATION_STEPS
            num_steps += 1
            
            if not is_accumulation_step:
                # Gradient clipping (essential for stable pre-training)
                torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
                optimizer.step()
                optimizer.zero_grad()
                optim_step += 1
                
                # ─── Logging with proper metric aggregation ───────────────
                if optim_step % 50 == 0:
                    # Aggregate loss across all ranks for accurate global metric
                    loss_tensor = torch.tensor([loss.item() * ACCUMULATION_STEPS], device=local_rank)
                    dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
                    log(f"  Epoch {epoch} | Optim step {optim_step} | Global Avg Loss: {loss_tensor.item():.4f}")
            
            # ─── Periodic checkpoint ──────────────────────────────────────
            global_step = epoch * len(loader) + step
            if global_step > 0 and global_step % CHECKPOINT_INTERVAL == 0:
                save_checkpoint(model, optimizer, epoch, step, loss.item(), CHECKPOINT_DIR)
        
        # Epoch-end aggregated loss
        epoch_loss = torch.tensor([total_loss / max(num_steps, 1)], device=local_rank)
        dist.all_reduce(epoch_loss, op=dist.ReduceOp.AVG)
        log(f"Epoch {epoch} complete | Global Avg Loss: {epoch_loss.item():.4f}")
    
    # ─── Final checkpoint ─────────────────────────────────────────────────
    save_checkpoint(model, optimizer, NUM_EPOCHS - 1, step, epoch_loss.item(), CHECKPOINT_DIR)
    cleanup()


if __name__ == "__main__":
    main()

To test locally on a single node with 4 GPUs:

torchrun --nproc_per_node=4 train_ddp.py