diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py index 674421669f3d..80d3cfa8c93b 100644 --- a/tests/models/testing_utils/common.py +++ b/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