Massively simplify expert_utils. Breaks checkpoints for models that use experts. Fixed bug in Parallelism, where caching devices were always used, even when none. Fixed bug in attention_lm, attention_lm_moe by setting the default norm_type to "layer" instead of "none".

PiperOrigin-RevId: 164869403

Author: nshazeer

tensor2tensor/layers/common_hparams.py

@@ -69,6 +69,11 @@ def basic_params1():
       sampling_method="argmax",  # "argmax" or "random"
       problem_choice="adaptive",  # "uniform", "adaptive", "distributed"
       multiply_embedding_mode="sqrt_depth",
+      # Parameters related to mixtures of experts.
+      moe_hidden_sizes="2048",  # hidden layer sizes (comma-separated)
+      moe_num_experts=64,  # number of experts per layer
+      moe_k=2,  # how many experts to use for each batch element
+      moe_loss_coef=1e-2,
       # Sequences of operations to perform on layer input and layer output.
       # Used by common_layers.layer_preprocess, common_layers.layer_postprocess
       # Each character repsesnts an operation:
@@ -83,7 +88,7 @@ def basic_params1():
       # dropout rate to use during layer_preprocess and layer_postprocess
       layer_prepostprocess_dropout=0.1,
       # What type of normalization to use
-      norm_type="none",  # "batch", layer", "noam", "none".
+      norm_type="layer",  # "batch", layer", "noam", "none".
       # epsilon parameter to normalization function
       norm_epsilon=1e-6,
       symbol_modality_num_shards=16,

tensor2tensor/layers/common_layers.py

@@ -193,7 +193,7 @@ def embedding(x, vocab_size, dense_size, name=None, reuse=None, multiplier=1.0):
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    embedding_var = eu.ConvertGradientToTensor(embedding_var)
+    embedding_var = eu.convert_gradient_to_tensor(embedding_var)
     emb_x = tf.gather(embedding_var, x)
     if multiplier != 1.0:
       emb_x *= multiplier
