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

Add small util to enable FSDP offloading quickly #3006

Merged
merged 5 commits into from
Aug 12, 2024
Merged

Conversation

muellerzr
Copy link
Collaborator

What does this PR do?

Enables users to quickly do:

from accelerate.utils import enable_fsdp_ram_efficient_loading
enable_fsdp_ram_efficient_loading()

model = AutoModel.from_*

(Also some tests got commented out when I was toying with things, hence the odd diff)

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@BenjaminBossan

Other info

Note that we have self.ram_efficient_loading in the FSDP Plugin. For clarity it should be cpu_ram_efficient_loading and also this was only added in the prior PR, so it's fine.

@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
Member

@BenjaminBossan BenjaminBossan 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 updating the functionality, LGTM. Just a small nit, feel free to ignore.

@@ -27,6 +27,23 @@
logger = get_logger(__name__)


def enable_fsdp_ram_efficient_loading():
"""
Enables RAM efficient loading of Hugging Face models for FSDP in the environment.
Copy link
Member

Choose a reason for hiding this comment

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

Nit: "in the environment" might not really be helpful for a user.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

After having Stas complain so much about magical env variables being set, I'd rather over document it hah

@muellerzr muellerzr merged commit 90d5023 into main Aug 12, 2024
27 checks passed
@muellerzr muellerzr deleted the cpu_offload_util branch August 12, 2024 15:53
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