Remove head_mask and attention weights from VideoGPT (facebookresearch#536)

Summary:

Remove dead `head_mask`, `return_attn_weights`, and `attention_weights` from the VideoGPT stack. These features were never used by any consumer — `head_mask` was always `None` or all-ones, and `return_attn_weights` was always `False` except in tests that verified the feature itself.

Reviewed By: OmarPavel

Differential Revision: D92927089

Author: stashuk-olek

examples/mugen/tests/generation/test_text_video_gpt.py

@@ -134,19 +134,17 @@ def test_lookup(model_fn, modality, expected_shape, expected_sum):
 def test_forward_no_pretrained(model_fn, video_seq_len, expected):
     test_params = {"video_seq_len": video_seq_len}
     kwargs = {**_model_params, **test_params}
-    n_head = kwargs["n_head"]
 
     x = torch.tensor([[1, 2, 3, 4]])
     y = torch.tensor([[5, 6, 7]])
     attn_mask = torch.tril(torch.ones(7, 7)).unsqueeze(0)  # (b, seq_len, seq_len)
-    head_mask = torch.ones(1, n_head, 7, 7)  # (b, h, seq_len, seq_len)
 
     model = model_fn(**kwargs)
     model.eval()
     num_tokens = model.num_in_tokens + model.num_out_tokens
     logits_mask = torch.ones(1, 7, num_tokens)  # (b, seq_len, num_tokens)
 
-    out = model(x, y, attn_mask=attn_mask, head_mask=head_mask, logits_mask=logits_mask)
+    out = model(x, y, attn_mask=attn_mask, logits_mask=logits_mask)
     actual = out.decoder_output.last_hidden_states
     assert_expected(actual.shape, (1, 7, 768))
     assert_expected(actual.sum().item(), expected, rtol=1e-5, atol=1e-4)
@@ -163,19 +161,17 @@ def test_forward_vqvae_pretrained(model_fn, video_seq_len, expected):
         "pretrained_video_vqvae_model_key": vqvae_model_key,
     }
     kwargs = {**_model_params, **test_params}
-    n_head = kwargs["n_head"]
 
     x = torch.tensor([[1, 2, 3, 4]])
     y = torch.tensor([[5, 6, 7]])
     attn_mask = torch.tril(torch.ones(7, 7)).unsqueeze(0)  # (b, seq_len, seq_len)
-    head_mask = torch.ones(1, n_head, 7, 7)  # (b, h, seq_len, seq_len)
 
     model = model_fn(**kwargs)
     model.eval()
     num_tokens = model.num_in_tokens + model.num_out_tokens
     logits_mask = torch.ones(1, 7, num_tokens)  # (b, seq_len, num_tokens)
 
-    out = model(x, y, attn_mask=attn_mask, head_mask=head_mask, logits_mask=logits_mask)
+    out = model(x, y, attn_mask=attn_mask, logits_mask=logits_mask)
     actual = out.decoder_output.last_hidden_states
     assert_expected(actual.shape, (1, 7, 768))
     assert_expected(actual.sum().item(), expected, rtol=1, atol=1e-4)
@@ -192,19 +188,17 @@ def test_forward_gpt_pretrained(model_fn, video_seq_len, expected):
         "pretrained_text_video_gpt_model_key": gpt_model_key,
     }
     kwargs = {**_model_params, **test_params}
-    n_head = kwargs["n_head"]
 
     x = torch.tensor([[1, 2, 3, 4]])
     y = torch.tensor([[5, 6, 7]])
     attn_mask = torch.tril(torch.ones(7, 7)).unsqueeze(0)  # (b, seq_len, seq_len)
-    head_mask = torch.ones(1, n_head, 7, 7)  # (b, h, seq_len, seq_len)
 
     model = model_fn(**kwargs)
     model.eval()
     num_tokens = model.num_in_tokens + model.num_out_tokens
     logits_mask = torch.ones(1, 7, num_tokens)  # (b, seq_len, num_tokens)
 
-    out = model(x, y, attn_mask=attn_mask, head_mask=head_mask, logits_mask=logits_mask)
+    out = model(x, y, attn_mask=attn_mask, logits_mask=logits_mask)
     actual = out.decoder_output.last_hidden_states
     assert_expected(actual.shape, (1, 7, 768))
     assert_expected(actual.sum().item(), expected, rtol=1, atol=1e-4)

tests/models/test_gpt.py

@@ -108,18 +108,6 @@ def _attn_mask(q_seq_len, k_seq_len=None):
     return _attn_mask
 
 
