tf-models-nightly 2.18.0.dev20241021__py2.py3-none-any.whl → 2.18.0.dev20241022__py2.py3-none-any.whl

official/nlp/modeling/layers/block_sparse_attention.py

@@ -48,6 +48,7 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
       tgt_block_size=None,
       use_sigmoid_attn=False,
       sigmoid_attn_bias=None,
+      num_kv_heads=None,
       **kwargs
   ):
     """Initializes the block sparse attention layer.
@@ -61,6 +62,8 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
       use_sigmoid_attn: If enabled, uses sigmoid instead of softmax to compute
         attn probs. https://arxiv.org/pdf/2409.04431
       sigmoid_attn_bias: Bias for sigmoid attn. Suggested value -ln(seq_len).
+      num_kv_heads: Number of key/value heads in the multi-head self attention.
+        Refer to multi_query_attention.py for more details.
       **kwargs: Args passed to the base class.
     """
     super().__init__(**kwargs)
@@ -68,6 +71,11 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
       raise ValueError("src_block_size must be specified.")
     self._src_block_size = src_block_size
     self._tgt_block_size = tgt_block_size or self._src_block_size
+    self._num_kv_heads = num_kv_heads
+    if num_kv_heads is not None and num_kv_heads != 1:
+      raise ValueError(
+          "num_kv_heads must be 1. Grouped-query attention is not supported."
+      )
     self._use_sigmoid_attn = use_sigmoid_attn
     self._sigmoid_attn_bias = sigmoid_attn_bias
     if self._use_sigmoid_attn:
@@ -117,22 +125,50 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
           name="query",
           **self._get_common_kwargs_for_sublayer(),
       )
-      self._key_dense = tf_keras.layers.EinsumDense(
-          proj_einsum_eqn,
-          output_shape=qk_output_shape,
-          bias_axes=bias_axes if self._use_bias else None,
-          name="key",
-          **self._get_common_kwargs_for_sublayer(),
-      )
-      self._value_dense = tf_keras.layers.EinsumDense(
-          proj_einsum_eqn,
-          output_shape=v_output_shape,
-          bias_axes=bias_axes if self._use_bias else None,
-          name="value",
-          **self._get_common_kwargs_for_sublayer(),
-      )
-      self._dot_product_equation = "BNLsH,BNLtH->BNLts"
-      self._combine_equation = "BNLts,BNLsH->BNLtH"
+      if self._num_kv_heads == 1:
+        self._key_dense = tf_keras.layers.EinsumDense(
+            "BTD,DH->BTH",
+            output_shape=[None, self._key_dim],
+            bias_axes="H" if self._use_bias else None,
+            name="key",
+            **self._get_common_kwargs_for_sublayer(),
+        )
+        self._value_dense = tf_keras.layers.EinsumDense(
+            "BTD,DH->BTH",
+            output_shape=[None, self._value_dim],
+            bias_axes="H" if self._use_bias else None,
+            name="value",
+            **self._get_common_kwargs_for_sublayer(),
+        )
+      else:
+        self._key_dense = tf_keras.layers.EinsumDense(
+            proj_einsum_eqn,
+            output_shape=qk_output_shape,
+            bias_axes=bias_axes if self._use_bias else None,
+            name="key",
+            **self._get_common_kwargs_for_sublayer(),
+        )
+        self._value_dense = tf_keras.layers.EinsumDense(
+            proj_einsum_eqn,
+            output_shape=v_output_shape,
+            bias_axes=bias_axes if self._use_bias else None,
+            name="value",
+            **self._get_common_kwargs_for_sublayer(),
+        )
+      if self._key_shape[-2] == self._tgt_block_size:
+        if self._num_kv_heads == 1:
+          self._dot_product_equation = "BsH,BNLtH->BNLts"
+          self._combine_equation = "BNLts,BsH->BNLtH"
+        else:
+          self._dot_product_equation = "BNsH,BNLtH->BNLts"
+          self._combine_equation = "BNLts,BNsH->BNLtH"
+      else:
+        if self._num_kv_heads == 1:
+          self._dot_product_equation = "BLsH,BNLtH->BNLts"
+          self._combine_equation = "BNLts,BLsH->BNLtH"
+        else:
+          self._dot_product_equation = "BNLsH,BNLtH->BNLts"
+          self._combine_equation = "BNLts,BNLsH->BNLtH"
       if self._output_shape:
         if not isinstance(self._output_shape, collections.abc.Sized):
           output_shape = [self._output_shape]
@@ -153,17 +189,25 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
     """Converts the attention mask to block diagonal."""
     # Uses the same key mask for the entire query sequence since softmax
     # is applied only on the key axis.
