Figure 2: Numerical precision formats and roundoff errors. Larger numbers have larger roundoff
errors. bfloat16 has up to 65,536x worse roundoff errors than float32. The router z-loss
encourages the absolute magnitude of numbers to be small, which doesn’t hinder model performance
and reduces roundoff errors. The router z-loss is most effective into functions where larger errors
can drastically change the relative output (e.g. exponential and sinusoidal functions).

4We also experimented with adding z-losses onto the attention logits which also improves model instability

without hurting model quality.

blob/master/mesh_tensorflow/

5See Mesh Tensorﬂow for implementation details: https://github.com/tensorflow/mesh/

6Matrix multiplications on TPUs perform multiplications in bfloat16 and accumulations in float32.

8
ST-MOE- DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS