Data x pipeline x expert parallelism 2 DP replicas / 2 PP stages / 2 EP ranks / 8 GPUs
Two data-parallel replicas, each containing two pipeline stages and two expert ranks Four input shards enter the first pipeline stage across two data replicas. Pipeline transfers run horizontally, expert All-to-All exchanges run vertically within each stage, and the backward pass uses separate expert-gradient and non-expert-gradient AllReduces across the data replicas. PP axis DP replica 0 / input shards 0, 1 PP stage 0 / early layers PP stage 1 / later layers EP 0 EP 1 EP axis rank 0 / DP 0 / PP 0 / EP 0 shared + router E0, E1 rank 2 / DP 0 / PP 1 / EP 0 shared + router E0, E1 rank 1 / DP 0 / PP 0 / EP 1 shared + router E2, E3 rank 3 / DP 0 / PP 1 / EP 1 shared + router E2, E3 input shard 0 input shard 1 residual send residual send loss L0 loss L1 ALL-TO-ALL dispatch + combine ALL-TO-ALL dispatch + combine loss grad dL0 loss grad dL1 input-gradient send input-gradient send input grad dx0 input grad dx1 ALL-TO-ALL token-gradient routing ALL-TO-ALL token-gradient routing DP axis independent batch shards / no forward communication EXPERT-GRADIENT ALL-REDUCE fixed (PP, EP), across DP replicas groups {0,4} {1,5} {2,6} {3,7} NON-EXPERT-GRADIENT ALL-REDUCE fixed PP, across DP and EP / shared + router parameters groups {0,1,4,5} {2,3,6,7} DP replica 1 / input shards 2, 3 PP stage 0 / early layers PP stage 1 / later layers EP 0 EP 1 EP axis rank 4 / DP 1 / PP 0 / EP 0 shared + router E0, E1 rank 6 / DP 1 / PP 1 / EP 0 shared + router E0, E1 rank 5 / DP 1 / PP 0 / EP 1 shared + router E2, E3 rank 7 / DP 1 / PP 1 / EP 1 shared + router E2, E3 input shard 2 input shard 3 residual send residual send loss L2 loss L3 ALL-TO-ALL dispatch + combine ALL-TO-ALL dispatch + combine loss grad dL2 loss grad dL3 input-gradient send input-gradient send input grad dx2 input grad dx3 ALL-TO-ALL token-gradient routing ALL-TO-ALL token-gradient routing DP replicas run independently; PP sends right; EP routes vertically within each stage PP gradients move left; EP routes vertically; the two parameter families use different AllReduce groups