Make Reformer config smaller and allow to use a single rng for all steps.

PiperOrigin-RevId: 268567430

Author: Lukasz Kaiser

tensor2tensor/trax/configs/transformer_revnet_imagenet64_8gb.gin

@@ -5,8 +5,8 @@ import tensor2tensor.trax.trax
 
 # Parameters for batch_fun:
 # ==============================================================================
-batch_fun.batch_size_per_device = 8
-batch_fun.eval_batch_size = 128
+batch_fun.batch_size_per_device = 2
+batch_fun.eval_batch_size = 16
 batch_fun.max_eval_length = 12288  # 64 * 64 * 3
 
 # Parameters for inputs:
@@ -41,24 +41,33 @@ DotProductCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.dropout = 0.0
 MemoryEfficientCausalAttention.loop_stride = 512
 
-# Parameters for DummyHashedAttention:
+# Parameters for MergedHashedCausalAttention:
 # ==============================================================================
-# DummyHashedAttention.dropout = 0.0
-# DummyHashedAttention.n_bins = 64
+MergedHashedCausalAttention.dropout = 0.0
+MergedHashedCausalAttention.n_bins = 32
+MergedHashedCausalAttention.bin_by_time = True
+MergedHashedCausalAttention.one_rng = False
+
+# Parameters for MergedMultiHashedCausalAttention:
+# ==============================================================================
+MergedMultiHashedCausalAttention.dropout = 0.0
+MergedMultiHashedCausalAttention.n_bins = 64
+MergedMultiHashedCausalAttention.n_hashes = 4
+MergedMultiHashedCausalAttention.bin_by_time = False
+MergedHashedCausalAttention.one_rng = True
 
 # Parameters for TransformerRevnetLM:
 # ==============================================================================
 TransformerRevnetLM.d_model = 1024
 TransformerRevnetLM.d_ff = 2048
-TransformerRevnetLM.d_attention_key = 32
-TransformerRevnetLM.d_attention_value = 32
+TransformerRevnetLM.d_attention_key = 64
+TransformerRevnetLM.d_attention_value = 64
 TransformerRevnetLM.dropout = 0.0
 TransformerRevnetLM.max_len = 12288  # 64 * 64 * 3
 TransformerRevnetLM.mode = 'train'
 TransformerRevnetLM.n_heads = 4
-TransformerRevnetLM.n_layers = 6
+TransformerRevnetLM.n_layers = 4
 TransformerRevnetLM.vocab_size = 256
 TransformerRevnetLM.n_chunks = 16
 TransformerRevnetLM.n_attention_chunks = 1
-TransformerRevnetLM.attention_type = @trax.layers.MemoryEfficientCausalAttention
-
+TransformerRevnetLM.attention_type = @trax.layers.MergedMultiHashedCausalAttention

tensor2tensor/trax/layers/attention.py

@@ -18,6 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import random
 import jax
 import numpy as onp
 
@@ -559,11 +560,17 @@ def body_fun(vals):  # pylint: disable=invalid-name
 class MergedHashedCausalAttention(BaseCausalAttention):
   """Hash-based causal attention."""
 
-  def __init__(self, dropout, mode, n_bins=64, bin_by_time=False):
+  def __init__(self, dropout, mode, n_bins=64,
+               bin_by_time=False, one_rng=False):
     del dropout, mode
     super(MergedHashedCausalAttention, self).__init__()
     self.n_bins = n_bins
     self.bin_by_time = bin_by_time
+    seed = random.randint(0, 2**31 - 1)
+    self._one_rng = one_rng
+    self._prng = None
+    if one_rng:
+      self._prng = backend.random.get_prng(seed)
 
   def call(self, inputs, params=(), state=(), **kwargs):
     del params
@@ -604,8 +611,12 @@ def hash_vectors(self, vecs, rng):
     # It's not clear whether sampling a different random rotation for each head
     # and batch element matters here, but see MergedMultiHashedCausalAttention.
     assert self.n_bins % 2 == 0
+    rot_rng = rng
+    if self._one_rng:
+      rot_rng = jax.lax.tie_in(vecs, self._prng)
     random_rotation = jax.random.normal(
-        rng, (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
+        rot_rng,
+        (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
     vecs = self.make_unit_length(vecs)
@@ -735,12 +746,18 @@ def binned_attn_vjp(sqk, sv, so_ct):  # pylint: disable=invalid-name
 class MergedMultiHashedCausalAttention(BaseCausalAttention):
   """Hash-based causal attention, with multiple hashes."""
 
-  def __init__(self, dropout, mode, n_bins=64, n_hashes=1, bin_by_time=False):
+  def __init__(self, dropout, mode, n_bins=64, n_hashes=1,
+               bin_by_time=False, one_rng=False):
     del dropout, mode
     super(MergedMultiHashedCausalAttention, self).__init__()
     self.n_bins = n_bins
     self.n_hashes = n_hashes
     self.bin_by_time = bin_by_time
+    seed = random.randint(0, 2**31 - 1)
+    self._one_rng = one_rng
+    self._prng = None
+    if one_rng:
+      self._prng = backend.random.get_prng(seed)
 
   def bin_vectors_by_time(self, vecs):
     seqlen = vecs.shape[-2]
@@ -770,8 +787,12 @@ def hash_vectors(self, vecs, rng):
     # of vecs. Applying multiple hashes to the same input is important because
     # it increases the probability of being in the same bin as relevant items.
     assert self.n_bins % 2 == 0
+    rot_rng = rng
+    if self._one_rng:
+      rot_rng = jax.lax.tie_in(vecs, self._prng)
     random_rotation = jax.random.normal(
-        rng, (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
+        rot_rng,
+        (vecs.shape[0], vecs.shape[-1], self.n_bins//2)).astype('float32')
 
     # TODO(kitaev): making the vectors unit-length here is probably redundant.
     vecs = self.make_unit_length(vecs)