Communication operations (allreduce and all2all) can signiﬁcantly impact sparse model
training throughput (see Table 1 for a description of the communication operations). allreduce
calls are executed along model and batch dimensions, typically dominated by the model dimension
allreduce calls that sum results of partial matrix multiplication operations from the workers.
These calls are needed when matrix multiplications are partitioned across multiple cores (e.g. model
parallelism). The gradient summation allreduce calls can be amortized away by training models
with larger batch sizes since the gradient accumulation allreduce communication cost is inde-
pendent of the batch size. To alleviate the memory issues of larger batch sizes, microbatches can
be used. Microbatches do this by splitting the batch into n evenly divisible chunks and computing
gradients on each sequentially, then summing.
To increase the allreduce throughput, more workers may need to be assigned to the model di-
ST-MOE- DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS