[Utils] `align_module_device` (#3204)
* implement align_module
* add docs
* move to modeling utils, integrate into existing source code
* update source, expose through utils
* Suggested docstring
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
* Rewrite for readability, add try finally
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
* Use try-finally when aligning with hook
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
* apply style
* improve get_state_dict_from_offload readability
* Update docstring
Co-authored-by: Benjamin Bossan <BenjaminBossan@users.noreply.github.com>
* rename to align_module_device, update docstring
---------
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
Co-authored-by: Benjamin Bossan <BenjaminBossan@users.noreply.github.com>
