Fix dtype when loading to meta model

[zucchini-nlp]

What does this PR do?

Before #36335, when loading the weights to meta model we kept the dtype of the module params. This PR fixes it by setting assign=True. I am not sure though, why strict=False, should I change it to True?

Before (no dtype is passed to accelerate):
set_module_tensor_to_device(model, param_name, param_device, **set_module_kwargs)

After:
module.load_state_dict({param_type: param[:].to(param_device)}, False, True)

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

[ArthurZucker]

Suggested change

	assign=False,
	assign=True,

we need to assign to get speed boost / leverage meta device

[zucchini-nlp]

oke, then we will need to either cast params to same dtype as module or add a more involved logic of passing dtype for composite model. Personally, I think the second way is will complicate things even more

Otherwise any time one wants to load different dtypes for each backbones with accelerate (device_map or low_cpu_mem-usage), one will get only one dtype for the whole model

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.