Distill large top-1 sparse models into small dense models

PiperOrigin-RevId: 325864101

Author: William Fedus

mesh_tensorflow/transformer/moe.py

@@ -345,6 +345,7 @@ def transformer_moe_layer_v1(
       group_size_dim.size,
       int((group_size_dim.size * capacity_factor) / experts_dim.size))
   expert_capacity = max(expert_capacity, 4)
+  tf.logging.info("expert_capacity: %d" % expert_capacity)
   expert_capacity_dim = mtf.Dimension("expert_capacity", expert_capacity)
 
   experts_dim_unsplit = mtf.Dimension("expert_unsplit", experts_dim.size)
@@ -704,8 +705,6 @@ def _rand_1_gating(
     inputs, outer_expert_dims, experts_dim, expert_capacity_dim,
     hparams, train, variable_dtype, importance=None, name="rand_1_gating"):
   """Compute a random top-1 gating."""
-  del importance
-
   # SELECT EXPERT
   if train:
     policy = hparams.moe_rand_1_policy_train
@@ -724,6 +723,10 @@ def _rand_1_gating(
                                  name=name)
   raw_gates = mtf.softmax(gate_logits, reduced_dim=experts_dim)
 
+  # The internals of this function run in float32.
+  #   bfloat16 seems to reduce quality.
+  raw_gates = mtf.to_float(raw_gates)
+
   if policy == "argmax" or policy == "input_dropout" or policy == "input_jitter":
     expert_gate, expert_index = mtf.top_1(raw_gates, reduced_dim=experts_dim)
   elif policy == "sample":
@@ -740,8 +743,14 @@ def _rand_1_gating(
   group_size_dim = inputs.shape[-2]
   density_1 = mtf.reduce_mean(expert_mask, reduced_dim=group_size_dim)
   density_1_proxy = mtf.reduce_mean(raw_gates, reduced_dim=group_size_dim)
-  loss = (mtf.reduce_mean(density_1_proxy * density_1)
-          * float(experts_dim.size * experts_dim.size))
+  if importance is not None:
+    expert_mask *= mtf.cast(mtf.equal(importance, 1.0), dtype=raw_gates.dtype)
+    expert_gate *= mtf.cast(mtf.equal(importance, 1.0), dtype=raw_gates.dtype)
+    density_1_proxy *= mtf.cast(
+        mtf.equal(importance, 1.0), dtype=raw_gates.dtype)
+  loss = (
+      mtf.reduce_mean(density_1_proxy * density_1) *
+      float(experts_dim.size * experts_dim.size))
 
   # Logging
   if train:
@@ -767,20 +776,25 @@ def _rand_1_gating(
   # the batch indices, to each expert, with position_in_expert
   position_in_expert = mtf.cumsum(
       expert_mask, group_size_dim, exclusive=True) * expert_mask
+  position_in_expert = mtf.cast(position_in_expert, dtype=raw_gates.dtype)
   # Keep only tokens that fit within expert_capacity.
   expert_capacity_float = float(expert_capacity_dim.size)
-  expert_mask *= mtf.to_float(mtf.less(position_in_expert,
-                                       expert_capacity_float))
+  expert_mask *= mtf.cast(
+      mtf.less(position_in_expert, expert_capacity_float),
+      dtype=raw_gates.dtype)
   expert_mask_flat = mtf.reduce_sum(expert_mask, reduced_dim=experts_dim)
 
   # Mask out the experts that have overflowed expert capacity. Sparsify the
   # expert_gate.
   expert_gate *= expert_mask_flat
 
   combine_tensor = (
-      expert_gate * expert_mask_flat
-      * mtf.one_hot(expert_index, experts_dim)
-      * mtf.one_hot(mtf.to_int32(position_in_expert), expert_capacity_dim))
+      expert_gate * expert_mask_flat *
+      mtf.one_hot(expert_index, experts_dim, dtype=raw_gates.dtype) *
+      mtf.one_hot(
+          mtf.to_int32(position_in_expert),
+          expert_capacity_dim,
+          dtype=raw_gates.dtype))
 
   # Match the inputs dtype.
   combine_tensor = mtf.cast(combine_tensor, inputs.dtype)