Four ways to split one training run across two GPUs

Four ways to split one training run across two GPUs An architecture diagram generated by Archify. DP · GPU 0 · full model · half batch · DP: split the batch DP · GPU 0 full model · half batch DP · GPU 1 · full model · half batch · DP: split the batch DP · GPU 1 full model · half batch FSDP · GPU 0 · 1/2 weights, grads, Adam · FSDP: split batch + state FSDP · GPU 0 1/2 weights, grads, Adam FSDP · GPU 1 · other 1/2 · half batch · FSDP: split batch + state FSDP · GPU 1 other 1/2 · half batch TP · GPU 0 · half of every matrix · TP: split inside a layer TP · GPU 0 half of every matrix TP · GPU 1 · other half · full batch · TP: split inside a layer TP · GPU 1 other half · full batch PP · GPU 0 · layers 1–16 · PP: split by layer PP · GPU 0 layers 1–16 PP · GPU 1 · layers 17–32 · PP: split by layer PP · GPU 1 layers 17–32 all-reduce gradients all-gather, reduce-scatter all-reduce every layer send activations DP: split the batch FSDP: split batch + state TP: split inside a layer PP: split by layer Legend GPU holds a full copy of the model GPU holds only part of the model