Pytorch can have its default dtype permanently set to the "wrong" value if there is an exception when loading a model
Author: Taytay
Pull request: no
Labels: name: Good Difficult Issue

### System Info

I just ran into the most head-scratching issue. My data collator was crashing because a tensor it made was in half precision (fp16). I couldn't figure out why, but then I realized my `torch.get_default_dtype()` was `torch.float16`!

Then I realized it's because my model code threw an exception in a previous run of a notebook cell. 
And if you look at this code PreTrainedModel:_from_config : [def _from_config(cls, config, **kwargs):](https://github.com/huggingface/transformers/blob/995a7ce9a80b80062ccfe0b2d78857fb17351e27/src/transformers/modeling_utils.py#L1256-L1294)

You can see that it tries to set the dtype back to the original value, but doesn't do so in a `finally` block:
```python
        # override default dtype if needed
        dtype_orig = None
        if torch_dtype is not None:
            dtype_orig = cls._set_default_torch_dtype(torch_dtype)
            
        # do some stuff here....maybe throw an exception...

        # restore default dtype if it was modified (assuming we get to this line)
        if dtype_orig is not None:
            torch.set_default_dtype(dtype_orig)

        return model
```

This would of course leave my torch default dtype in whatever it was in when I was trying to load the model.

We could sprinkle some `finally` blocks around, or we could write a class like this:
```python

class temporily_set_default_torch_dtype:
    def __init__(self, dtype):
        self.new_dtype = dtype
        if dtype is not None:
            self.original_dtype = torch.get_default_dtype()
        else:
            # try to make this a no-op
            self.original_dtype = None

    def __enter__(self):
        if self.new_dtype is not None:
            torch.set_default_dtype(self.new_dtype)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.original_dtype is not None:
            torch.set_default_dtype(self.original_dtype)

```

And use it like so:
```python

torch.set_default_dtype(torch.float32)
print(f"default dtype is this before: {torch.get_default_dtype()}")

try:
  with temporily_set_default_torch_dtype(torch.float16):
      print(f"default dtype is now this inside: {torch.get_default_dtype()}")
      raise ValueError("Throwing an exception to make sure it works")
except ValueError as e:
  print("We caught the exception")
  pass

print(f"default dtype is this after: {torch.get_default_dtype()}")

# prints:
# default dtype is this before: torch.float32
# default dtype is now this inside: torch.float16
# default dtype is this after: torch.float32

```





### Who can help?

Think @ArthurZucker and @younesbelkada are correct here?

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1: Run a notebook cell that loads a model from_pretrained, in dtype=float16, and throws an exception while doing so.
2: Note that your torch.get_default_dtype() is still set to float16.

(This causes a real problem when things like the `DataCollatorForLanguageModeling` calls `torch_mask_tokens`, and then:
```python
# this will accidentally create a float16 tensor:
probability_matrix = torch.full(labels.shape, self.mlm_probability)
#...
probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
masked_indices = torch.bernoulli(probability_matrix).bool()
```

An exception gets thrown when you try to call `bernoulli` on a cpu tensor at half precision:
`RuntimeError: "bernoulli_tensor_cpu_self_" not implemented for 'Half'`


### Expected behavior

My default torch dtype should not get "corrupted" even if the model loading code throws an exception

Comments:
thanks for the deep investigation the proposed solution! 
I think having a context manager that handles that looks like a valid solution and might cover all edge cases including yours - I would be curious to hear what you think @amyeroberts @ArthurZucker ?
I think that @Taytay is right here in the sense that if any error occurs here: https://github.com/huggingface/transformers/blob/995a7ce9a80b80062ccfe0b2d78857fb17351e27/src/transformers/modeling_utils.py#L1272-L1287 for some reason, the original dtype will never get set back, leading to a "corrupted" `torch.get_default_dtype()`, and having a context manager with a proper `__exit__` method will correctly set back the original dtype after an exception has been thrown



We have a similar call in `from_pretrained`, I'd be in favor of a context manager but internal:
```python 
with _temp_default_dtype():
            config = copy.deepcopy(config)  # We do not want to modify the config inplace in _from_config.
        config._attn_implementation = kwargs.pop("attn_implementation", None)
        config = cls._autoset_attn_implementation(
            config, use_flash_attention_2=use_flash_attention_2, check_device_map=False
        )

        if is_deepspeed_zero3_enabled():
            import deepspeed

            logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
            # this immediately partitions the model across all gpus, to avoid the overhead in time
            # and memory copying it on CPU or each GPU first
            with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
                model = cls(config, **kwargs)
        else:
            model = cls(config, **kwargs)
```
same for https://github.com/huggingface/transformers/blob/07f5cdcac14b672ea6934b16da432518717f5b74/src/transformers/modeling_utils.py#L3734

try catch is also fine since we only put them in two places

@Taytay Thanks for raising this issue and the detailed investigation and proposal! Yes, I think a context manager here sounds like a good idea! 

This is a very serious issue and needs more attention. It causes problems in downstream applications all the time.

PRs are welcome! 🤗 I agree with you that this is welcome, anyone is free to open a PR for a fix and ping me for review. I did not have time to tackle this yet!