@@ -823,71 +823,6 @@ def decompress_seqcnn(x,
     return tf.layers.dense(outputs, targets_vocab_size)
 
 
-def moe_layer(data_parallelism,
-              ps_devices,
-              xs,
-              train,
-              model_hidden_size,
-              expert_hidden_size,
-              n1,
-              n2,
-              loss_coef,
-              autoscale=True,
-              name=None):
-  """A mixture of experts layer.
-
-  Args:
-    data_parallelism: a expert_utils.Parallelism object.
-    ps_devices: a list of strings
-    xs: a list of input tensors.
-    train: a boolean scalar.
-    model_hidden_size: an integer (input/output size for this layer)
-    expert_hidden_size: an integer (size of each expert's hidden layer)
-    n1: an integer - number of experts (or # of groups for hierarchical MoE)
-    n2: optional integer - size of each group of experts for hierarchical MoE
-    loss_coef: a scalar - multiplier on load-balancing losses
-    autoscale: a boolean
-    name: a string
-
-  Returns:
-    ys: a list of tensors:
-    extra_training_loss: a scalar
-  """
-  dp = data_parallelism
-  with tf.variable_scope(name, default_name="moe"):
-    # Set up the hyperparameters for the gating networks.
-    primary_gating_hp = eu.NoisyTopKGatingParams()
-    primary_gating_hp.num_experts = n1
-    if n2:
-      # hierarchical MoE containing moe_n1 groups of moe_n2 experts.
-      assert n2 > 1
-      secondary_gating_hp = eu.NoisyTopKGatingParams()
-      secondary_gating_hp.num_experts = n2
-    else:
-      # flat mixture of moe_n1 experts.
-      secondary_gating_hp = None
-    # Set up the hyperparameters for the expert networks.
-    # Each expert contains a hidden RELU layer of size filter_size
-    expert_hp = eu.FeedForwardExpertParams()
-    expert_hp.autoscale = autoscale
-    expert_hp.hidden_layer_sizes = [expert_hidden_size]
-    # Create the mixture of experts.
-    moe = eu.DistributedMixtureOfExperts(primary_gating_hp, secondary_gating_hp,
-                                         expert_hp, model_hidden_size,
-                                         model_hidden_size, ps_devices, "moe")
-    # MoE expects input tensors to be 2d.
-    #  Flatten out spatial dimensions.
-    xs_2d = dp(tf.reshape, xs, [[-1, model_hidden_size]] * dp.n)
-    # Call the MoE
-    moe_out_2d, importance, load, _, _ = moe.Eval(
-        dp.devices, xs_2d, train, identifiers=None)
-    # Reshape the output to the original shape.
-    moe_out = dp(tf.reshape, moe_out_2d, dp(tf.shape, xs))
-    # These losses encourage equal load on the different experts.
-    loss = loss_coef * (eu.CVSquared(importance) + eu.CVSquared(load))
-    return moe_out, loss
-
-
 def simple_attention(target, source, bias=None):
   """A simple attention function.
 

tensor2tensor/layers/modalities.py

@@ -70,7 +70,7 @@ def _get_weights(self):
       ret = shards[0]
     else:
       ret = tf.concat(shards, 0)
-    ret = eu.ConvertGradientToTensor(ret)
+    ret = eu.convert_gradient_to_tensor(ret)
     return ret
 
   def bottom_simple(self, x, name, reuse):

tensor2tensor/models/attention_lm_moe.py

@@ -32,6 +32,7 @@
 from tensor2tensor.layers import common_attention
 from tensor2tensor.layers import common_hparams
 from tensor2tensor.layers import common_layers
+from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
 
@@ -61,6 +62,7 @@ def postprocess(x, y):
     x = dp(tf.nn.dropout, decoder_input,
            1.0 - hparams.layer_prepostprocess_dropout)
     extra_loss = 0.0
+    moe_hidden_sizes = [int(s) for s in hparams.moe_hidden_sizes.split(",")]
     for layer in xrange(hparams.num_hidden_layers):
       with tf.variable_scope("layer_%d" % layer):
         with tf.variable_scope("attention"):
@@ -78,11 +80,18 @@ def postprocess(x, y):
           x = postprocess(x, y)
         with tf.variable_scope("ffn"):
           if str(layer) in hparams.moe_layers.split(","):
-            y, loss = common_layers.moe_layer(
-                dp, self._ps_devices, preprocess(x),
+            y, loss = expert_utils.distributed_moe(
+                dp,
+                self._ps_devices,
+                preprocess(x),
                 hparams.mode == tf.contrib.learn.ModeKeys.TRAIN,
-                hparams.hidden_size, hparams.moe_hidden_size, hparams.moe_n1,
-                hparams.moe_n2, hparams.moe_loss_coef)
+                input_size=hparams.hidden_size,
+                expert_fn=expert_utils.ffn_expert_fn(
+                    hparams.hidden_size, moe_hidden_sizes,
+                    hparams.hidden_size),
+                num_experts=hparams.moe_num_experts,
+                k=hparams.moe_k,
+                loss_coef=hparams.moe_loss_coef)
             extra_loss += loss
           else:
             y = dp(
@@ -149,16 +158,7 @@ def attention_lm_moe_base():
   hparams.label_smoothing = 0.0
   hparams.shared_embedding_and_softmax_weights = int(False)
   hparams.add_hparam("filter_size", 2048)  # Add new ones like this.
-  # comma-separated list of layer numbers.
-  # At each of these layers, we replace the ffn with a mixture of experts.
-  hparams.add_hparam("moe_layers", "2")
-  # If moe_n2 is None, then use a flat MoE with moe_n1 experts.
-  # If moe_n2 is an integer, then use a hierarchical MoE
-  #   consisting of moe_n1 groups of moe_n2 experts each.
-  hparams.add_hparam("moe_n1", 32)
-  hparams.add_hparam("moe_n2", 0)
-  hparams.add_hparam("moe_hidden_size", 2048)
-  hparams.add_hparam("moe_loss_coef", 1e-2)
+  hparams.moe_num_experts = 32
   # attention-related flags
   hparams.add_hparam("num_heads", 8)
   hparams.add_hparam("attention_key_channels", 0)
@@ -168,6 +168,7 @@ def attention_lm_moe_base():
   hparams.add_hparam("attention_dropout", 0.0)
   hparams.add_hparam("relu_dropout", 0.0)
   hparams.add_hparam("pos", "timing")  # timing, none
+  hparams.add_hparam("moe_layers", "2")  # comma separated list of layer numbers
   return hparams
 
 
@@ -188,9 +189,20 @@ def attention_lm_moe_small():
   hparams.num_hidden_layers = 4
   hparams.hidden_size = 512
   hparams.filter_size = 2048
-  hparams.moe_n1 = 128
+  hparams.moe_num_experts = 128
   hparams.moe_layers = "2"
-  hparams.moe_hidden_size = 2048
+  return hparams
+
+
+@registry.register_hparams
+def attention_lm_moe_tiny():
+  """Cheap model for debugging.
+
+  Returns:
+    an hparams object.
+  """
+  hparams = attention_lm_moe_small()
+  hparams.moe_num_experts = 32
   return hparams
 
 
@@ -233,7 +245,7 @@ def attention_lm_moe_large():
   hparams.hidden_size = 1024
   hparams.num_heads = 16
   hparams.filter_size = 4096
-  hparams.moe_hidden_size = 4096
-  hparams.moe_n1 = 128
+  hparams.moe_hidden_sizes = "4096"
+  hparams.moe_num_experts = 128
   hparams.layer_prepostprocess_dropout = 0.2
   return hparams