Get bf16 activations working with mixture-of-experts.

nshazeer · Copybara-Service · commit ee89d7bb18bd · 2018-10-12T15:49:43.000-07:00
PiperOrigin-RevId: 216935230
diff --git a/tensor2tensor/mesh_tensorflow/research/moe.py b/tensor2tensor/mesh_tensorflow/research/moe.py
@@ -475,6 +475,10 @@ def _top_2_gating(
       inputs, experts_dim, use_bias=False,
       expert_dims=outer_expert_dims), experts_dim)
 
+  # The internals of this function run in float32.
+  #   bfloat16 seems to reduce quality.
+  raw_gates = mtf.to_float(raw_gates)
+
   expert_capacity_f = float(expert_capacity_dim.size)
 
   # FIND TOP 2 EXPERTS PER POSITON
@@ -591,6 +595,9 @@ def _top_2_gating(
       * mtf.one_hot(index_2, experts_dim)
       * mtf.one_hot(mtf.to_int32(position_in_expert_2), expert_capacity_dim))
 
+  combine_tensor = mtf.cast(combine_tensor, inputs.dtype)
+  loss = mtf.cast(loss, inputs.dtype)
+
   dispatch_tensor = mtf.cast(
       mtf.cast(combine_tensor, tf.bool), combine_tensor.dtype)
 
