Fix dtype bug when offload_state_dict=True and dtype is specified

[fxmarty]

As per title.

I don't have time to provide a full repro & add a test - could do next week.

But for general context, the issue is the following:

When loading a state dict with offload_state_dict=True, we go into

accelerate/src/accelerate/utils/modeling.py

Line 1403 in bd72a5f

elif param_device == "cpu" and offload_state_dict:

and we are missing a check similar to

accelerate/src/accelerate/utils/modeling.py

Line 292 in bd72a5f

elif not str(value.dtype).startswith(("torch.uint", "torch.int", "torch.bool")):

when the argument dtype is specified.

The byproduct is that all parameters are casted to dtype, although they should not.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[SunMarc]

LGTM ! Left a few comments. I will run a few tests about bnb cpu offload to see if nothing is broken then merge.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[muellerzr]

Gentle ping @fxmarty 🤗

[fxmarty]

This typing is wrong as None is not an accepted value in this function (raises an error at

accelerate/src/accelerate/utils/modeling.py

Line 436 in bd72a5f

if param in device_map:

)

[fxmarty]

I leave it as is because changing to device_map: Dict[str, Union[int, str, torch.device]] raises errors in the CI.

[SunMarc]

Thx for noticing this. This should be fixed here.

[muellerzr]

Thanks for iterating! Overall looks great to me bar one suggestion. cc @SunMarc for a final look through?

[muellerzr]

@fxmarty could you run make style; make quality? Thanks!

[fxmarty]

Done thank you!

[SunMarc]

LGTM ! Thanks for adding a test.