-@pytest.fixture
-def head_mask(n_head):
-    def _head_mask(q_seq_len, k_seq_len=None):
-        if k_seq_len is None:
-            k_seq_len = q_seq_len
-        masked = torch.zeros(1, q_seq_len, k_seq_len)
-        unmasked = torch.ones(n_head - 1, q_seq_len, k_seq_len)
-        return torch.cat((masked, unmasked), dim=0)  # (h, q_seq_len, k_seq_len)
-
-    return _head_mask
-
-
 @pytest.fixture
 def logits_mask(in_seq_len, out_seq_len, num_in_tokens, num_out_tokens):
     total_seq_len = in_seq_len + out_seq_len
@@ -379,19 +367,16 @@ def test_forward(
         in_tokens,
         out_tokens,
         attn_mask,
-        head_mask,
     ):
         gpt = gpt()
 
         b, in_seq_len = in_tokens.shape
         b, out_seq_len = out_tokens.shape
         attn_mask = attn_mask(in_seq_len + out_seq_len)
-        head_mask = head_mask(in_seq_len + out_seq_len)
         actual = gpt(
             in_tokens=in_tokens,
             out_tokens=out_tokens,
             attn_mask=attn_mask,
-            head_mask=head_mask,
             use_cache=True,
             causal=True,
             right_shift=True,
@@ -401,10 +386,9 @@ def test_forward(
             "decoder_output": {
                 "last_hidden_states": (
                     torch.Size([1, 7, 4]),  # (b, seq_len, d_model)
-                    64.5348,
+                    64.3909,
                 ),
                 "hidden_states": None,
-                "attention_weights": None,
                 "past_key_values": (
                     (
                         {
@@ -424,20 +408,17 @@ def test_forward_logits_mask(
         in_tokens,
         out_tokens,
         attn_mask,
-        head_mask,
         logits_mask,
     ):
         gpt = gpt()
 
         b, in_seq_len = in_tokens.shape
         b, out_seq_len = out_tokens.shape
         attn_mask = attn_mask(in_seq_len + out_seq_len)
-        head_mask = head_mask(in_seq_len + out_seq_len)
         out = gpt(
             in_tokens=in_tokens,
             out_tokens=out_tokens,
             attn_mask=attn_mask,
-            head_mask=head_mask,
             use_cache=True,
             causal=True,
             right_shift=True,
@@ -476,7 +457,6 @@ def test_forward_in_modality(self, mm_decoder, in_modality):
                 0.2222,
             ),  # (b, in_seq_len, d_model)
             "hidden_states": None,
-            "attention_weights": None,
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
@@ -491,7 +471,6 @@ def test_forward_out_modality(self, mm_decoder, out_modality):
                 5.2093,
             ),  # (b, out_seq_len, d_model)
             "hidden_states": None,
-            "attention_weights": None,
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
@@ -509,7 +488,6 @@ def test_forward_two_modality(self, mm_decoder, in_modality, out_modality):
                 7.9519,
             ),  # (b, in_seq_len + out_seq_len, d_model)
             "hidden_states": None,
-            "attention_weights": None,
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
@@ -537,7 +515,6 @@ def test_forward_eval_right_shift_on(
                 7.9519,
             ),  # (b, in_seq_len + out_seq_len, d_model)
             "hidden_states": None,
-            "attention_weights": None,
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
@@ -564,7 +541,6 @@ def test_forward_eval_right_shift_off(
                 10.1681,
             ),  # (b, in_seq_len + out_seq_len, d_model)
             "hidden_states": None,
-            "attention_weights": None,
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
@@ -583,37 +559,31 @@ def test_optional_pos_ids(self, mm_decoder, in_modality):
 
 
 class TestTransformerDecoder:
-    def test_forward_mask_extended(
-        self, decoder, decoder_input, attn_mask, head_mask, num_layers
-    ):
+    def test_forward_mask_extended(self, decoder, decoder_input, attn_mask, num_layers):
         b, seq_len, _ = decoder_input.shape
         attn_mask = attn_mask(seq_len).unsqueeze(0)  # add batch dim
-        head_mask = head_mask(seq_len).unsqueeze(0)
-        actual = decoder(decoder_input, attn_mask, head_mask)
+        actual = decoder(decoder_input, attn_mask)
         assert isinstance(actual, TransformerDecoderOutput)
         assert_expected(actual.last_hidden_states.shape, torch.Size([1, 3, 4]))
 
-    def test_forward(self, decoder, decoder_input, attn_mask, head_mask, num_layers):
+    def test_forward(self, decoder, decoder_input, attn_mask, num_layers):
         b, seq_len, _ = decoder_input.shape
         attn_mask = attn_mask(seq_len)
-        head_mask = head_mask(seq_len)
-        actual = decoder(decoder_input, attn_mask, head_mask)
+        actual = decoder(decoder_input, attn_mask)
         assert isinstance(actual, TransformerDecoderOutput)
         assert_expected(actual.last_hidden_states.shape, torch.Size([1, 3, 4]))
 
     def test_forward_additional_output(self, decoder, decoder_input, num_layers):
         actual = decoder(
             decoder_input,
             use_cache=True,
-            return_attn_weights=True,
             return_hidden_states=True,
         )
         assert isinstance(actual, TransformerDecoderOutput)
         assert_expected(actual.last_hidden_states.shape, torch.Size([1, 3, 4]))
         assert_expected(
             len(actual.hidden_states), num_layers + 1
         )  # +1 to include the input hidden_states
-        assert_expected(len(actual.attention_weights), num_layers)
         assert_expected(len(actual.past_key_values), num_layers)
 
 
@@ -623,33 +593,26 @@ def test_forward(self, decoder_layer, decoder_input):
         assert isinstance(actual, TransformerLayerOutput)
         expected = {
             "hidden_states": (torch.Size([1, 3, 4]), 5.1956),  # (b, seq_len, d_model)
-            "attention_weights": None,
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
 
-    def test_forward_masked(self, decoder_layer, decoder_input, attn_mask, head_mask):
+    def test_forward_masked(self, decoder_layer, decoder_input, attn_mask):
         b, seq_len, _ = decoder_input.shape
         attn_mask = attn_mask(seq_len)
-        head_mask = head_mask(seq_len)
-        actual = decoder_layer(decoder_input, attn_mask, head_mask)
+        actual = decoder_layer(decoder_input, attn_mask)
         assert isinstance(actual, TransformerLayerOutput)
         expected = {
-            "hidden_states": (torch.Size([1, 3, 4]), 7.0397),  # (b, seq_len, seq_len)
-            "attention_weights": None,
+            "hidden_states": (torch.Size([1, 3, 4]), 5.6602),  # (b, seq_len, seq_len)
             "past_key_values": None,
         }
         assert_expected_namedtuple(actual, expected, rtol=1e-5, atol=1e-4)
 
     def test_forward_additional_output(self, decoder_layer, decoder_input):
-        actual = decoder_layer(decoder_input, use_cache=True, return_attn_weights=True)
+        actual = decoder_layer(decoder_input, use_cache=True)
         assert isinstance(actual, TransformerLayerOutput)
         expected = {
             "hidden_states": (torch.Size([1, 3, 4]), 5.1956),  # (b, seq_len, seq_len)
-            "attention_weights": (
-                torch.Size([1, 2, 3, 3]),
-                6.0,
-            ),  # (b, h, seq_len, seq_len)
             "past_key_values": {
                 "k": ([1, 2, 3, 2], 4.8075),
                 "v": ([1, 2, 3, 2], -5.6613),

tests/models/test_video_gpt.py

@@ -95,12 +95,9 @@ def test_forward(self, model_fn, model_params):
         model = model_fn(**model_params)
         model.eval()
 
-        n_head = model_params["n_head"]
-
         x = torch.tensor([[1, 2, 3, 4]])  # (b, in_seq_len)
         y = torch.tensor([[5, 6, 7]])  # (b, out_seq_len)
         attn_mask = torch.tril(torch.ones(7, 7)).unsqueeze(0)  # (b, seq_len, seq_len)
-        head_mask = torch.ones(1, n_head, 7, 7)  # (b, h, seq_len, seq_len)
 
         num_tokens = model.num_in_tokens + model.num_out_tokens
         logits_mask = torch.ones(1, 7, num_tokens)  # (b, seq_len, num_tokens)
@@ -109,11 +106,9 @@ def test_forward(self, model_fn, model_params):
             x,
             y,
             attn_mask=attn_mask,
-            head_mask=head_mask,
             logits_mask=logits_mask,
             use_cache=True,
             causal=True,
-            return_attn_weights=True,
             return_hidden_states=True,
         )
         actual = out.decoder_output.last_hidden_states

tests/models/test_video_vqvae.py

@@ -72,7 +72,7 @@ def axial_attn(self):
 
     def test_axial_attention(self, axial_attn, q, kv):
         k = v = kv
-        actual, _ = axial_attn(q, k, v)
+        actual = axial_attn(q, k, v)
         expected = torch.tensor(
             [
                 [