Fix test

Author: mseeger

tests/test_model.py

@@ -1438,7 +1438,7 @@ def test_sdpa_choice(config):
 
     torch.set_default_dtype(torch.float16)
 
-    def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logit_softcapping):
+    def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
         # SDPAParams gained an additional argument in PyTorch 2.5
         args = []
         assert k_and_v.both_in_parallel()
@@ -1456,7 +1456,7 @@ def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logi
             assert math_sdp_enabled(), "math_sdp_enabled() is False"
         else:
             raise NotImplementedError
-        return original_fn(query, k_and_v, scale, mask, attention_logit_softcapping)
+        return original_fn(query, k_and_v, mask, return_scores)
 
     config["n_layer"] = 1
     config = config_module.Config(**config)
@@ -1469,10 +1469,9 @@ def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logi
         # best effort, if the GPU can load it
         pytest.xfail()
 
-    for h in model.transformer.h:
-        litgpt.attention.scaled_dot_product_attention = partial(
-            assert_sdpa_backend, litgpt.attention.scaled_dot_product_attention
-        )
+    model.mha.scaled_dot_product_attention = partial(
+        assert_sdpa_backend, model.mha.scaled_dot_product_attention,
+    )
 
     if SUPPORTS_FLASH_ATTENTION:
         expected = SDPBackend.FLASH_ATTENTION
@@ -1490,7 +1489,7 @@ def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logi
 def test_sdpa_choice_kv_cache(config):
     torch.set_default_dtype(torch.float16)
 
-    def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logit_softcapping):
+    def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
         # SDPAParams gained an additional argument in PyTorch 2.5
         args = []
         assert k_and_v.both_in_parallel()
@@ -1507,7 +1506,7 @@ def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logi
             assert math_sdp_enabled()
         else:
             raise NotImplementedError
-        return original_fn(query, k_and_v, scale, mask, attention_logit_softcapping)
+        return original_fn(query, k_and_v, mask, return_scores)
 
     config["n_layer"] = 1
     config = config_module.Config(**config)
@@ -1522,10 +1521,9 @@ def assert_sdpa_backend(original_fn, query, k_and_v, scale, mask, attention_logi
         # best effort, if the GPU can load it
         pytest.xfail()
 
-    for h in model.transformer.h:
-        litgpt.attention.scaled_dot_product_attention = partial(
-            assert_sdpa_backend, litgpt.attention.scaled_dot_product_attention
-        )
+    model.mha.scaled_dot_product_attention = partial(
+        assert_sdpa_backend, model.mha.scaled_dot_product_attention,
+    )
 
     if SUPPORTS_FLASH_ATTENTION:
         # flash attention does not support an attention mask