Model parallelism and data parallelism are two approaches used to distribute the training of a deep learning model across multiple processors or devices.
In data parallelism, the model is replicated on each device, and each device trains the model on a different batch of data at the same time. The gradients are then aggregated and applied to the original model, updating its parameters. This approach is particularly useful for models that can fit entirely in the memory of individual devices, as it allows for efficient use of available resources. Data parallelism can be scaled horizontally (by adding more devices), which increases the amount of data that can be processed in parallel, and therefore the overall training speed.
Here’s an example of data parallelism in PyTorch:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data.distributed as dist
# Define model
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
x = nn.functional.softmax(x, dim=1)
return x
# Initialize model and optimizer
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# Initialize distributed data parallel training
world_size = torch.cuda.device_count()
dist.init_process_group(backend='nccl', init_method='tcp://localhost:1234', world_size=world_size, rank=torch.cuda.current_device())
model = nn.parallel.DistributedDataParallel(model)
# Load data
train_data = torch.utils.data.DataLoader(...)
test_data = torch.utils.data.DataLoader(...)
# Train model
for epoch in range(num_epochs):
total_loss = 0.0
for data in train_data:
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = nn.functional.cross_entropy(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Epoch {epoch + 1}: Loss={total_loss/len(train_data)}")
# Evaluate model
with torch.no_grad():
total_correct = 0
for data in test_data:
inputs, labels = data
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total_correct += (predicted == labels).sum().item()
accuracy = total_correct / len(test_data.dataset)
print(f"Accuracy={accuracy}")In model parallelism, on the other hand, the model is divided into multiple parts, and each part is assigned to a different device. The forward and backward passes are then executed by passing the intermediate activations or gradients between devices. This approach is particularly useful for large models that cannot fit entirely in the memory of individual devices, as it allows for the distribution of both the model and the data across multiple devices.
Here’s an example of model parallelism in PyTorch:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
# Define model
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
self.relu2 = nn.ReLU(inplace=True)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
self.relu3 = nn.ReLU(inplace=True)
self.conv4 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
self.relu4 = nn.ReLU(inplace=True)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv5 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
self.relu5 = nn.ReLU(inplace=True)
self.conv6 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.relu6 = nn.ReLU(inplace=True)
self.conv7 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.relu7 = nn.ReLU(inplace=True)
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv8 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
self.relu8 = nn.ReLU(inplace=True)
self.conv9 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.relu9 = nn.ReLU(inplace=True)
self.conv10 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.relu10 = nn.ReLU(inplace=True)
self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv11 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.relu11 = nn.ReLU(inplace=True)
self.conv12 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.relu12 = nn.ReLU(inplace=True)
self.conv13 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.relu13 = nn.ReLU(inplace=True)
self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(512 * 7 * 7, 4096)
self.relu14 = nn.ReLU(inplace=True)
self.dropout1 = nn.Dropout()
self.fc2 = nn.Linear(4096, 4096)
self.relu15 = nn.ReLU(inplace=True)
self.dropout2 = nn.Dropout()
self.fc3 = nn.Linear(4096, 1000)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
x = self.conv1(x)
x = self.relu1(x)
x = self.conv2(x)
x = self.relu2(x)
x = self.pool1(x)
x = self.conv3(x)
x = self.relu3(x)
x = self.conv4(x)
x = self.relu4(x)
x = self.pool2(x)
x = self.conv5(x)
x = self.relu5(x)
x = self.conv6(x)
x = self.relu6(x)
x = self.conv7(x)
x = self.relu7(x)
x = self.pool3(x)
x = self.conv8(x)
x = self.relu8(x)
x = self.conv9(x)
x = self.relu9(x)
x = self.conv10(x)
x = self.relu10(x)
x = self.pool4(x)
x = self.conv11(x)
x = self.relu11(x)
x = self.conv12(x)
x = self.relu12(x)
x = self.conv13(x)
x = self.relu13(x)
x = self.pool5(x)
x = x.view(-1, 512 * 7 * 7)
x = self.fc1(x)
x = self.relu14(x)
x = self.dropout1(x)
x = self.fc2(x)
x = self.relu15(x)
x = self.dropout2(x)
x = self.fc3(x)
x = self.softmax(x)
return x
# Initialize model and optimizer
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# Initialize distributed model parallel training
world_size = torch.cuda.device_count()
rank = torch.distributed.get_rank()
backend = 'nccl'
torch.distributed.init_process_group(backend=backend, init_method='tcp://localhost:23456', world_size=world_size, rank=rank)
torch.cuda.set_device(rank)
chunk_sizes = [53, 96, 96, 96, 96, 96, 96, 96, 96, 96, 96, 96, 16]
devices = list(range(world_size))
model.to(devices[rank])
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank, find_unused_parameters=True)
input_shape = (32, 3, 224, 224)
# Train model
for epoch in range(num_epochs):
total_loss = 0.0
for i, chunk_size in enumerate(chunk_sizes):
inputs = torch.randn(*input_shape[rank], chunk_size, requires_grad=True)
labels = torch.randint(0, 1000, (chunk_size,))
optimizer.zero_grad()
outputs = model(inputs)
loss = nn.functional.cross_entropy(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
total_loss_tensor = torch.tensor(total_loss).to(device=rank)
dist.all_reduce(total_loss_tensor, op=dist.ReduceOp.SUM)
print(f"Epoch {epoch + 1}: Loss={total_loss_tensor.item()/len(chunk_sizes)}")
# Evaluate model
with torch.no_grad():
total_correct = 0
for i, chunk_size in enumerate(chunk_sizes):
inputs = torch.randn(*input_shape[rank], chunk_size, requires_grad=False)
labels = torch.randint(0, 1000, (chunk_size,))
outputs = model(inputs)
if rank == 0:
_, predicted = torch.max(outputs.data, 1)
total_correct += (predicted == labels).sum().item()
accuracy = total_correct / sum(chunk_sizes)
print(f"Accuracy={accuracy}")In summary, data parallelism and model parallelism are two different approaches used to scale the training of deep learning models in PyTorch. Data parallelism is used to split the dataset across multiple devices and train the model on each split in parallel, while model parallelism is used to split the model across multiple devices and compute the forward and backward passes in parallel. Both approaches have different advantages and challenges, and the choice of approach depends on the specific requirements of the model and the hardware available.