Nail in edge case of torch dtype being overriden permantly in the case of an error

src/transformers/modeling_utils.py

@@ -98,6 +98,7 @@
     logging,
     replace_return_docstrings,
     strtobool,
+    test_injection,
 )
 from .utils.hub import create_and_tag_model_card, get_checkpoint_shard_files
 from .utils.import_utils import (
@@ -245,6 +246,25 @@ def set_zero3_state():
         _is_ds_init_called = False
 
 
+def restore_default_torch_dtype(func):
+    """
+    Decorator to restore the default torch dtype
+    at the end of the function. Serves
+    as a backup in case loading a model in raises
+    an error.
+    """
+
+    @wraps(func)
+    def _wrapper(*args, **kwargs):
+        old_dtype = torch.get_default_dtype()
+        try:
+            return func(*args, **kwargs)
+        finally:
+            torch.set_default_dtype(old_dtype)
+
+    return _wrapper
+
+
 def get_parameter_device(parameter: Union[nn.Module, "ModuleUtilsMixin"]):
     try:
         return next(parameter.parameters()).device
@@ -1401,6 +1421,7 @@ def add_model_tags(self, tags: Union[List[str], str]) -> None:
                 self.model_tags.append(tag)
 
     @classmethod
+    @restore_default_torch_dtype
     def _from_config(cls, config, **kwargs):
         """
         All context managers that the model should be initialized under go here.
@@ -1422,6 +1443,7 @@ def _from_config(cls, config, **kwargs):
         dtype_orig = None
         if torch_dtype is not None:
             dtype_orig = cls._set_default_torch_dtype(torch_dtype)
+        test_injection()
 
         config = copy.deepcopy(config)  # We do not want to modify the config inplace in _from_config.
 
@@ -3138,6 +3160,7 @@ def float(self, *args):
             return super().float(*args)
 
     @classmethod
+    @restore_default_torch_dtype
     def from_pretrained(
         cls: Type[SpecificPreTrainedModelType],
         pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
@@ -4059,6 +4082,7 @@ def from_pretrained(
                 for key in config.sub_configs.keys():
                     value = getattr(config, key)
                     value.torch_dtype = default_dtype
+            test_injection()
 
             # Check if `_keep_in_fp32_modules` is not None
             use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and (

src/transformers/utils/__init__.py

@@ -320,3 +320,10 @@ def get_available_devices() -> FrozenSet[str]:
         devices.add("musa")
 
     return frozenset(devices)
+
+
+def test_injection():
+    """
+    An injection point for testing hard-to-mock functions
+    """
+    return True

tests/utils/test_modeling_utils.py

@@ -1818,6 +1818,45 @@ def test_cache_when_needed_at_train_time(self):
         self.assertIsNone(model_outputs.past_key_values)
         self.assertTrue(model.training)
 
+    def test_restore_default_torch_dtype_from_pretrained(self):
+        """
+        Tests that the default torch dtype is restored
+        when an error happens during the loading of a model.
+        """
+        old_dtype = torch.get_default_dtype()
+        # set default type to float32
+        torch.set_default_dtype(torch.float32)
+        # Mock injection point which is right after the call to `_set_default_torch_dtype`
+        with mock.patch("transformers.modeling_utils.test_injection", side_effect=RuntimeError()):
+            with self.assertRaises(RuntimeError):
+                _ = AutoModelForCausalLM.from_pretrained(TINY_MISTRAL, device_map="auto", torch_dtype=torch.float16)
+        # default should still be float32
+        assert torch.get_default_dtype() == torch.float32
+        torch.set_default_dtype(old_dtype)
+
+    def test_restore_default_torch_dtype_from_config(self):
+        """
+        Tests that the default torch dtype is restored
+        when an error happens during the loading of a model.
+        """
+        old_dtype = torch.get_default_dtype()
+        # set default type to float32
+        torch.set_default_dtype(torch.float32)
+
+        config = AutoConfig.from_pretrained(
+            TINY_MISTRAL,
+        )
+        # Mock injection point which is right after the call to `_set_default_torch_dtype`
+        with mock.patch("transformers.modeling_utils.test_injection", side_effect=RuntimeError()):
+            with self.assertRaises(RuntimeError):
+                config.torch_dtype = torch.float16
+                _ = AutoModelForCausalLM.from_config(
+                    config,
+                )
+        # default should still be float32
+        assert torch.get_default_dtype() == torch.float32
+        torch.set_default_dtype(old_dtype)
+
 
 @slow
 @require_torch