DistributedDataParallel in PyTorch - ML Experiment: Train & Evaluate

Experiment - DistributedDataParallel

Problem:You have a deep learning model training on a single GPU. Training is slow and you want to speed it up by using multiple GPUs with PyTorch's DistributedDataParallel (DDP). Currently, the model trains on one GPU with 75% accuracy after 10 epochs.

Current Metrics:Training accuracy: 75%, Validation accuracy: 72%, Training loss: 0.6, Validation loss: 0.65

Issue:Training is slow because it uses only one GPU. You want to use multiple GPUs to speed up training without losing accuracy.

Your Task

Modify the training code to use PyTorch's DistributedDataParallel to train on 2 GPUs, aiming to reduce training time by at least 40% while maintaining validation accuracy above 70%.

Use only 2 GPUs available on the machine.

Keep the model architecture and dataset unchanged.

Do not change batch size or learning rate.

Hint 1

Hint 2

Hint 3

Hint 4

Hint 5

Solution

PyTorch

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
from torchvision.models import resnet18

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)
    torch.manual_seed(42)
    device = torch.device(f'cuda:{rank}')

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    train_dataset = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
    val_dataset = datasets.CIFAR10('./data', train=False, download=True, transform=transform)

    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
    train_loader = DataLoader(train_dataset, batch_size=32, sampler=train_sampler)
    val_loader = DataLoader(val_dataset, batch_size=32)

    model = resnet18(num_classes=10).to(device)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    epochs = 10
    for epoch in range(epochs):
        model.train()
        train_sampler.set_epoch(epoch)
        total_loss = 0.0
        correct = 0
        total = 0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            correct += (predicted == labels).sum().item()
            total += labels.size(0)
        # Aggregate training metrics across GPUs
        loss_tensor = torch.tensor(total_loss, device=device)
        correct_tensor = torch.tensor(correct, device=device)
        total_tensor = torch.tensor(total, device=device)
        dist.all_reduce(loss_tensor)
        dist.all_reduce(correct_tensor)
        dist.all_reduce(total_tensor)
        train_loss = loss_tensor.item() / total_tensor.item()
        train_acc = correct_tensor.item() / total_tensor.item() * 100

        if rank == 0:
            model.eval()
            val_loss = 0.0
            val_correct = 0
            val_total = 0
            with torch.no_grad():
                for inputs, labels in val_loader:
                    inputs, labels = inputs.to(device), labels.to(device)
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                    val_loss += loss.item() * inputs.size(0)
                    _, predicted = torch.max(outputs, 1)
                    val_correct += (predicted == labels).sum().item()
                    val_total += labels.size(0)
            val_loss /= val_total
            val_acc = val_correct / val_total * 100
            print(f'Epoch {epoch+1}/{epochs} - Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%, Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')

    cleanup()

if __name__ == '__main__':
    world_size = 2
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

Initialized distributed process group with 'nccl' backend for GPU communication.

Used torch.multiprocessing.spawn to launch one process per GPU.

Wrapped the model with DistributedDataParallel and moved it to the correct GPU.

Used DistributedSampler to split the training dataset across GPUs.

Ensured each process sets its own device and synchronizes training epochs.

Results Interpretation

Before: Training accuracy: 75%, Validation accuracy: 72%, Training loss: 0.6, Validation loss: 0.65, Training time: 100%

After: Training accuracy: 74%, Validation accuracy: 71%, Training loss: 0.62, Validation loss: 0.66, Training time: 55%

Using DistributedDataParallel allows training on multiple GPUs simultaneously, which speeds up training significantly while maintaining similar accuracy and loss. This shows how parallelism can improve efficiency in deep learning.

Bonus Experiment

Try training with 4 GPUs using DistributedDataParallel and observe how training time and accuracy change.

💡 Hint

Adjust world_size to 4 and ensure your machine has 4 GPUs available. Use the same code structure but increase the number of processes.

Practice

1. What is the main purpose of using DistributedDataParallel in PyTorch?

easy

A. To automatically tune hyperparameters during training

B. To reduce the size of the model for deployment

C. To convert a model to run only on CPU

D. To speed up training by using multiple GPUs or machines simultaneously

Solution

1. Step 1: Understand the role of DistributedDataParallel

   DistributedDataParallel (DDP) is designed to run training across multiple GPUs or machines to speed up the process.

2. Step 2: Compare options with DDP's purpose

   Options B, C, and D do not relate to parallel training or speed-up. Only To speed up training by using multiple GPUs or machines simultaneously correctly describes DDP's main purpose.

3. Final Answer:

   To speed up training by using multiple GPUs or machines simultaneously -> Option D

4. Quick Check:

   DDP speeds training = A [OK]

