Data parallelism vs model parallelism in MLOps - CLI Comparison

Introduction

When training large machine learning models, it can take a long time and use a lot of computer power. Data parallelism and model parallelism are two ways to split the work across multiple computers or processors to make training faster and more efficient.

When your dataset is very large and you want to split it across multiple GPUs to train faster.

When your model is too big to fit into the memory of a single GPU and needs to be split across multiple GPUs.

When you want to reduce training time by using multiple processors working together.

When you want to scale your training to use cloud resources efficiently.

When debugging how your model behaves when split across devices.

Commands

This Python code shows how to use data parallelism with PyTorch. It splits the input data batches across multiple GPUs to speed up training. The model is wrapped with nn.DataParallel to do this automatically.

Terminal

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Create a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)
    def forward(self, x):
        return self.linear(x)

# Create dummy data
data = torch.randn(100, 10)
labels = torch.randn(100, 5)
dataset = TensorDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=10)

# Initialize model and optimizer
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Use Data Parallelism
model = nn.DataParallel(model)

# Training loop
for inputs, targets in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = nn.MSELoss()(outputs, targets)
    loss.backward()
    optimizer.step()

print('Training with data parallelism complete')

Expected Output

Training with data parallelism complete

This code shows model parallelism by splitting parts of the model across two GPUs manually. The first layer runs on GPU 0 and the second on GPU 1. Data moves between GPUs during forward pass.

Terminal

import torch
import torch.nn as nn
import torch.optim as optim

# Define a model split across two GPUs
class ModelParallel(nn.Module):
    def __init__(self):
        super(ModelParallel, self).__init__()
        self.part1 = nn.Linear(10, 20).to('cuda:0')
        self.part2 = nn.Linear(20, 5).to('cuda:1')
    def forward(self, x):
        x = self.part1(x.to('cuda:0'))
        return self.part2(x.to('cuda:1'))

model = ModelParallel()
optimizer = optim.SGD(model.parameters(), lr=0.01)

inputs = torch.randn(10, 10)
labels = torch.randn(10, 5).to('cuda:1')

optimizer.zero_grad()
outputs = model(inputs)
loss = nn.MSELoss()(outputs, labels)
loss.backward()
optimizer.step()

print('Training with model parallelism complete')

Expected Output

Training with model parallelism complete

Key Concept

If you remember nothing else, remember: data parallelism splits the data across devices, while model parallelism splits the model itself across devices.

Code Example

MLOps

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)
    def forward(self, x):
        return self.linear(x)

# Dummy data
data = torch.randn(100, 10)
labels = torch.randn(100, 5)
dataset = TensorDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=10)

# Initialize model and optimizer
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Use Data Parallelism if multiple GPUs available
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

model.to('cuda' if torch.cuda.is_available() else 'cpu')

# Training loop
for inputs, targets in dataloader:
    inputs = inputs.to('cuda' if torch.cuda.is_available() else 'cpu')
    targets = targets.to('cuda' if torch.cuda.is_available() else 'cpu')
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = nn.MSELoss()(outputs, targets)
    loss.backward()
    optimizer.step()

print('Training complete')

Common Mistakes

Trying to use data parallelism when the model is too large to fit on one GPU.

Data parallelism requires the whole model to fit on each device, so it fails if the model is too big.

Use model parallelism to split the model across multiple GPUs instead.

Not moving inputs and outputs to the correct devices in model parallelism.

If tensors are not on the right GPU, the code will raise errors or run slowly due to data transfer overhead.

Explicitly move tensors to the correct device before computation.

Assuming data parallelism automatically speeds up training without checking GPU availability.

If only one GPU is available, data parallelism adds overhead without speed benefit.

Check the number of GPUs and use data parallelism only if multiple GPUs are present.

Summary

Data parallelism splits input data batches across multiple devices to speed up training when the model fits on each device.

Model parallelism splits the model itself across devices to handle models too large for one device.

Use data parallelism by wrapping the model with nn.DataParallel and model parallelism by manually assigning model parts to devices.