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

destroy process group in end_training #3012

Merged
merged 4 commits into from
Aug 15, 2024
Merged

destroy process group in end_training #3012

merged 4 commits into from
Aug 15, 2024

Conversation

SunMarc
Copy link
Member

@SunMarc SunMarc commented Aug 14, 2024

What does this PR do ?

With the latest version of torch, when we use multi-gpu, torch will trigger a warning asking us to call destoy_process_group(). This PR fixes this by adding that in end_training method. Note that this function needs to be called on all process.

For trackers, we are already only executing the methods on the main process. See here. So it should be safe to remove the on_main_process decorator in the end_training method.
However, I see for WandBTracker that main_process_only = False. Is there a specific reason @muellerzr ?

@SunMarc SunMarc requested a review from muellerzr August 14, 2024 11:28
@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.

With the removal of @on_main_process, would tracker.finish() not be called on all processes where previously that was only called on the main process?

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.

Nice!

@muellerzr
Copy link
Collaborator

All of the trackers come loaded with @on_main_process decorators already for their functions (if needed/applicable) so this is fine how it is/the right method: https://github.com/huggingface/accelerate/blob/main/src/accelerate/tracking.py#L571-L576

@@ -2678,11 +2678,10 @@ def log(self, values: dict, step: int | None = None, log_kwargs: dict | None = {
for tracker in self.trackers:
tracker.log(values, step=step, **log_kwargs.get(tracker.name, {}))

@on_main_process
def end_training(self):
Copy link
Collaborator

Choose a reason for hiding this comment

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

Note that eventually we'll get to a point where users will want to pass in their own pg's and init them, e.g. this is needed for hybrid shard w/ FSDP. So we may need to take that in there potentially. Will see how it ends up looking

@muellerzr muellerzr merged commit 589fddd into main Aug 15, 2024
28 checks passed
@muellerzr muellerzr deleted the destroy_process_group branch August 15, 2024 12:31
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.

4 participants