-    attention_mask = tf.cast(attention_mask[:, 0, :], dtype=dtype)
     tgt_num_blocks = self._key_shape[-2] // self._tgt_block_size
-    attention_mask = tf.reshape(
-        attention_mask,
-        [
-            -1,
-            tgt_num_blocks,
-            self._tgt_block_size,
-        ],
-    )
-    return tf.einsum("BLQ,BLK->BLQK", attention_mask, attention_mask)
+    if tgt_num_blocks == 1:
+      src_num_blocks = self._query_shape[-2] // self._src_block_size
+      result = tf.reshape(
+          attention_mask,
+          [-1, src_num_blocks, self._src_block_size, self._tgt_block_size],
+      )
+    else:
+      attention_mask = tf.cast(attention_mask[:, 0, :], dtype=dtype)
+      attention_mask = tf.reshape(
+          attention_mask,
+          [
+              -1,
+              tgt_num_blocks,
+              self._tgt_block_size,
+          ],
+      )
+      result = tf.einsum("BLQ,BLK->BLQK", attention_mask, attention_mask)
+    return result
 
   def _masked_softmax(self, attention_scores, attention_mask=None):
     # Normalize the attention scores to probabilities.
@@ -217,7 +261,7 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
     src_num_blocks = self._query_shape[-2] // self._src_block_size
     tgt_num_blocks = self._key_shape[-2] // self._tgt_block_size
 
-    if src_num_blocks != tgt_num_blocks:
+    if src_num_blocks != tgt_num_blocks and tgt_num_blocks != 1:
       raise ValueError(
           "src_num_blocks must be equal to tgt_num_blocks."
       )
@@ -230,20 +274,37 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
         self._src_block_size,
         self._key_dim,
     ])
-    key_blocks = tf.reshape(key, [
-        -1,
-        self._num_heads,
-        tgt_num_blocks,
-        self._tgt_block_size,
-        self._key_dim,
-    ])
-    value_blocks = tf.reshape(value, [
-        -1,
-        self._num_heads,
-        tgt_num_blocks,
-        self._tgt_block_size,
-        self._value_dim,
-    ])
+    if tgt_num_blocks != 1 and self._num_kv_heads != 1:
+      key_blocks = tf.reshape(key, [
+          -1,
+          self._num_heads,
+          tgt_num_blocks,
+          self._tgt_block_size,
+          self._key_dim,
+      ])
+      value_blocks = tf.reshape(value, [
+          -1,
+          self._num_heads,
+          tgt_num_blocks,
+          self._tgt_block_size,
+          self._value_dim,
+      ])
+    elif tgt_num_blocks != 1 and self._num_kv_heads == 1:
+      key_blocks = tf.reshape(key, [
+          -1,
+          tgt_num_blocks,
+          self._tgt_block_size,
+          self._key_dim,
+      ])
+      value_blocks = tf.reshape(value, [
+          -1,
+          tgt_num_blocks,
+          self._tgt_block_size,
+          self._value_dim,
+      ])
+    else:
+      key_blocks = key
+      value_blocks = value
     if attention_mask is not None:
       attention_mask = self._block_diagonal_mask(attention_mask, key.dtype)
     # pytype: disable=attribute-error

official/nlp/modeling/layers/block_sparse_attention_test.py

@@ -27,9 +27,36 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("key_value_same_proj", None, None, [40, 80]),
+      ("key_value_same_proj_mqa", None, None, [40, 80], False, 1),
+      ("key_value_same_proj_multi_query_blocks", None, None, [40, 80], True),
+      (
+          "key_value_same_proj_multi_query_blocks_mqa",
+          None,
+          None,
+          [40, 80],
+          True,
+          1,
+      ),
       ("key_value_different_proj", 32, 60, [40, 60]),
+      ("key_value_different_proj_mqa", 32, 60, [40, 60], False, 1),
+      ("key_value_different_proj_multi_query_blocks", 32, 60, [40, 60], True),
+      (
+          "key_value_different_proj_multi_query_blocks_mqa",
+          32,
+          60,
+          [40, 60],
+          True,
+          1,
+      ),
   )
