Repository: huggingface/transformers
Pull request: Merge test_keep_in_fp32_modules and test_keep_in_fp32_modules_strict

The `test_keep_in_fp32_modules` issues in #44052 are because the test assumes a model has **either** `_keep_in_fp32_modules` or `_keep_in_fp32_modules_strict` **but not both.**

The only model that uses both is `glm_moe_dsa`, so this is the only model that fails the tests. This PR refactors things and merges both tests into a single `test_keep_in_fp32_modules` test. This test checks that:

1) When we load in `fp16`, the layers in `_keep_in_fp32_modules` and `_keep_in_fp32_modules_strict` are upcasted, and no other layers
2) When we load in `bf16`, the layers in `_keep_in_fp32_modules_strict` are upcasted, and no other layes.
3) There is no overlap between `_keep_in_fp32_modules` and `_keep_in_fp32_modules_strict` (because allowing this just creates weird edge cases, and there's no reason to)

This was always the intended behaviour, and now the test covers it correctly!

Fixes #44052, cc @arthurzucker

Changed files:
filename: tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py
patch: @@ -122,14 +122,6 @@ def test_training_overfit(
     ):
         pass
 
-    @unittest.skip("I am in a rush, will check it out later on")
-    def test_keep_in_fp32_modules_strict(self):
-        pass
-
-    @unittest.skip("I am in a rush, will check it out later on")
-    def test_keep_in_fp32_modules(self):
-        pass
-
     @require_torch_accelerator
     @slow
     def test_flash_attn_2_inference_equivalence_right_padding(self):

filename: tests/test_modeling_common.py
patch: @@ -914,61 +914,39 @@ def test_keep_in_fp32_modules_exist(self):
                     )
 
     def test_keep_in_fp32_modules(self):
-        """Test that the flag `_keep_in_fp32_modules` is correctly respected."""
+        """Test that the flag `_keep_in_fp32_modules` and `_keep_in_fp32_modules_strict`  is correctly respected."""
         config, _ = self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__):
                 model = model_class(copy.deepcopy(config))
-                if len(model._keep_in_fp32_modules) == 0:
+                if len(model._keep_in_fp32_modules) == 0 and len(model._keep_in_fp32_modules_strict) == 0:
                     self.skipTest(
-                        reason=f"{model_class.__name__} class has no _keep_in_fp32_modules attribute defined"
+                        reason=f"{model_class.__name__} class has no _keep_in_fp32_modules or _keep_in_fp32_modules_strict attribute defined"
                     )
 
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     model.save_pretrained(tmpdirname)
 
-                    # Test when reloading in fp16 -> should be upcasted to fp32
                     model = model_class.from_pretrained(tmpdirname, dtype=torch.float16)
-                    for name, param in model.state_dict().items():
-                        if any(re.search(rf"(?:^|\.){k}(?:\.|$)", name) for k in model._keep_in_fp32_modules):
-                            self.assertTrue(param.dtype == torch.float32, f"{name} not upcasted to fp32")
-                        else:
-                            self.assertTrue(param.dtype == torch.float16, f"{name} was upcasted but it should NOT")
-
-                    # Test when reloading in bf16 -> should stay bf16
-                    model = model_class.from_pretrained(tmpdirname, dtype=torch.bfloat16)
-                    for name, param in model.state_dict().items():
-                        self.assertTrue(param.dtype == torch.bfloat16, f"{name} was upcasted but it should NOT")
-
-    def test_keep_in_fp32_modules_strict(self):
-        """Test that the flag `_keep_in_fp32_modules_strict` is correctly respected."""
-        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
-        for model_class in self.all_model_classes:
-            with self.subTest(model_class.__name__):
-                model = model_class(copy.deepcopy(config))
-                if len(model._keep_in_fp32_modules_strict) == 0:
-                    self.skipTest(
-                        reason=f"{model_class.__name__} class has no _keep_in_fp32_modules_strict attribute defined"
+                    self.assertFalse(
+                        model._keep_in_fp32_modules & model._keep_in_fp32_modules_strict,
+                        "We found a layer in both the `_keep_in_fp32_modules` and `_keep_in_fp32_modules_strict` lists. Please remove it from one of the two lists.",
                     )
-
-                with tempfile.TemporaryDirectory() as tmpdirname:
-                    model.save_pretrained(tmpdirname)
-
-                    # Test when reloading in fp16 -> should be upcasted to fp32
-                    model = model_class.from_pretrained(tmpdirname, dtype=torch.float16)
+                    # When reloading in fp16, keep_in_fp32_modules AND keep_in_fp32_modules_strict should be upcasted
+                    all_fp32_modules = model._keep_in_fp32_modules | model._keep_in_fp32_modules_strict
                     for name, param in model.state_dict().items():
-                        if any(re.search(rf"(?:^|\.){k}(?:\.|$)", name) for k in model._keep_in_fp32_modules_strict):
+                        if any(re.search(rf"(?:^|\.){k}(?:\.|$)", name) for k in all_fp32_modules):
                             self.assertTrue(param.dtype == torch.float32, f"{name} not upcasted to fp32")
                         else:
-                            self.assertTrue(param.dtype == torch.float16, f"{name} was upcasted but it should NOT")
+                            self.assertTrue(param.dtype == torch.float16, f"{name} was upcasted but it should NOT be")
 
-                    # Test when reloading in bf16 -> should also be upcasted to fp32
+                    # When reloading in bf16, only keep_in_fp32_modules_strict should be upcasted
                     model = model_class.from_pretrained(tmpdirname, dtype=torch.bfloat16)
                     for name, param in model.state_dict().items():
                         if any(re.search(rf"(?:^|\.){k}(?:\.|$)", name) for k in model._keep_in_fp32_modules_strict):
                             self.assertTrue(param.dtype == torch.float32, f"{name} not upcasted to fp32")
                         else:
-                            self.assertTrue(param.dtype == torch.bfloat16, f"{name} was upcasted but it should NOT")
+                            self.assertTrue(param.dtype == torch.bfloat16, f"{name} was upcasted but it should NOT be")
 
     def test_save_load_keys_to_ignore_on_save(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()