Set allreduce_in_bfloat16_max_group_size to 2 to avoid OOM during inference/eval.

adarob · Mesh TensorFlow Team · commit d46ff8751f38 · 2020-07-23T06:46:54.000-07:00
PiperOrigin-RevId: 322776425
diff --git a/mesh_tensorflow/simd_mesh_impl.py b/mesh_tensorflow/simd_mesh_impl.py
@@ -41,7 +41,7 @@ def __init__(self,
                devices=None,
                device_assignment=None,
                logical_to_physical=None,
-               allreduce_in_bfloat16_max_group_size=0):
+               allreduce_in_bfloat16_max_group_size=8):
     """Create a SimdMeshImpl.
 
     Args:
