State dictionary retrieval from offloaded modules

[blbadger]

What does this PR do?

Adds a function to retrieve a state dictionary from an offloaded module, loading the parameters into a specified device.

See transformers PR #27412 for more information

Working on the test too:)

TODO: replace disk offload conditional check with expand_device_map

@SunMarc

[SunMarc]

Thanks for the work @blbadger ! Let's add a test and try it with the transformers PR ! After we manage to make it work on accelerate and transformers, we can merge this! I left a few comments

[blbadger]

Thanks for the review @SunMarc ! I have implemented your suggestions and wrote the test, and refactored the transformers code to use this method.

I have written the execution device to switch to the cpu by default, although if you think it would be a better idea to keep this as a gpu (if available) I would be happy to change. The transformers test is a little more straightforward if the cpu is default, although I know that fitting code for a test is not a great design consideration:)

[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.

[blbadger]

Just a heads up, I got around to integrating this module with transformers PR #27412 and testing looks good for large models there too.

[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.

[SunMarc]

Awesome @blbadger ! I will merge this asap after we merge the PR that fixes the CI ! I've tested this PR with transformers and it works great !