Model parallelism trains a model that cannot fit on one device by splitting its neural network into sequential subnetworks and placing each subnetwork on a different GPU or CPU. Intermediate outputs move from one device to the next during the forward computation, allowing the full model to run across a group of devices. section 4.1 section 4.2.1
Figure 4.6 Split a four-layer, fully connected deep learning network into four subgroups; each group has one layer, and each subgroup runs on one GPU.
A four-layer network can be partitioned into four single-layer subnetworks, with one subnetwork assigned to each GPU. section 4.4.2
Two-GPU PyTorch layout
Initialize two subnetworks, assign each to a separate GPU, run the input through the first subnetwork, then copy its output to the second GPU for the next computation. This layout is architecture-specific: different network structures require different partitioning code. section 4.4.2
Listing 4.2 A sample model parallelism implementation in PyTorch
gpu1 = 1
gpu2 = 2
class a_large_model(nn.Module):
def __init__(self):
super().__init__()
# initialize the network as two subnetworks.
self.subnet1 = ...
self.subnet2 = ...
# put subnetwork 1 and 2 to two different GPUs
self.subnet1.cuda(gpu1)
self.subnet2.cuda(gpu2)
def forward(x):
# load data to GPU 1 and calculate output for
# subnet 1, GPU 2 is idle at the moment.
x = x.cuda(gpu1)
x = self.subnet1(x)
# move the output of subnet 1 to GPU 2 and calculate
# output for subnet 2. GPU 1 is idle
x = x.cuda(gpu2)
x = self.sub_network2(x)
return x
The listing shows explicit GPU placement for two subnetworks and explicit transfer of intermediate output between them. section 4.4.2
Use model parallelism with its operational trade-offs in mind:
- Choose it when a neural network exceeds the memory limit of a single GPU; data parallelism requires every device to hold the complete model. section 4.2.1
- Expect substantial, case-by-case code changes to split the network, assign devices, and copy intermediate compute outputs; there is no general splitting paradigm. section 4.1
- Account for low device utilization: sequential dependencies can leave all but one GPU idle. Pipeline parallelism improves this by processing different microbatches concurrently. section 4.1 section 4.4.2
- Consider memory-saving approaches such as gradient accumulation or memory swap as alternatives when a model does not fit, recognizing that they are easier to implement but slower. section 4.5
