What if a model is too large to fit on one GPU? Data parallelism cannot solve that problem by itself: it puts a complete copy of the model on each device. Model parallelism instead divides the model across devices, so each holds some parameters and performs its part of the forward and backward passes. section 16.5
The parts still have to work together. In a simple three-device setup, device 0 processes a batch and sends its intermediate output, or activations, to device 1. Device 1 processes those activations and sends the result to device 2. This adds communication to an otherwise familiar forward pass. section 16.5
Listing 16.10 Implementing model parallelism using send and recv
... #1
if rank == 0: #2
model_part1 = nn.Linear(1, 2)
input_batch = torch.tensor([[1.0]])
part1_activations = model_part1(input_batch)
print(f"Process {rank} activations from part 1: {part1_activations}")
dist.send(part1_activations, dst=1)
elif rank == 1:
model_part2 = nn.Linear(2, 3)
part1_activations = torch.zeros(2)
dist.recv(part1_activations, src=0) #3
part2_activations = model_part2(part1_activations)
dist.send(part2_activations, dst=2)
print(f"Process {rank} activations from part 2: {part2_activations}")
elif rank == 2:
model_part3 = nn.Linear(3, 1)
part2_activations = torch.zeros(3)
dist.recv(part2_activations, src=1)
part3_output = model_part3(part2_activations)
print(f"Process {rank} output from part 3: {part3_output}")
See the three linear model parts and the torch.distributed send() and recv() calls that pass activations between ranks. section 16.5
Partitioning a model creates useful options, but the simplest setup is not always the fastest:
- Basic model parallelism can leave devices idle while another device handles its stage. section 16.5.1
- Pipeline parallelism splits input into microbatches so different stages can work concurrently. section 16.5.1
- Tensor parallelism partitions individual layers or tensors; n-dimensional parallelism combines strategies such as model and data parallelism. section 16.12
Partition and replicate
Imagine two model replicas, each split across four devices. One mesh dimension identifies replicas; the other identifies model partitions. PyTorch's DeviceMesh can organize the process groups needed for communication within those dimensions. section 16.6
