Session 23¶

Topic 3: Neural networks and PyTorch¶

• introduction to PyTorch
• automatic differentiation
• basic neural network models and training
• optimization hyperparameters
• distributed training with data parallelism 👈
• distributed training on multiple compute nodes

Today¶

Distributed training with data parallelism¶

1. basics of distributed training
2. distributed training with PyTorch DDP
3. hands-on: converting single-GPU code to multi-GPU

Distributed Training with PyTorch DDP¶

From learning → scaling¶

so far:

• tensors … store data
• autograd … automatic differentiation
• models … parameterized functions
• training … how the model learns
• optimization param’s … control training

→ now we ask: how to scale to larger datasets

Why do we need distributed training?¶

training is only getting more expensive:

• bigger datasets and
• larger models …
• lead to longer training times ⏳

goal: accelerate training by using more compute:

1. one GPU
2. several GPUs
3. multiple compute nodes with several GPUs each

Parallelize the training dataset¶

we already know this sequence in the training loop:

→ can we run this on multiple GPUs simultaneously?

Idea: data parallelism¶

Setup: copy the same model to all GPUs.

To train, split the data across GPUs:

1. each GPU gets one mini-batch (data shard)
2. each GPU computes its gradient locally (backprop.)
3. then all GPUs synchronize their gradients
4. each GPU updates model parameters using the same gradient

→ all models stay consistent

What is PyTorch DDP?¶

Distributed Data Parallel (DDP) is PyTorch’s solution for multi-GPU training / multi-process training.

Example:

2-node × 2-GPU-per-node setup → world size = 4

DDP execution model¶

for each GPU:

• separate Python process
• separate data shard
• identical model copy

across GPUs:

• gradients are averaged across processes (i.e., GPUs)
• same parameter updates for each process → models stay identical

Workflow in PyTorch code¶

we:

1. initialize distributed processes
2. split dataset across GPUs
3. run training loop (independently)

PyTorch handles most of the hard parts like:

• synchronize gradients automatically

Step 1: launch one process per GPU¶

torch.multiprocessing.spawn(
    worker_fn,       # our own training function
    nprocs=args.num_gpus,
    args=(args,)  # extra arguments for worker_fn()
)

• each process automatically gets a local rank
• a local rank is mapped to a specific GPU

Step 2: initialize process group¶

now, inside worker_fn()

setup required for communication across GPUs:

def worker_fn(local_rank, args):
    global_rank = args.node_id * args.num_gpus + local_rank
    dist.init_process_group(
        backend="nccl",  # use NVIDIA Collective Communication Library (NCCL)
        world_size=WORLD_SIZE,
        rank=global_rank,
        device_id=local_rank,
    )
    atexit.register(dist.destroy_process_group)

• world_size → total number of processes
• rank → unique ID per process

Step 3: assign GPU to process¶

assign each process a different GPU:

def worker_fn(local_rank, args):
    ...

    device = torch.device(f"cuda:{local_rank}")

Step 4: distribute the data¶

train_sampler = DistributedSampler(
    train_set,
    num_replicas=WORLD_SIZE,
    rank=global_rank
)

• num_replicas → total number of processes
• rank → unique ID per process

→ each process sees different data; avoids duplication

Step 5: wrap model with DDP¶

first, ensure all models start with the same weights

torch.manual_seed(base_seed)
model = WideResNet(num_classes)

...

model = nn.parallel.DistributedDataParallel(
    model,
    device_ids=[local_rank]
)

then, wrap the model with DDP for automatic gradient synchronization

Step 6: core training loop (unchanged!)¶

outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()

✨ DDP automatically:

• averages gradients across GPUs
• keeps models in sync

Subtle but important details¶

• call train_sampler.set_epoch(epoch) at each epoch → reshuffles data each epoch
• use dist.barrier() → synchronizes processes
• reduce metrics across GPUs, for example,

# sum the image/sec over all workers (result on main rank only)
torch.distributed.reduce(images_per_sec, MAIN_RANK)

# average validation accuracy across workers (result on all ranks)
torch.distributed.all_reduce(v_accuracy, op=dist.ReduceOp.AVG)

Demo: Fashion-MNIST with DDP¶

code/pytorch/fashion_mnist_ddp

we will:

• start with a singe-GPU code
• write the code for distributed training
• train a CNN on Fashion-MNIST

Hands-on: what to experiment with¶

vary:

• number of GPUs (num_gpus)
• batch size (batch_size)
• learning rate (lr)

… and observe scaling behavior

questions to explore:

• does training get faster?
• does accuracy change?

Key takeaways¶

• DDP = multi-process + synchronized gradients
• each GPU runs its own training loop
• PyTorch handles communication

Looking ahead¶

• training neural nets: from a single machine → data center

Prepare for the next session 🌱¶

sessions/24-todo