Add an hparam use_global_position_in_packed_sequence in mtf_transformer2.

If True (default), then we use the global position in the packed example as the input to the positional embedding.  If False, then we use the position in the individual sequence.
It is counterintuitive why we want to make True the default, since False seems to make more sense.
However, the previous submitted CL had the effect of changing from True to False, which caused some models to diverge.  This CL restores the previous working state.
TODO(noam): investigate why the models diverge with False.

PiperOrigin-RevId: 233427027

Author: nshazeer

tensor2tensor/data_generators/wiki_multi_problems.py

@@ -80,26 +80,30 @@ class LanguagemodelMultiWikiTranslatePacked1k(
   """Wiki-LM, Translation, MNLI, SQUAD mixed problem class."""
 
   def __init__(self, was_reversed=False, was_copy=False):
-    problems = [
-        # TODO(noam): uncommonet once data is generated
-        wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k(),
-        wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k(was_reversed=True),
-        translate_ende.TranslateEndeWmtMulti64kPacked1k(),
-        translate_ende.TranslateEndeWmtMulti64kPacked1k(was_reversed=True),
-        translate_enfr.TranslateEnfrWmtMulti64kPacked1k(),
-        translate_enfr.TranslateEnfrWmtMulti64kPacked1k(was_reversed=True),
-        translate_enro.TranslateEnroWmtMultiTiny64kPacked1k(),
-        translate_enro.TranslateEnroWmtMultiTiny64kPacked1k(was_reversed=True),
-        cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k(),
-        cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k(was_reversed=True),
-        multinli.MultiNLIText2textMulti64kPacked1k(),
-        squad.SquadText2textMulti64kPacked1k(),
-    ]
-    schedule = multi_problem_v2.constant_schedule(
-        multi_problem_v2.epoch_rates_to_pmf(problems))
+    problems = []
+    rates = []
+    for rate, also_reverse, cls in self.problems_and_rates:
+      for r in [False, True] if also_reverse else [False]:
+        problems.append(cls(was_reversed=r))
+        rates.append(rate)
+    pmf = multi_problem_v2.epoch_rates_to_pmf(problems, epoch_rates=rates)
+    schedule = multi_problem_v2.constant_schedule(pmf)
     super(LanguagemodelMultiWikiTranslatePacked1k, self).__init__(
         problems, schedule, was_reversed=was_reversed, was_copy=was_copy)
 
+  @property
+  def problems_and_rates(self):
+    """Returns a list of (weight, also_reverse, problem_class) triples."""
+    return [
+        (1.0, True, wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k),
+        (1.0, True, translate_ende.TranslateEndeWmtMulti64kPacked1k),
+        (1.0, True, translate_enfr.TranslateEnfrWmtMulti64kPacked1k),
+        (1.0, True, translate_enro.TranslateEnroWmtMultiTiny64kPacked1k),
+        (1.0, True, cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k),
+        (1.0, False, multinli.MultiNLIText2textMulti64kPacked1k),
+        (1.0, False, squad.SquadText2textMulti64kPacked1k),
+    ]
+
   @property
   def has_inputs(self):
     return True
@@ -117,6 +121,25 @@ def packed_length(self):
     return 1024
 
 
+@registry.register_problem
+class LanguagemodelMultiWikiTranslatePacked1kV2(
+    LanguagemodelMultiWikiTranslatePacked1k):
+  """Higher rates for rarer problems."""
+
+  @property
+  def problems_and_rates(self):
+    """Returns a list of (weight, also_reverse, problem_class) triples."""
+    return [
+        (1.0, True, wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k),
+        (3.0, True, translate_ende.TranslateEndeWmtMulti64kPacked1k),
+        (1.0, True, translate_enfr.TranslateEnfrWmtMulti64kPacked1k),
+        (100.0, True, translate_enro.TranslateEnroWmtMultiTiny64kPacked1k),
+        (1.0, True, cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k),
+        (10.0, False, multinli.MultiNLIText2textMulti64kPacked1k),
+        (10.0, False, squad.SquadText2textMulti64kPacked1k),
+    ]
+
+
 @registry.register_problem
 class LanguagemodelEnWikiLMMultiNLISubwords64k(multi_problem.MultiProblem):
   """Wiki LM and MNLI mixed problem class."""

tensor2tensor/models/mtf_transformer2.py

@@ -136,13 +136,16 @@ def import_feature(key):
       return self._import_feature(features, mesh, key)
     targets = import_feature("targets")
     sequence_id = import_feature("targets_segmentation")
-    position = import_feature("targets_position")
+    if hparams.use_global_position_in_packed_sequence:
+      position = None
+    else:
+      position = import_feature("targets_position")
     if self.autoregressive:
       inputs = mtf.shift(
           targets, offset=1, dim=self.length_dim, wrap=False)
-      if position is not None:
-        # first input in later sequences should be 0
-        inputs *= mtf.to_int32(mtf.not_equal(position, 0))
+      # We should have a 0 at the beginning of each sequence rather than the
+      # shifted EOS (1) from the previous sequence.
+      inputs -= mtf.to_int32(mtf.equal(inputs, 1))
     else:
       inputs = import_feature("inputs")
       # TODO(noam): options for bert-style masking here?
@@ -248,8 +251,12 @@ def import_feature(key):
     decoder_sequence_id = import_feature("targets_segmentation")
     if decoder_sequence_id is None:
       decoder_sequence_id = mtf.to_int32(mtf.not_equal(targets, 0))
-    encoder_position = import_feature("inputs_position")
-    decoder_position = import_feature("targets_position")
+    if hparams.use_global_position_in_packed_sequence:
+      encoder_position = None
+      decoder_position = None
+    else:
+      encoder_position = import_feature("inputs_position")
+      decoder_position = import_feature("targets_position")
     model = self.model()
     logits, loss = model.call_simple(
         inputs=inputs,
@@ -349,7 +356,7 @@ def layer_stack_from_hparams(hparams, prefix):
   """Create a layer stack based on the hyperparameter values."""
   layers = hparams.get(prefix + "layers")
   return transformer.LayerStack(
-      [layers_registry.get(l)(hparams, prefix) for l in layers],
+      [layers_registry[l](hparams, prefix) for l in layers],
       dropout_rate=hparams.layer_prepostprocess_dropout,
       norm_epsilon=hparams.norm_epsilon)
 
@@ -418,6 +425,14 @@ def mtf_transformer2_base():
       "targets": modalities.ModalityType.IDENTITY_SYMBOL,
   }
   hparams.add_hparam("beam_size", 1)
+
+  # If this is True, then in a packed dataset (where exaples are concatenated
+  # to form longer examples) we use the global position (within the concatenated
+  # sequence) to compute the positional embedding, instead of the position
+  # within the individual sequence.  This is counterintuitive, but for some
+  # reason, it keeps the model from diverging.
+  hparams.add_hparam("use_global_position_in_packed_sequence", True)
+
   return hparams
 
 
@@ -837,12 +852,3 @@ def mtr_tr_ende_deep():
   hparams.encoder_num_layers = 12
   hparams.decoder_num_layers = 12
   return hparams
-
-
-@registry.register_hparams
-def ogm_dense_0():
-  hparams = mtr_tr_dense(0)
-  hparams.max_length = 1024
-  hparams.batch_size = 128
-  hparams.shared_embedding_and_softmax_weights = True
-  return hparams

tensor2tensor/utils/trainer_lib.py

@@ -132,7 +132,8 @@ def create_session_config(log_device_placement=False,
       gpu_options=gpu_options,
       log_device_placement=log_device_placement,
       inter_op_parallelism_threads=inter_op_parallelism_threads,
-      intra_op_parallelism_threads=intra_op_parallelism_threads)
+      intra_op_parallelism_threads=intra_op_parallelism_threads,
+      isolate_session_state=True)
   return config