Fixed some bugs for Synthesizer model.

PiperOrigin-RevId: 323770262

Author: Mesh TensorFlow Team

mesh_tensorflow/transformer/attention.py

@@ -209,9 +209,7 @@ def synthetic_attention(q,
       tf.logging.info("Using Random Synthesizers")
       r_shape = mtf.Shape([mtf.Dimension("length", max_length),
                            mtf.Dimension("heads", num_heads.size),
-                           mtf.Dimension("memory_length",
-                                         num_heads, max_length)])
-      initializer = tf.random_uniform_initializer()
+                           mtf.Dimension("memory_length", max_length)])
       r = mtf.get_variable(context.mesh, "R", r_shape,
                            initializer=None,
                            dtype=context.variable_dtype)
@@ -235,12 +233,11 @@ def synthetic_attention(q,
       r_shape = mtf.Shape([mtf.Dimension("length", 512),
                            mtf.Dimension("heads", num_heads.size),
                            mtf.Dimension("memory_length", 512)])
-      initializer = tf.random_normal_initializer()
       r1 = mtf.get_variable(context.mesh, "R1", r1_shape,
-                            initializer=initializer,
+                            initializer=None,
                             dtype=context.variable_dtype)
       r2 = mtf.get_variable(context.mesh, "R2", r2_shape,
-                            initializer=initializer,
+                            initializer=None,
                             dtype=context.variable_dtype)
       r = mtf.einsum([r1, r2], r_shape)
       r = mtf.slice(r, 0, memory_length_dim.size, memory_length_dim.name)
@@ -324,6 +321,7 @@ def synthetic_attention(q,
       outputs_shape = mtf.Shape(q.shape.dims[:-1] + [num_heads, value_dim])
   else:
     outputs_shape = q.shape - [key_dim] + value_dim
+
   outputs = mtf.einsum([weights, v], outputs_shape)
   return outputs
 

mesh_tensorflow/transformer/transformer_layers.py

@@ -438,7 +438,7 @@ def __init__(self,
       self.no_query = False
     else:
       self.shared_kv = True
-      self.shared_q = True
+      self.no_query = True
 
   def make_params(self, context):
     return attention_params(context=context,
@@ -451,7 +451,6 @@ def make_params(self, context):
   def call(self, context, x, losses=None):
     """Call the layer."""
     params = self.make_params(context)
-    q = params.compute_q(x)
     memory_length = self.memory_length(context)
     if context.mode == "incremental":
       m = x
@@ -467,26 +466,23 @@ def call(self, context, x, losses=None):
       q = x
     else:
       q = params.compute_q(x)
+    if self.shared_kv:
+      k = kv
+      v = kv
     if context.mode == "incremental":
       one_hot = mtf.one_hot(
           context.position, memory_length, dtype=context.activation_dtype)
       inv_one_hot = 1.0 - one_hot
-      if self.shared_kv:
-        old_kv = context.get_states(1)
-        kv = old_kv * inv_one_hot + kv * one_hot
-      else:
-        old_k, old_v = context.get_states(2)
-        k = old_k * inv_one_hot + k * one_hot
-        v = old_v * inv_one_hot + v * one_hot
+      old_k, old_v = context.get_states(2)
+      k = old_k * inv_one_hot + k * one_hot
+      v = old_v * inv_one_hot + v * one_hot
       memory_position = mtf.range(context.mesh, memory_length, tf.int32)
     else:
       memory_position = self.rename_length_to_memory_length(
           context.position, context)
     if context.mode == "incremental" or context.mode == "first_part":
-      context.record_new_states([kv] if self.shared_kv else [k, v])
-    if self.shared_kv:
-      k = kv
-      v = kv
+      context.record_new_states([k, v])
+
     o = attention.synthetic_attention(q, k, v, memory_length,
                                       self.kv_dim, self.kv_dim,
                                       self.compute_bias(context,