Token, Meet Expert

A Mixture-of-Experts layer routes each token to a specific expert that may live on a different GPU, so each GPU must send a different, unevenly-sized chunk of tokens to every other GPU. Which collective communication primitive matches this pattern, as opposed to the all-reduce used for gradient synchronization in data parallelism?

A

Reduce-scatter

B

All-to-all

C

All-reduce

D

Broadcast