[Utils] align_module_device

[kylesayrs]

Purpose

• Add an align_module utility function which temporarily moves a module's parameters to its execution device. The execution device may be overwritten and defaults to the execution device of the AlignDevicesHook if present
• Prepare for other utility functions to be implemented in future PRs

Changes

• Implement utility function in accelerate.utils.modeling, exposed through accelerate.utils
• Replace repeat code use in accelerate.utils.modeling and with newly added function

Testing

• Added tests in tests/test_modeling_utils.py.py

Who can review?

@SunMarc
@LysandreJik
@muellerzr
@mgoin

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

[muellerzr]

Thanks for adding the tests! Overall this looks great to me, I left a few readability suggestions 🤗

[kylesayrs]

@muellerzr Thank you for the suggestions!

[BenjaminBossan]

Thanks a lot for providing this new utility function, it makes a lot of sense to me. Also, nice tests.

I just some smaller comments, no blocker.

[SunMarc]

Very nice utility function ! Thanks for adding it !