Splitting tokens when routing

PiperOrigin-RevId: 378316002

Author: Mesh TensorFlow Team

mesh_tensorflow/transformer/moe.py

@@ -61,7 +61,9 @@ def __init__(self,
                ntlb_top_k=4,
                output_dim=None,
                use_experts_attention=False,
-               z_loss=None):
+               z_loss=None,
+               num_hidden_splits=None,
+               split_hidden_before_routing=False):
     self._hparams = HParams(
         moe_gating=moe_gating,
         moe_num_experts=num_experts,
@@ -85,7 +87,9 @@ def __init__(self,
         moe_output_dim=output_dim,
         moe_ntlb_top_k=ntlb_top_k,
         moe_use_experts_attention=use_experts_attention,
-        moe_z_loss=z_loss)
+        moe_z_loss=z_loss,
+        moe_num_hidden_splits=num_hidden_splits,
+        moe_split_hidden_before_routing=split_hidden_before_routing)
     self._activation = activation
 
   def call(self, context, x, losses=None):
@@ -327,8 +331,8 @@ def transformer_moe_layer_v1(
   # We "cheat" here and look at the mesh shape and layout. This is to ensure
   # that the number of groups is a multiple of the mesh dimension
   # over which those groups are split.
-  batch_and_length_dims, input_dim = (orig_inputs.shape.dims[:-1],
-                                      orig_inputs.shape.dims[-1])
+  batch_and_length_dims, orig_input_dim = (
+      orig_inputs.shape.dims[:-1], orig_inputs.shape.dims[-1])
   # Hack: we assume that
   #   "outer_batch" == replication of experts
   #   mesh_dim_size can be derived from mesh_shape and orig_batch_dim
@@ -348,16 +352,57 @@ def transformer_moe_layer_v1(
 
   n = n // outer_batch_dim.size
 
-  mesh_dim_size = mtf.tensor_dim_to_mesh_dim_size(layout, mesh_shape,
-                                                  orig_batch_dim)
-  num_groups, group_size = _split_into_groups(n, hparams.moe_group_size,
-                                              mesh_dim_size)
+  # Create num_groups and group_size dimensions
+  mesh_dim_size = mtf.tensor_dim_to_mesh_dim_size(
+      layout, mesh_shape, orig_batch_dim)
+  num_groups, group_size = _split_into_groups(
+      n, hparams.moe_group_size, mesh_dim_size)
+  orig_group_size_dim = mtf.Dimension("group", group_size)
+  orig_num_groups_dim = mtf.Dimension(orig_batch_dim.name, num_groups)
+
+  # The original dimensions correspond to those before splitting tokens
+  # into subtokens
+  group_size_dim = orig_group_size_dim
+  num_groups_dim = orig_num_groups_dim
+  input_dim = orig_input_dim
+
+  split_hidden_before_routing = False
+  split_hidden_after_routing = False
+  if hparams.moe_num_hidden_splits is not None:
+    if orig_input_dim.size % hparams.moe_num_hidden_splits:
+      raise ValueError("num_hidden_splits {} must divide input_dim {}".format(
+          hparams.moe_num_hidden_splits, input_dim.size))
+    if output_dim.size % hparams.moe_num_hidden_splits:
+      raise ValueError("num_hidden_splits {} must divide input_dim {}".format(
+          hparams.moe_num_hidden_splits, input_dim.size))
+    split_hidden_before_routing = hparams.moe_split_hidden_before_routing
+    split_hidden_after_routing = not hparams.moe_split_hidden_before_routing
+    hidden_dim = mtf.Dimension(
+        "expert_hidden",
+        hparams.moe_hidden_size // hparams.moe_num_hidden_splits)
+    sub_output_dim = mtf.Dimension(
+        output_dim.name, output_dim.size // hparams.moe_num_hidden_splits)
+    num_splits_dim = mtf.Dimension(
+        "num_splits", hparams.moe_num_hidden_splits)
+
+  if split_hidden_before_routing:
+    input_dim = mtf.Dimension(
+        input_dim.name, input_dim.size // hparams.moe_num_hidden_splits)
+
+    # Split into groups and subtokens
+    inputs = mtf.reshape(
+        inputs, [outer_batch_dim, num_groups_dim, group_size_dim,
+                 num_splits_dim, input_dim])
 
-  group_size_dim = mtf.Dimension("group", group_size)
-  num_groups_dim = mtf.Dimension(orig_batch_dim.name, num_groups)
+    inputs = mtf.transpose(
+        inputs, [outer_batch_dim, num_groups_dim, num_splits_dim,
+                 group_size_dim, input_dim])
 
+    num_groups_dim = mtf.Dimension(
+        orig_batch_dim.name, num_groups * hparams.moe_num_hidden_splits)
+
+  # [outer_batch_dim, num_groups_dim.B, group_size_dim, input_dim]
   moe_input_dims = [outer_batch_dim, num_groups_dim, group_size_dim, input_dim]
-  # OGSM Tensor
   inputs = mtf.reshape(inputs, moe_input_dims)
 
   # Each sequence sends expert_capacity positions to each expert.
@@ -373,156 +418,138 @@ def transformer_moe_layer_v1(
   expert_capacity_dim = mtf.Dimension("expert_capacity", expert_capacity)
   experts_dim_unsplit = mtf.Dimension("expert_unsplit", experts_dim.size)
   batch_dim_unsplit = mtf.Dimension("batch_unsplit", num_groups_dim.size)
+
   if nonpadding is not None:
     nonpadding = mtf.zeros(
         inputs.mesh, batch_and_length_dims, dtype=inputs.dtype) + nonpadding
+
+    if split_hidden_before_routing:
+      nonpadding = mtf.reshape(
+          nonpadding,
+          [outer_batch_dim, orig_num_groups_dim, orig_group_size_dim])
+
+      # Tile num_hidden_splits times with an einsum
+      tiling_tensor = mtf.ones(inputs.mesh, [num_splits_dim])
+      nonpadding = mtf.einsum(
+          [nonpadding, tiling_tensor],
+          output_shape=[outer_batch_dim, orig_num_groups_dim, num_splits_dim,
+                        orig_group_size_dim])
+
     nonpadding = mtf.reshape(nonpadding, moe_input_dims[:-1])
-  if hparams.moe_gating == "top_2":
-    # combine_tensor,
-    # dispatch_tensor  OG`SEC Tensors
-    # (G is generally split along mesh dim)
-    dispatch_tensor, combine_tensor, loss = _top_2_gating(
-        inputs=inputs,
-        outer_expert_dims=None,
-        experts_dim=experts_dim_unsplit,
-        expert_capacity_dim=expert_capacity_dim,
-        hparams=hparams,
-        train=train,
-        variable_dtype=variable_dtype,
-        importance=nonpadding,
-        num_microbatches=num_microbatches)
-  elif hparams.moe_gating == "switch":
-    dispatch_tensor, combine_tensor, loss = _switch_gating(
-        inputs=inputs,
-        outer_expert_dims=None,
-        experts_dim=experts_dim_unsplit,
-        expert_capacity_dim=expert_capacity_dim,
-        hparams=hparams,
-        train=train,
-        variable_dtype=variable_dtype,
-        importance=nonpadding,
-        num_microbatches=num_microbatches)
-  elif hparams.moe_gating == "ntlb":
-    dispatch_tensor, combine_tensor, loss = _ntlb_gating(
-        inputs=inputs,
-        outer_expert_dims=None,
-        experts_dim=experts_dim_unsplit,
-        expert_capacity_dim=expert_capacity_dim,
-        hparams=hparams,
-        train=train,
-        variable_dtype=variable_dtype,
-        importance=nonpadding,
-        num_microbatches=num_microbatches)
-  elif hparams.moe_gating == "switch_max":
-    dispatch_tensor, combine_tensor, loss = _switch_max_gating(
-        inputs=inputs,
-        outer_expert_dims=None,
-        experts_dim=experts_dim_unsplit,
-        expert_capacity_dim=expert_capacity_dim,
-        hparams=hparams,
-        train=train,
-        variable_dtype=variable_dtype,
-        importance=nonpadding,
-        num_microbatches=num_microbatches)
-  elif hparams.moe_gating == "expert_selection":
-    dispatch_tensor, combine_tensor, loss = _expert_selection_gating(
-        inputs=inputs,
-        outer_expert_dims=None,
-        experts_dim=experts_dim_unsplit,
-        group_size_dim=group_size_dim,
-        expert_capacity_dim=expert_capacity_dim,
-        hparams=hparams,
-        train=train,
-        variable_dtype=variable_dtype,
-        importance=nonpadding,
-        name="expert_selection_gating",
-        num_microbatches=num_microbatches)
-  else:
-    raise ValueError("unknown hparams.moe_gating=%s" % hparams.moe_gating)
 
-  expert_inputs = mtf.einsum([inputs, dispatch_tensor],
-                             mtf.Shape([
-                                 outer_batch_dim, experts_dim_unsplit,
-                                 num_groups_dim, expert_capacity_dim, input_dim
-                             ]))
+  # [outer_batch_dim, num_groups_dim.B, group_size_dim,
+  #  experts_dim_unsplit, expert_capacity_dim]
+  gating_fn = get_gating_fn(hparams.moe_gating)
+  dispatch_tensor, combine_tensor, loss = gating_fn(
+      inputs=inputs,
+      outer_expert_dims=None,
+      experts_dim=experts_dim_unsplit,
+      expert_capacity_dim=expert_capacity_dim,
+      hparams=hparams,
+      train=train,
+      variable_dtype=variable_dtype,
+      importance=nonpadding,
+      num_microbatches=num_microbatches)
+
+  # Dispatch to the experts by reducing group_size_dim
+  # inputs: [outer_batch_dim, num_groups_dim.B, group_size_dim, input_dim]
+  # dispatch_tensor: [outer_batch_dim, num_groups_dim.B, group_size_dim,
+  #  experts_dim_unsplit, expert_capacity_dim]
+  # expert_inputs: [outer_batch_dim, experts_dim_unsplit, num_groups_dim.B,
+  #  expert_capacity_dim, input_dim]
+  expert_inputs_shape = [
+      outer_batch_dim, experts_dim_unsplit, num_groups_dim,
+      expert_capacity_dim, input_dim]
+  expert_inputs = mtf.einsum([inputs, dispatch_tensor], expert_inputs_shape)
 
+  # Split over batch -> split over experts
   # Extra reshape reduces communication cost for model-parallel versions.
   # For model-parallel versions, this reshape causes an mtf.slice and for non-
   # model-parallel versions, this has no effect.
+  # expert_inputs: [outer_batch_dim, experts_dim.B, batch_dim_unsplit,
+  #  expert_capacity_dim, input_dim or input_dim.M]
   d_model_split_dim = mtf.Dimension("d_model_split", input_dim.size)
-  expert_inputs = mtf.reshape(
-      expert_inputs,
-      mtf.Shape([
-          outer_batch_dim, experts_dim, batch_dim_unsplit, expert_capacity_dim,
-          d_model_split_dim
-      ]))
-
-  # Split over batch -> split over experts
-  expert_inputs = mtf.reshape(
-      expert_inputs,
-      mtf.Shape([
-          outer_batch_dim, experts_dim, batch_dim_unsplit, expert_capacity_dim,
-          input_dim
-      ]))
-
-  # Now feed the expert inputs through the experts.
-  h = mtf.layers.dense_product(
-      expert_inputs,
-      reduced_dims=expert_inputs.shape.dims[-1:],
-      new_dims=[hidden_dim],
-      expert_dims=[experts_dim],
-      activation_functions=activation, use_bias=False,
-      variable_dtype=variable_dtype, name="wi")
-
-  if hparams.moe_dropout_rate != 0.0:
-    h = mtf.dropout(h, is_training=train,
-                    keep_prob=1.0 - hparams.moe_dropout_rate)
-
-  def _compute_output(hidden, layer_name):
-    """Compute the output of the attention layer from the hidden vector."""
+  expert_inputs_shape = [
+      outer_batch_dim, experts_dim, batch_dim_unsplit,
+      expert_capacity_dim, d_model_split_dim]
+  expert_inputs = mtf.reshape(expert_inputs, expert_inputs_shape)
+
+  expert_inputs_shape = [
+      outer_batch_dim, experts_dim, batch_dim_unsplit,
+      expert_capacity_dim, input_dim]
+  expert_inputs = mtf.reshape(expert_inputs, expert_inputs_shape)
+
+  def _apply_experts(x, output_dim, hidden_dim):
+    # x: [outer_batch_dim, experts_dim.B, batch_dim_unsplit,
+    # expert_capacity_dim, input_dim]
+    h = mtf.layers.dense_product(
+        x,
+        reduced_dims=x.shape.dims[-1:],
+        new_dims=[hidden_dim],
+        expert_dims=[experts_dim],
+        activation_functions=activation, use_bias=False,
+        variable_dtype=variable_dtype, name="wi")
+
+    if hparams.moe_dropout_rate != 0.0:
+      h = mtf.dropout(h, is_training=train,
+                      keep_prob=1.0 - hparams.moe_dropout_rate)
     expert_output = mtf.layers.dense(
-        hidden, output_dim, expert_dims=[experts_dim], use_bias=False,
-        reduced_dims=hidden.shape.dims[-1:], variable_dtype=variable_dtype,
-        name=layer_name)
-
-    # Extra reshape reduces communication cost for model-parallel versions.
-    # For model-parallel versions, this reshape causes an mtf.slice and for non-
-    # model-parallel versions, this has no effect.
-    expert_output = mtf.reshape(
-        expert_output,
-        mtf.Shape([
-            outer_batch_dim, experts_dim_unsplit, num_groups_dim,
-            expert_capacity_dim, d_model_split_dim
-        ]))
-
-    # Split over experts -> split over batch
+        h, output_dim, expert_dims=[experts_dim], use_bias=False,
+        reduced_dims=h.shape.dims[-1:], variable_dtype=variable_dtype,
+        name="wo")
+
+    return expert_output
+
+  if split_hidden_after_routing:
+    input_dim = mtf.Dimension(
+        input_dim.name, input_dim.size // hparams.moe_num_hidden_splits)
+    expert_inputs = mtf.reshape(
+        expert_inputs, expert_inputs.shape[:-1] + [num_splits_dim, input_dim])
+    expert_output = _apply_experts(expert_inputs, sub_output_dim, hidden_dim)
+    # Concat sub_tokens into tokens
     expert_output = mtf.reshape(
-        expert_output,
-        mtf.Shape([
-            outer_batch_dim,
-            experts_dim_unsplit,
-            num_groups_dim,
-            expert_capacity_dim,
-            output_dim,
-        ]))
-    moe_output_dims = moe_input_dims[:-1] + [output_dim]
-    output = mtf.einsum([expert_output, combine_tensor],
-                        mtf.Shape(moe_output_dims))
-    output = mtf.reshape(output, batch_and_length_dims + [output_dim])
-    return output
-
-  if hparams.moe_use_experts_attention:
-    # We share k_h and v_h with no degradation in performance
-    q_h, k_h = h, h
-    outputs = []
-    q = _compute_output(q_h, layer_name="q_wo")
-    k = _compute_output(k_h, layer_name="k_wo")
-    outputs.append(q)
-    outputs.append(k)
-    return outputs, loss * hparams.moe_loss_coef
+        expert_output, expert_output.shape[:-2] + [output_dim])
+  elif split_hidden_before_routing:
+    expert_output = _apply_experts(expert_inputs, sub_output_dim, hidden_dim)
   else:
-    output = _compute_output(h, layer_name="wo")
-    return output, loss * hparams.moe_loss_coef
+    expert_output = _apply_experts(expert_inputs, output_dim, hidden_dim)
+
+  # Extra reshape reduces communication cost for model-parallel versions.
+  # For model-parallel versions, this reshape causes an mtf.slice and for non-
+  # model-parallel versions, this has no effect.
+  expert_output_shape = [
+      outer_batch_dim, experts_dim_unsplit, num_groups_dim,
+      expert_capacity_dim, d_model_split_dim]
+  expert_output = mtf.reshape(expert_output, expert_output_shape)
+
+  # Split over experts -> split over batch
+  expert_output_shape = [
+      outer_batch_dim, experts_dim_unsplit, num_groups_dim,
+      expert_capacity_dim, expert_output.shape[-1]]
+  expert_output = mtf.reshape(expert_output, expert_output_shape)
+
+  # Combine by reducing experts_dim_unsplit and expert_capacity_dim
+  # expert_output: [outer_batch_dim, experts_dim_unsplit, num_groups_dim,
+  #   expert_capacity_dim, output_dim]
+  # combine_tensor: [outer_batch_dim, num_groups_dim.B, group_size_dim,
+  #  experts_dim_unsplit, expert_capacity_dim]
+  # output: [outer_batch_dim, num_groups_dim.B, group_size_dim, input_dim]
+  moe_output_dims = moe_input_dims[:-1] + [expert_output.shape[-1]]
+  output = mtf.einsum([expert_output, combine_tensor], moe_output_dims)
+  # import pdb; pdb.set_trace()  # pylint:disable=g-import-not-at-top
+
+  if split_hidden_before_routing:
+    output = mtf.reshape(
+        output, [output.shape[0], orig_num_groups_dim, num_splits_dim] + (
+            output.shape[-2:]))
+    output = mtf.transpose(
+        output, output.shape[:2] + [
+            group_size_dim, num_splits_dim, output.shape[-1]])
+    output = mtf.reshape(output, output.shape[:3] + [output_dim])
+
+  output = mtf.reshape(output, batch_and_length_dims + [output_dim])
+
+  return output, loss * hparams.moe_loss_coef
 
 
 def transformer_moe_layer_v2(
@@ -801,6 +828,22 @@ def transformer_moe_layer_v2(
   return output, (loss_outer + loss_inner) * hparams.moe_loss_coef
 
 
+def get_gating_fn(moe_gating):
+  """Factory for gating functions."""
+  if moe_gating == "top_2":
+    return _top_2_gating
+  elif moe_gating == "switch":
+    return _switch_gating
+  elif moe_gating == "ntlb":
+    return _ntlb_gating
+  elif moe_gating == "switch_max":
+    return _switch_max_gating
+  elif moe_gating == "expert_selection":
+    return _expert_selection_gating
+  else:
+    raise ValueError("unknown hparams.moe_gating=%s" % moe_gating)
+
+
 def _ntlb_gating(inputs,
                  outer_expert_dims,
                  experts_dim,