Hint: DDP means training faster with many GPUs [OK]

Common Mistakes:

• Thinking DDP reduces model size
• Confusing DDP with CPU-only training
• Assuming DDP tunes hyperparameters automatically

2. Which of the following is the correct way to wrap a PyTorch model with DistributedDataParallel assuming model is your neural network and device is the GPU device?

easy

A. model = torch.nn.DataParallel(model, device_ids=[device])

B. model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device])

C. model = torch.nn.DistributedDataParallel(model, device=device)

D. model = torch.nn.parallel.DistributedDataParallel(model, devices=device)

Solution

1. Step 1: Recall the correct syntax for DistributedDataParallel

   The correct class is torch.nn.parallel.DistributedDataParallel and it takes device_ids as a list of devices.

2. Step 2: Check each option for syntax correctness

   model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device]) uses the correct class and parameter device_ids=[device]. model = torch.nn.DataParallel(model, device_ids=[device]) uses DataParallel, which is different. Options C and D use incorrect parameter names or class names.

3. Final Answer:

   model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device]) -> Option B

4. Quick Check:

   Correct class and device_ids list = A [OK]

Hint: Use torch.nn.parallel.DistributedDataParallel with device_ids list [OK]

Common Mistakes:

• Using DataParallel instead of DistributedDataParallel
• Passing device instead of device_ids list
• Wrong class or parameter names

3. Consider the following code snippet for a distributed training setup:

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

def setup(rank, world_size):
    dist.init_process_group('gloo', rank=rank, world_size=world_size)

model = torch.nn.Linear(10, 1).to(rank)
model = DDP(model, device_ids=[rank])

print(model.module.weight.device)

What will be the output device of model.module.weight.device if rank=1 and GPUs are available?

medium

A. cuda:1

B. cpu

C. cuda:0

D. cuda:rank

Solution

1. Step 1: Understand device assignment in the code

   The model is moved to rank device with .to(rank), so for rank=1, it is on cuda:1.

2. Step 2: Check DistributedDataParallel wrapping

   DDP wraps the model and uses device_ids=[rank], so the model stays on cuda:1.

3. Final Answer:

   cuda:1 -> Option A

4. Quick Check:

   Model device matches rank GPU = B [OK]

Hint: Model device matches rank GPU in DDP setup [OK]

Common Mistakes:

• Assuming model stays on CPU
• Confusing device index with rank variable
• Thinking device is always cuda:0

4. You run a distributed training script using DistributedDataParallel but get the error: RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. What is the most likely cause?

medium

A. You forgot to call optimizer.step() after loss.backward()

B. You did not wrap the model with DistributedDataParallel

C. You used the wrong backend in init_process_group

D. You did not call dist.destroy_process_group() at the end

Solution

1. Step 1: Understand the error about unfinished reduction

   This error usually means gradients were not synchronized properly before the next iteration.

2. Step 2: Identify missing optimizer step

   Not calling optimizer.step() after loss.backward() can cause gradients to not be updated and reduction to remain unfinished.

3. Final Answer:

   You forgot to call optimizer.step() after loss.backward() -> Option A

4. Quick Check:

   Missing optimizer.step() causes reduction error = C [OK]

Hint: Always call optimizer.step() after backward() in DDP [OK]

Common Mistakes:

• Skipping optimizer.step() call
• Ignoring proper model wrapping
• Confusing backend errors with sync errors

5. You want to train a model on 4 GPUs using DistributedDataParallel across 2 machines, each with 2 GPUs. Which of the following is the correct approach to set up the training?

hard

A. Launch 4 processes on one machine only, ignoring the second machine

B. Launch 2 processes total, each process uses 2 GPUs, wrap model with DataParallel on each process

C. Launch 4 processes total, each process uses one GPU, initialize process group with correct rank and world_size, wrap model with DDP on each process

D. Launch 1 process per machine, use DDP with device_ids=[0,1] on each machine

Solution

1. Step 1: Understand multi-machine multi-GPU DDP setup

   Each GPU should have one process. For 4 GPUs across 2 machines, launch 4 processes total.

2. Step 2: Confirm process group and model wrapping

   Initialize the process group with correct rank and world_size=4. Wrap the model with DDP on each process using the assigned GPU.

3. Final Answer:

   Launch 4 processes total, each process uses one GPU, initialize process group with correct rank and world_size, wrap model with DDP on each process -> Option C

4. Quick Check:

   One process per GPU with correct init = D [OK]

Hint: One process per GPU, init process group correctly [OK]

Common Mistakes:

• Using fewer processes than GPUs
• Using DataParallel instead of DDP for multi-machine
• Ignoring second machine in setup