Multi-GPU training in PyTorch - ML Experiment: Train & Evaluate

Experiment - Multi-GPU training

Problem:Train a neural network on the CIFAR-10 dataset using a single GPU. The model trains well but is slow due to limited GPU memory and compute.

Current Metrics:Training accuracy: 85%, Validation accuracy: 80%, Training time per epoch: 120 seconds

Issue:Training is slow because only one GPU is used. We want to speed up training without losing accuracy.

Your Task

Use multi-GPU training to reduce training time per epoch by at least 40% while maintaining validation accuracy above 78%.

Use PyTorch's DataParallel or DistributedDataParallel for multi-GPU training.

Do not change the model architecture or dataset.

Keep batch size per GPU the same as before.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import time

# Check if multiple GPUs are available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
num_gpus = torch.cuda.device_count()
effective_gpus = max(1, num_gpus)
train_batch_size = 128  # Keep batch size per GPU the same

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

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=train_batch_size * effective_gpus, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

# Define a simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleCNN()

# Use DataParallel if multiple GPUs are available
if num_gpus > 1:
    model = nn.DataParallel(model)

model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def train_one_epoch():
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = 100 * correct / total
    return epoch_loss, epoch_acc

def validate():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in testloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    val_acc = 100 * correct / total
    return val_acc

# Measure training time and accuracy
start_time = time.time()
train_loss, train_acc = train_one_epoch()
end_time = time.time()
train_time = end_time - start_time
val_acc = validate()

print(f'Training loss: {train_loss:.4f}, Training accuracy: {train_acc:.2f}%')
print(f'Validation accuracy: {val_acc:.2f}%')
print(f'Training time for one epoch: {train_time:.2f} seconds')

Wrapped the model with torch.nn.DataParallel to enable multi-GPU training.

Moved model and data tensors to the correct device (GPU).

Kept batch size per GPU the same to maintain training stability.

Measured training time per epoch to compare speed improvements.

Results Interpretation

Before Multi-GPU: Training accuracy: 85%, Validation accuracy: 80%, Training time: 120s per epoch

After Multi-GPU: Training accuracy: 84%, Validation accuracy: 79%, Training time: 70s per epoch

Using multiple GPUs with DataParallel speeds up training significantly (about 40% faster) while maintaining similar accuracy. This shows how parallel computing can help train models faster without losing performance.

Bonus Experiment

Try using torch.nn.DistributedDataParallel instead of DataParallel for potentially better multi-GPU scaling.

💡 Hint

DistributedDataParallel requires setting up a process group and launching one process per GPU. It is more efficient for multi-GPU training but needs more setup.

Practice

1. What is the main purpose of using torch.nn.DataParallel in PyTorch?

easy

A. To reduce the model size by pruning layers

B. To convert the model to a CPU-only version

C. To split the model and data across multiple GPUs for faster training

D. To visualize the model architecture

Solution

1. Step 1: Understand the role of DataParallel

   torch.nn.DataParallel is designed to distribute the model and input data across multiple GPUs to speed up training.

2. Step 2: Compare options

   Reducing the model size by pruning layers, converting the model to a CPU-only version, and visualizing the model architecture do not relate to multi-GPU training. Only splitting the model and data across multiple GPUs for faster training correctly describes the purpose.

3. Final Answer:

   To split the model and data across multiple GPUs for faster training -> Option C

4. Quick Check:

   Multi-GPU training = DataParallel usage [OK]

Hint: DataParallel splits model and data across GPUs for speed [OK]

Common Mistakes:

• Thinking DataParallel reduces model size
• Confusing DataParallel with CPU conversion
• Assuming DataParallel is for visualization

2. Which of the following is the correct way to wrap a PyTorch model for multi-GPU training?

easy

A. model = torch.nn.DataParallel(model)

B. model = torch.nn.Sequential(model)

C. model = torch.nn.ModuleList(model)

D. model = torch.nn.Parameter(model)

Solution

1. Step 1: Identify the correct wrapper for multi-GPU

   To enable multi-GPU training, PyTorch provides torch.nn.DataParallel which wraps the model.

2. Step 2: Check other options

   Sequential chains layers, ModuleList stores layers, and Parameter is for tensors, none enable multi-GPU.

3. Final Answer:

   model = torch.nn.DataParallel(model) -> Option A

4. Quick Check:

   Wrap model with DataParallel for multi-GPU [OK]

Hint: Use DataParallel(model) to enable multi-GPU training [OK]

Common Mistakes:

• Using Sequential or ModuleList instead of DataParallel
• Trying to wrap model with Parameter
• Forgetting to wrap the model at all

3. Consider this code snippet for multi-GPU training:

import torch
model = MyModel()
model = torch.nn.DataParallel(model)
model.cuda()
inputs = inputs.cuda()
outputs = model(inputs)
print(outputs.device)

What will be printed if multiple GPUs are available?

medium

A. cuda:0

B. cpu

C. cuda:1

D. cuda:all

Solution

1. Step 1: Understand DataParallel device behavior

   When using DataParallel, the model replicates on all GPUs but the output is gathered on the default device, usually cuda:0.

2. Step 2: Analyze the print statement

   The outputs.device will show the device of the gathered output tensor, which is cuda:0, not all GPUs or CPU.

3. Final Answer:

   cuda:0 -> Option A

4. Quick Check:

   DataParallel output device = cuda:0 [OK]

Hint: DataParallel gathers output on cuda:0 device [OK]

Common Mistakes:

• Expecting output on all GPUs simultaneously
• Thinking output is on CPU
• Assuming output device is cuda:1 or others

4. You wrote this code for multi-GPU training but get an error:

model = MyModel()
model = torch.nn.DataParallel(model)
inputs = inputs.cuda()
outputs = model(inputs)

The error says: 'Expected all tensors to be on the same device'. What is the likely cause?

medium

A. Model should be wrapped with Sequential instead

B. Model is not moved to GPU before wrapping with DataParallel

C. DataParallel requires inputs on CPU

D. Inputs are not on CPU

Solution

1. Step 1: Check model and input device placement

   The model must be moved to GPU (e.g., model.cuda()) before wrapping with DataParallel to ensure all parts are on the same device.

2. Step 2: Analyze error cause

   If the model is on CPU but inputs are on GPU, the error about device mismatch occurs.

3. Final Answer:

   Model is not moved to GPU before wrapping with DataParallel -> Option B

4. Quick Check:

   Model and inputs must be on same device [OK]

Hint: Move model to GPU before DataParallel wrap [OK]

Common Mistakes:

• Wrapping model before moving to GPU
• Assuming inputs must be on CPU
• Using wrong wrapper like Sequential

5. You want to train a model on 4 GPUs using PyTorch. Which of these steps is NOT necessary for correct multi-GPU training?

hard

A. Move input data to GPU before passing to the model

B. Move the model to GPU before wrapping

C. Wrap the model with torch.nn.DataParallel

D. Manually split input batches and send each to a different GPU

Solution

1. Step 1: Review multi-GPU training steps

   Using DataParallel automatically splits input batches and sends them to multiple GPUs, so manual splitting is not needed.

2. Step 2: Confirm necessary steps

   Wrapping the model, moving it to GPU, and moving inputs to GPU are all required for proper multi-GPU training.

3. Final Answer:

   Manually split input batches and send each to a different GPU -> Option D

4. Quick Check:

   DataParallel handles input splitting automatically [OK]

Hint: DataParallel auto-splits inputs; manual split not needed [OK]

Common Mistakes:

• Trying to manually split inputs across GPUs
• Forgetting to move model to GPU
• Not wrapping model with DataParallel