-  def test_non_masked_attention(self, value_dim, output_shape, output_dims):
+  def test_non_masked_attention(
+      self,
+      value_dim,
+      output_shape,
+      output_dims,
+      multi_query_blocks=False,
+      num_kv_heads=None,
+  ):
     """Test that the attention layer can be created without a mask tensor."""
     test_layer = block_sparse_attention.MultiHeadAttention(
         num_heads=12,
@@ -37,7 +64,8 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
         value_dim=value_dim,
         output_shape=output_shape,
         src_block_size=10,
-        tgt_block_size=5,
+        tgt_block_size=20 if multi_query_blocks else 5,
+        num_kv_heads=num_kv_heads,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     query = tf_keras.Input(shape=(40, 80))
@@ -57,13 +85,24 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("with_bias", True),
+      ("with_bias_mqa", True, False, False, 1),
+      ("with_bias_multi_query_blocks", True, False, True),
+      ("with_bias_multi_query_blocks_mqa", True, False, True, 1),
       ("no_bias", False),
+      ("no_bias_mqa", False, False, False, 1),
+      ("no_bias_multi_query_blocks", False, False, True),
+      ("no_bias_multi_query_blocks_mqa", False, False, True, 1),
       ("with_sigmoid_attn", True, True),
+      ("with_sigmoid_attn_mqa", True, True, False, 1),
+      ("with_sigmoid_attn_multi_query_blocks", True, True, True),
+      ("with_sigmoid_attn_multi_query_blocks_mqa", True, True, True, 1),
   )
   def test_masked_attention(
       self,
       use_bias,
       use_sigmoid_attn=False,
+      multi_query_blocks=False,
+      num_kv_heads=None,
   ):
     """Test with a mask tensor."""
     if use_sigmoid_attn:
@@ -75,9 +114,10 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
         key_dim=2,
         use_bias=use_bias,
         src_block_size=2,
-        tgt_block_size=1,
+        tgt_block_size=2 if multi_query_blocks else 1,
         use_sigmoid_attn=use_sigmoid_attn,
         sigmoid_attn_bias=sigmoid_attn_bias,
+        num_kv_heads=num_kv_heads,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     batch_size = 3

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -238,22 +238,30 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     self._sigmoid_attn_bias = sigmoid_attn_bias
     self._linformer_dim = linformer_dim
     self._linformer_shared_kv_projection = linformer_shared_kv_projection
-    if self._num_kv_heads is not None and self._src_block_size is not None:
+    if (
+        self._src_block_size is not None
+        and self._num_kv_heads is not None
+        and self._num_kv_heads != 1
+    ):
       raise ValueError(
-          "Block sparse attention does not support Multi-query attention."
-          " Specify only one of them."
+          "Block sparse attention only supports Multi-query attention.Please"
+          " set num_kv_heads to 1 to enable MQA with block sparse attention."
       )
     if attention_initializer:
       self._attention_initializer = tf_keras.initializers.get(
-          attention_initializer)
+          attention_initializer
+      )
     else:
       self._attention_initializer = tf_utils.clone_initializer(
-          self._kernel_initializer)
+          self._kernel_initializer
+      )
     self._attention_axes = attention_axes
 
     if self._diff_q_kv_att_layer_norm and not self._norm_first:
-      raise ValueError("Setting `diff_q_and_kv_attention_layer_norm` to True"
-                       "when `norm_first` is False is invalid.")
+      raise ValueError(
+          "Setting `diff_q_and_kv_attention_layer_norm` to True"
+          "when `norm_first` is False is invalid."
+      )
 
   def build(self, input_shape):
     if isinstance(input_shape, tf.TensorShape):
@@ -303,6 +311,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
           tgt_block_size=self._tgt_block_size,
           use_sigmoid_attn=self._use_sigmoid_attn,
           sigmoid_attn_bias=self._sigmoid_attn_bias,
+          num_kv_heads=self._num_kv_heads,
           name="block_sparse_attention",
       )
       attention_fn = block_sparse_attention.MultiHeadAttention

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -755,9 +755,11 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ('use_softmax_attn', False),
+      ('use_softmax_attn_mqa', False, 1),
       ('use_sigmoid_attn', True),
+      ('use_sigmoid_attn_mqa', True, 1),
   )
-  def test_block_sparse_attention(self, use_sigmoid_attn):
+  def test_block_sparse_attention(self, use_sigmoid_attn, num_kv_heads=None):
     num_attention_heads = 8
     sequence_length = 21
     width = 80
@@ -771,6 +773,7 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
         return_attention_scores=True,
         src_block_size=src_block_size,
         tgt_block_size=tgt_block_size,
+        num_kv_heads=num_kv_heads,
         use_sigmoid_attn=use_sigmoid_attn,
         sigmoid_attn_bias=-math.log(sequence_length)
         if use_sigmoid_attn

{tf_models_nightly-2.18.0.dev20241021.dist-info → tf_models_nightly-2.18.0.dev20241022.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.18.0.dev20241021
+Version: 2.18.0.dev20241022
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.18.0.dev20241021.dist-info → tf_models_nightly-2.18.0.dev20241022.dist-info}/RECORD

