For large groups, allreduce in float32 even when the input/output is bfloat16. This will decrease the performance of some models, but it is numerically safer not to sequentially add large numbers of values in bfloat16. The option is controlled by an argument allreduce_in_bfloat16_max_group_size to SimdMeshImpl(). Default value is 64, meaning that allreduce with group size >64 happens in float32. Raise at your own risk (to improve performance), but check that the model quality is not affected.
For large groups, allreduce in float32 even when the input/output is bfloat16.
This will decrease the performance of some models, but it is numerically safer not to sequentially add large numbers of values in bfloat16.
The option is controlled by an argument allreduce_in_bfloat16_max_group_size to SimdMeshImpl().
Default value is 64, meaning that allreduce with group size >64 happens in float32.
Raise at your own risk (to improve performance), but check that the model quality is not affected.

PiperOrigin-RevId: 239862275

tensorflow/mesh,tensorflow/mesh