As in most modern distributed Transformers we train with mixed precision (Micikevicius et al.,
2017) 5. Weights are stored in float32 for gradient updates and then converted to bfloat16
when doing matrix multiplications in the forward and backward pass6. Furthermore, all activations
are stored and operated on in bfloat16 and allreduce communications can be done in either
bfloat16 or float32 numerical precision. For the largest model explored in this work (ST-
MoE-32B presented later) we ﬁnd speed-ups halving the numerical precision of the allreduce,
however this also can destabilize the training so we keep this as float32 throughout this work.
A lower precision format enables more efﬁcient models by reducing (a) communication costs be-
tween processors and memory, (b) computation costs, (c) memory for storing tensors (e.g. activa-
tions). However, lower precision formats come at the expense of larger roundoff errors which can
lead to irrecoverable training instabilities.
ST-MOE- DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS