handle torch_dtype in low cpu mem usage
Author: patil-suraj
Pull request: yes

# What does this PR do?

The `torch_dtype` argument in `from_pretrained` is not respected when loading the model with `low_cpu_mem_usage=True`.  This is because `_load_pretrained_model_low_mem` creates and assigns new `Paramter` rather than loading directly in model `state_dict`, so the `torch_dtype` is ignored. as can be seen [here](https://github.com/huggingface/transformers/issues/16378#issuecomment-1086502209).

https://github.com/huggingface/transformers/blob/013a7dbe3d8af8f16f2b6cb60b6e21258a9e1399/src/transformers/modeling_utils.py#L2168

This PR casts each tensor in the `state_dict` by retrieving the `data_type` from the model's meta parameters. 

Should be merged after #16548

Comments:
_The documentation is not available anymore as the PR was closed or merged._

The first chunk is definitely the right fix, unrelated to dtype

The second chunk is different. what you want is to move this code:

https://github.com/huggingface/transformers/blob/be9474bd3551d6a89fd788a0063895a9b316e7e0/src/transformers/modeling_utils.py#L1846-L1849

to before this code:
https://github.com/huggingface/transformers/blob/be9474bd3551d6a89fd788a0063895a9b316e7e0/src/transformers/modeling_utils.py#L1893-L1894

which would restore the scope of non-default dtype. can merge the 2 `if from_pt`.

Once this is done pytorch will automatically operate with `torch_dtype` for tensor allocation.

Please let me know if what I suggested is easy to follow. The first snipped I quoted is the closing of the special scope.

And the obvious question - should we need to add a new test?

 
 

@stas00 I already tried the approach that you suggested. But it didn't work. It seems `torch.set_default_dtype` doesn't affect `torch.load`. So all params in loaded `state_dict` will have the `dtype` with which they were saved, irrespective of `torch.set_default_dtype`.

> The first chunk is definitely the right fix, unrelated to dtype

This is actually fixed in https://github.com/huggingface/transformers/pull/16548, I added it in this PR just to be able to test the changes.

> And the obvious question - should we need to add a new test?

We should, IMO.

I run some experiments and you're correct, Suraj, on all accounts.

It's a bummer since someone who is low on CPU memory will have enough to be able to start the model training in lower precision but then not be able to finetune from the single precision checkpoint from the hub, targeting the same lower precision, since `torch.load` will force a full fp32 model copy, requiring 3x memory of the half-precision model normally and 2x with the `_load_pretrained_model_low_mem` hack.

I'm pretty sure we will have to abandon `torch.load` altogether as the models get bigger unless it becomes more flexible. There is no reason whatsoever to load all of the model's weights into memory at once. And more so when different GPUs should load different chunks of the model. It should be possible to load them in segments. But, of course, this is not the right forum to discuss that.

I made a feature request: https://github.com/pytorch/pytorch/issues/75242