test: cover padded SDPA outputs gradients and fallback paths

tests/test_attention_sdpa.py

@@ -0,0 +1,160 @@
+"""Regression tests for padded SDPA; no model weights or datasets are needed.
+
+Run: python -m unittest discover -s tests -p test_attention_sdpa.py -v
+"""
+import copy
+import unittest
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+from model.model_minimind import Attention, MiniMindConfig, precompute_freqs_cis
+
+
+class AttentionSDPATest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.old_threads = torch.get_num_threads()
+        torch.set_num_threads(min(cls.old_threads, 2))
+
+    @classmethod
+    def tearDownClass(cls):
+        torch.set_num_threads(cls.old_threads)
+
+    def make_inputs(self, device="cpu", dtype=torch.float32, dropout=0.0):
+        torch.manual_seed(42)
+        config = MiniMindConfig(hidden_size=64, num_hidden_layers=1,
+                               num_attention_heads=4, num_key_value_heads=2,
+                               head_dim=16, dropout=dropout)
+        attention = Attention(config).to(device=device, dtype=dtype).eval()
+        inputs = torch.randn(2, 12, 64, device=device, dtype=dtype)
+        positions = tuple(t.to(device) for t in precompute_freqs_cis(16, end=12))
+        return attention, inputs, positions
+
+    def check_parity(self, device, dtype):
+        for causal in (True, False):
+            for padding in ("none", "ones", "left", "right", "interior"):
+                with self.subTest(device=device, causal=causal, padding=padding):
+                    candidate, inputs, positions = self.make_inputs(device, dtype)
+                    candidate.is_causal = causal
+                    reference = copy.deepcopy(candidate)
+                    candidate.flash, reference.flash = True, False
+                    mask = torch.ones(2, 12, device=device, dtype=torch.long)
+                    if padding == "left":
+                        mask[0, :4], mask[1, :2] = 0, 0
+                    elif padding == "right":
+                        mask[0, 8:], mask[1, 10:] = 0, 0
+                    elif padding == "interior":
+                        mask[0, 3:5], mask[1, 6:8] = 0, 0
+                    valid = mask.bool()
+                    passed_mask = None if padding == "none" else mask
+                    outputs, gradients = [], []
+                    target = torch.randn_like(inputs).float()
+                    for attention in (reference, candidate):
+                        x = inputs.detach().clone().requires_grad_(True)
+                        output, _ = attention(x, positions, attention_mask=passed_mask)
+                        # Fully masked padded query rows are not equivalent to the
+                        # old finite -1e9 mask, and must not contribute to the loss.
+                        loss = (output[valid].float() - target[valid]).square().mean()
+                        loss.backward()
+                        self.assertTrue(torch.isfinite(output).all().item())
+                        gradient = torch.cat([x.grad.float().flatten()] + [
+                            p.grad.float().flatten() for p in attention.parameters()
+                            if p.grad is not None
+                        ])
+                        self.assertTrue(torch.isfinite(gradient).all().item())
+                        outputs.append(output.detach()[valid].float())
+                        gradients.append(gradient)
+                    bf16 = dtype == torch.bfloat16
+                    torch.testing.assert_close(outputs[1], outputs[0],
+                                               atol=1e-2 if bf16 else 1e-5,
+                                               rtol=3e-2 if bf16 else 1e-5)
+                    # Compare the gradient vectors, not just their norms.
+                    relative_error = ((gradients[1] - gradients[0]).norm()
+                                      / gradients[0].norm().clamp_min(1e-12)).item()
+                    self.assertLess(relative_error, 3e-2 if bf16 else 1e-4)
+
+    def test_cpu_outputs_and_full_gradients(self):
+        self.check_parity("cpu", torch.float32)
+
+    @unittest.skipUnless(torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
+                         "CUDA BF16 is optional")
+    def test_cuda_bf16_outputs_and_full_gradients(self):
+        self.check_parity("cuda", torch.bfloat16)
+
+    def test_padded_inputs_use_combined_sdpa_mask(self):
+        attention, inputs, positions = self.make_inputs()
+        mask = torch.ones(2, 12, dtype=torch.long)
+        mask[0, :3], mask[1, 9:] = 0, 0
+        for causal in (True, False):
+            attention.is_causal = causal
+            with self.subTest(causal=causal), patch(
+                "model.model_minimind.F.scaled_dot_product_attention",
+                wraps=F.scaled_dot_product_attention,
+            ) as sdpa:
+                attention(inputs, positions, attention_mask=mask)
+                sdpa.assert_called_once()
+                kwargs = sdpa.call_args.kwargs
+                expected = torch.ones(12, 12, dtype=torch.bool)
+                if causal:
+                    expected = expected.tril()
+                expected = expected & mask.bool()[:, None, None, :]
+                torch.testing.assert_close(kwargs["attn_mask"], expected)
+                self.assertFalse(kwargs["is_causal"])
+
+    def test_unpadded_inputs_keep_original_causal_sdpa(self):
+        attention, inputs, positions = self.make_inputs()
+        for mask in (None, torch.ones(2, 12, dtype=torch.long)):
+            with self.subTest(mask_is_none=mask is None), patch(
+                "model.model_minimind.F.scaled_dot_product_attention",
+                wraps=F.scaled_dot_product_attention,
+            ) as sdpa:
+                attention(inputs, positions, attention_mask=mask)
+                sdpa.assert_called_once()
+                self.assertIsNone(sdpa.call_args.kwargs["attn_mask"])
+                self.assertTrue(sdpa.call_args.kwargs["is_causal"])
+
+    def test_manual_fallbacks_and_cached_continuation(self):
+        attention, inputs, positions = self.make_inputs()
+        mask = torch.ones(2, 12, dtype=torch.long)
+        mask[0, :2] = 0
+        attention.flash = False
+        with torch.no_grad():
+            _, cache = attention(inputs[:, :8], tuple(t[:8] for t in positions),
+                                 use_cache=True, attention_mask=mask[:, :8])
+        for length, past, enabled in ((12, None, False), (1, None, True),
+                                      (1, cache, True), (4, cache, True)):
+            with self.subTest(length=length, cached=past is not None, flash=enabled):
+                start = 8 if past is not None else 0
+                x = inputs[:, start:start + length]
+                pos = tuple(t[start:start + length] for t in positions)
+                passed_mask = mask[:, :start + length]
+                attention.flash = enabled
+                with patch("model.model_minimind.F.scaled_dot_product_attention",
+                           wraps=F.scaled_dot_product_attention) as sdpa:
+                    output, _ = attention(x, pos, past_key_value=past,
+                                          attention_mask=passed_mask)
+                    sdpa.assert_not_called()
+                reference = copy.deepcopy(attention)
+                reference.flash = False
+                expected, _ = reference(x, pos, past_key_value=past,
+                                        attention_mask=passed_mask)
+                torch.testing.assert_close(output, expected, atol=0, rtol=0)
+
+    def test_dropout_probability_tracks_training_mode(self):
+        attention, inputs, positions = self.make_inputs(dropout=0.2)
+        mask = torch.ones(2, 12, dtype=torch.long)
+        mask[:, -2:] = 0
+        for training, probability in ((True, 0.2), (False, 0.0)):
+            attention.train(training)
+            with self.subTest(training=training), patch(
+                "model.model_minimind.F.scaled_dot_product_attention",
+                wraps=F.scaled_dot_product_attention,
+            ) as sdpa:
+                attention(inputs, positions, attention_mask=mask)
+                self.assertEqual(sdpa.call_args.kwargs["dropout_p"], probability)
+
+
+if __name__ == "__main__":
+    unittest.main()