Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Utils] align_module #3204

Open
wants to merge 9 commits into
base: main
Choose a base branch
from

Conversation

kylesayrs
Copy link
Contributor

@kylesayrs kylesayrs commented Oct 29, 2024

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.

Copy link
Collaborator

@muellerzr muellerzr left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

src/accelerate/utils/modeling.py Outdated Show resolved Hide resolved
src/accelerate/utils/modeling.py Outdated Show resolved Hide resolved
src/accelerate/utils/modeling.py Outdated Show resolved Hide resolved
kylesayrs and others added 4 commits October 30, 2024 15:33
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
Co-authored-by: Zach Mueller <muellerzr@gmail.com>
@kylesayrs
Copy link
Contributor Author

@muellerzr Thank you for the suggestions!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants