[tests] tighten keep_in_32 modules tests (#14399)

* tighten keep_in_32 modules tests

* remove unneeded tests

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Author: sayakpaul

tests/models/testing_utils/common.py

@@ -499,11 +499,6 @@ def test_getattr_is_correct(self, caplog):
 
         assert str(error.value) == f"'{type(model).__name__}' object has no attribute 'does_not_exist'"
 
-    @require_accelerator
-    @pytest.mark.skipif(
-        torch_device not in ["cuda", "xpu"],
-        reason="float16 and bfloat16 can only be used with an accelerator",
-    )
     def test_keep_in_fp32_modules(self, tmp_path):
         model = self.model_class(**self.get_init_dict())
         fp32_modules = model._keep_in_fp32_modules
@@ -516,11 +511,71 @@ def test_keep_in_fp32_modules(self, tmp_path):
         model.save_pretrained(tmp_path)
         model = self.model_class.from_pretrained(tmp_path, torch_dtype=torch.float16).to(torch_device)
 
+        # The rule is applied to every floating point checkpoint tensor, so persistent buffers are covered too.
+        # Non-persistent buffers are not in the checkpoint — they are regenerated by `__init__` and left alone.
+        for name, tensor in named_persistent_module_tensors(model, recurse=True):
+            if not tensor.is_floating_point():
+                continue
+            if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules):
+                assert tensor.dtype == torch.float32, f"{name} should be float32 but got {tensor.dtype}"
+            else:
+                assert tensor.dtype == torch.float16, f"{name} should be float16 but got {tensor.dtype}"
+
+    def test_keep_in_fp32_modules_as_str(self, tmp_path, monkeypatch):
+        model = self.model_class(**self.get_init_dict())
+        fp32_modules = model._keep_in_fp32_modules
+
+        if fp32_modules is None or len(fp32_modules) == 0:
+            pytest.skip("Model does not have _keep_in_fp32_modules defined.")
+
+        # Pick an entry that owns at least one parameter of the tiny test config, otherwise the assertions below
+        # would hold trivially.
+        parameter_name_parts = [name.split(".") for name, _ in model.named_parameters()]
+        fp32_module = next(
+            (module for module in fp32_modules if any(module in parts for parts in parameter_name_parts)), None
+        )
+        if fp32_module is None:
+            pytest.skip("No _keep_in_fp32_modules entry owns a parameter of this model.")
+
+        # `from_pretrained` also accepts `_keep_in_fp32_modules` declared as a bare string.
+        monkeypatch.setattr(self.model_class, "_keep_in_fp32_modules", fp32_module)
+
+        model.save_pretrained(tmp_path)
+        model = self.model_class.from_pretrained(tmp_path, torch_dtype=torch.float16).to(torch_device)
+
+        for name, param in model.named_parameters():
+            expected_dtype = torch.float32 if fp32_module in name.split(".") else torch.float16
+            assert param.dtype == expected_dtype, f"Parameter {name} should be {expected_dtype} but got {param.dtype}"
+
+    def test_keep_in_fp32_modules_layerwise_casting(self):
+        # Lives here rather than next to the other layerwise casting tests because it asserts
+        # `_keep_in_fp32_modules` semantics (`enable_layerwise_casting` folds it into the skip patterns) and needs
+        # no accelerator, while the layerwise casting mixin is accelerator-gated.
+        model = self.model_class(**self.get_init_dict())
+        fp32_modules = model._keep_in_fp32_modules
+
+        if fp32_modules is None or len(fp32_modules) == 0:
+            pytest.skip("Model does not have _keep_in_fp32_modules defined.")
+
+        if all(
+            any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules)
+            for name, _ in model.named_parameters()
+        ):
+            pytest.skip("Every parameter is kept in fp32, so layerwise casting has nothing to cast.")
+
+        # float16 storage instead of float8 so the assertions hold on every device — the skip patterns are applied
+        # the same way whatever the storage dtype is.
+        model.enable_layerwise_casting(storage_dtype=torch.float16, compute_dtype=torch.float32)
+
         for name, param in model.named_parameters():
             if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules):
                 assert param.dtype == torch.float32, f"Parameter {name} should be float32 but got {param.dtype}"
-            else:
-                assert param.dtype == torch.float16, f"Parameter {name} should be float16 but got {param.dtype}"
+
+        # The skip patterns keep several other modules in fp32 as well, so the loop above cannot check the
+        # complement. Assert that casting happened at all instead, otherwise it would pass on an untouched model.
+        assert any(param.dtype == torch.float16 for param in model.parameters()), (
+            "No parameter was cast to the storage dtype, so the assertions above hold trivially"
+        )
 
     def test_to_keep_in_fp32_modules_warns(self, caplog):
         fp32_modules = self.model_class._keep_in_fp32_modules