handle torch_dtype in low cpu mem usage

[patil-suraj]

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.

transformers/src/transformers/modeling_utils.py

Line 2168 in 013a7db

new_val = state_dict[k]

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

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

LGTM but make sure to have @stas00 review before merging as he knows this code better than I do!

[stas00]

The first chunk is definitely the right fix, unrelated to dtype

The second chunk is different. what you want is to move this code:

transformers/src/transformers/modeling_utils.py

Lines 1846 to 1849 in be9474b

	if from_pt:
	# restore default dtype
	if dtype_orig is not None:
	torch.set_default_dtype(dtype_orig)

to before this code:

transformers/src/transformers/modeling_utils.py

Lines 1893 to 1894 in be9474b

	# make sure token embedding weights are still tied if needed
	model.tie_weights()

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?

[patil-suraj]

@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 #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.

[stas00]

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: pytorch/pytorch#75242

[stas00]

Thank you for fixing this, Suraj.

Perhaps let's make a call to community to add tests for all these new features and combinations of those where they are lacking? i.e. First Good Issue?