@@ -305,8 +305,8 @@ official/nlp/modeling/layers/bigbird_attention.py,sha256=dzutgRoQt2DFsYMpMILv_QF
 official/nlp/modeling/layers/bigbird_attention_test.py,sha256=cBYwK5k1rnykZ0gif-n7VaByLIoElA-N0_svCRKASoU,2206
 official/nlp/modeling/layers/block_diag_feedforward.py,sha256=FDEt-J_QjOxwar3eT5yjMs4hR41Ppke1zj7iswsZR4M,7243
 official/nlp/modeling/layers/block_diag_feedforward_test.py,sha256=wcg8In6FIOCxcKqe5rucftjJ_kUWTi9Ei7eEmlVCYpE,4181
-official/nlp/modeling/layers/block_sparse_attention.py,sha256=eY6jkSI-TrnL0JkP_9B-0DCxzppZdK_c8qp6Uw6yiD0,9923
-official/nlp/modeling/layers/block_sparse_attention_test.py,sha256=KSQENNhRG7Y1qDpdW_O3Ws6nPC4se7zv1UcxF2o7blI,15037
+official/nlp/modeling/layers/block_sparse_attention.py,sha256=8Jyinyf5XuU6nuxblGRfNnAviBEZIltMSoNJzAVjAYo,12233
+official/nlp/modeling/layers/block_sparse_attention_test.py,sha256=9YiKtv4YCrKIyUbv27P2xcTXFohoaRxq2K6vsOUi4zU,16447
 official/nlp/modeling/layers/cls_head.py,sha256=0X_gdjnAt6TZVrH_xkDcQCpwLuVz5Pb7d04wEVN_Kn8,16208
 official/nlp/modeling/layers/cls_head_test.py,sha256=01oMmiuyp1lDEXBYa9r3krn6BtH-QuSedGOca9LViEc,8888
 official/nlp/modeling/layers/factorized_embedding.py,sha256=4oFRYJbpoaSxqv8hTWY2JPGPllp-zhniz99IyRtlzV8,2902
@@ -363,8 +363,8 @@ official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=J52mXzoiuaXfR61kh
 official/nlp/modeling/layers/tn_transformer_expand_condense.py,sha256=gbGJOrgxJd1SyMGB6ME04FSxuZfHqsi94Xxt23l7368,11032
 official/nlp/modeling/layers/tn_transformer_test.py,sha256=Fh-EDRoAkhO7ccD3w3FsJHC51MnZySv8jBlHYnvKZMc,8893
 official/nlp/modeling/layers/transformer.py,sha256=yofIEOjZpcvDmHbcjBmkZrl5iSe6pLtMsetNbXmxDnY,20087
-official/nlp/modeling/layers/transformer_encoder_block.py,sha256=kiCQ4yGejmwRsJBKpmrwA1As4rFUekNYf9xGS052kyU,24766
-official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=cIunagl03W1tPkkt1BDVpGEpd-7ZwCqc3sPdzQOmpuc,32269
+official/nlp/modeling/layers/transformer_encoder_block.py,sha256=14dgbg6z9xeXl2trEJkxsVyQPguQ9m7U20aDAmOVDQE,24930
+official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=eTIDHGbTZobWIyMswPp9K_tgyzWTLFJ9j1ujXY3EXvY,32406
 official/nlp/modeling/layers/transformer_scaffold.py,sha256=m8TF4geBkm8-VJQiTpzMI6FSJZry6oa2vPO3FXCCClE,15704
 official/nlp/modeling/layers/transformer_scaffold_test.py,sha256=pqUGldhmAKROrd4eoCWmHNtKOdCO6PH_-EigcYnvIpE,19920
 official/nlp/modeling/layers/transformer_test.py,sha256=kC_9NcLbJnBbuTaE_7BW60EF8xG_QUoICj0t0gS7O4Q,5522
@@ -1222,9 +1222,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.18.0.dev20241021.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.18.0.dev20241021.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.18.0.dev20241021.dist-info/METADATA,sha256=RwWDKM5onfUl4FjQtHMFuJAKRJ5avUVrgN1TTVUXxYU,1432
-tf_models_nightly-2.18.0.dev20241021.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.18.0.dev20241021.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.18.0.dev20241021.dist-info/RECORD,,
+tf_models_nightly-2.18.0.dev20241022.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.18.0.dev20241022.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.18.0.dev20241022.dist-info/METADATA,sha256=Xq17BE4FrMRbkZCvlH9T5EFnFy7FIPjl2th1UY1LuUA,1432
+tf_models_nightly-2.18.0.dev20241022.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.18.0.dev20241022.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.18.0.dev20241022.dist-info/RECORD,,