Fix

Author: mseeger

tests/test_model.py

@@ -1435,6 +1435,8 @@ def test_sdpa_choice(config):
         pytest.skip("Gemma 2 doesn't support SDPA")
 
     torch.set_default_dtype(torch.float16)
+    config["n_layer"] = 1
+    config = config_module.Config(**config)
     enable_gqa = config["n_query_groups"] < config["n_head"]
 
     def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
@@ -1457,9 +1459,6 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             raise NotImplementedError
         return original_fn(query, k_and_v, mask, return_scores)
 
-    config["n_layer"] = 1
-    config = config_module.Config(**config)
-
     try:
         with torch.device("cuda"):
             model = GPT(config)
@@ -1488,6 +1487,8 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
 @torch.inference_mode()
 def test_sdpa_choice_kv_cache(config):
     torch.set_default_dtype(torch.float16)
+    config["n_layer"] = 1
+    config = config_module.Config(**config)
     enable_gqa = config["n_query_groups"] < config["n_head"]
 
     def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
@@ -1509,9 +1510,6 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             raise NotImplementedError
         return original_fn(query, k_and_v, mask, return_scores)
 
-    config["n_layer"] = 1
-    config = config_module.Config(**config)
-
     try:
         with torch.device("cuda"):
             model = GPT(config)