Data x expert parallelism
top-2 / 2 DP replicas / 2 EP ranks / 4 GPUs
Two data-parallel replicas, each implemented by a two-rank expert-parallel group
Four ranks receive four different input shards. Forward All-to-All exchanges stay inside each data replica. Backward uses two distinct reductions: expert gradients are AllReduced between matching expert owners across data replicas, while non-expert gradients are AllReduced across all four ranks.
DP replica 0 / EP group 0
DP replica 1 / EP group 1
EP rank 0
EP rank 1
GPU 0 / DP 0 / EP 0
shared + router
E0, E1
input shard 0
GPU 2 / DP 1 / EP 0
E0, E1
shared + router
input shard 2
GPU 1 / DP 0 / EP 1
shared + router
E2, E3
input shard 1
GPU 3 / DP 1 / EP 1
E2, E3
shared + router
input shard 3
ALL-TO-ALL
dispatch + combine
ALL-TO-ALL
dispatch + combine
DP replicas are independent
during forward
each replica performs its own top-2 expert routing
NON-EXPERT-GRADIENT ALL-REDUCE
shared + router parameters / all 4 ranks
ALL-TO-ALL
token-gradient routing
ALL-TO-ALL
token-gradient routing
EXPERT-GRADIENT
ALL-REDUCE
E0, E1 replicas
EXPERT-GRADIENT
ALL-REDUCE
E2, E3 replicas
two reducers follow the two parameter-replication patterns
forward
backward