Data x pipeline parallelism
2 DP replicas / 3 PP stages / 6 GPUs
Two data-parallel replicas, each split into three pipeline stages
Forward splits a global batch into two shards and pipelines each shard across three GPUs. Backward sends activation gradients in reverse and AllReduces parameter gradients between matching pipeline stages.
PP stage 0
W0 replicated
PP stage 1
W1 replicated
PP stage 2
W2 replicated
DP replica 0 / shard B0
DP replica 1 / shard B1
GPU 0 / stage 0
layers
0-1
layers
2-3
GPU 1 / stage 1
layers
4-5
layers
6-7
GPU 2 / stage 2
layers
8-9
layers
10-11
GPU 3 / stage 0
layers
0-1
layers
2-3
GPU 4 / stage 1
layers
4-5
layers
6-7
GPU 5 / stage 2
layers
8-9
layers
10-11
global batch
shard B0
shard B1
activation
activation
activation
activation
loss
L0
loss
L1
each batch shard follows its own three-stage pipeline
loss grad
dL0
loss grad
dL1
input-gradient send
input-gradient send
input-gradient send
input-gradient send
ALL-REDUCE
gradients dW0
ALL-REDUCE
gradients dW1
ALL-REDUCE
gradients dW2
matching PP stages synchronize parameter gradients across DP replicas
forward
backward