Pipeline x expert parallelism 2 PP stages / 2 EP ranks / 4 GPUs
Two pipeline stages, each implemented by a two-rank expert-parallel group Pipeline communication runs horizontally between matching expert ranks. Within each pipeline stage, expert dispatch and return run vertically between expert ranks. Backward reverses those exchanges, sends input gradients to the previous pipeline stage, and AllReduces shared parameter gradients within each stage. PP axis PP stage 0 / layers 0-5 PP stage 1 / layers 6-11 EP axis EP rank 0 EP rank 1 GPU 0 / PP 0 / EP 0 shared layers + router / replicated expert E0 expert E1 GPU 2 / PP 1 / EP 0 shared layers + router / replicated expert E0 expert E1 GPU 1 / PP 0 / EP 1 shared layers + router / replicated expert E2 expert E3 GPU 3 / PP 1 / EP 1 shared layers + router / replicated expert E2 expert E3 input shard x0 input shard x1 residual-stream send residual-stream send loss L0 loss L1 ALL-TO-ALL x2 dispatch + return ALL-TO-ALL x2 dispatch + return horizontal = pipeline sends / vertical = expert-group exchanges loss grad dL0 loss grad dL1 input-gradient send input-gradient send input grad dx0 input grad dx1 INVERSE A2A x2 expert-token gradients ALL-REDUCE shared + router grads expert gradients stay local INVERSE A2A x2 expert-token gradients ALL-REDUCE shared + router grads expert gradients stay local pipeline gradients move left; EP communication stays vertical inside each stage