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.

Sign up for GitHub

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

Jump to bottom

destroy process group in end_training #3012

Merged
merged 4 commits into from
Aug 15, 2024
Merged

destroy process group in end_training #3012

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
db2a773
destroy process group
SunMarc Aug 14, 2024
472fe6f
rephrase
SunMarc Aug 14, 2024
a0a4d8b
style
SunMarc Aug 14, 2024
2856b97
fix on_main_process
SunMarc Aug 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/accelerate/accelerator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

"""
Runs any special end training behaviors, such as stopping trackers on the main process only. Should always be
called at the end of your script if using experiment tracking.
Runs any special end training behaviors, such as stopping trackers on the main process only or destoying
process group. Should always be called at the end of your script if using experiment tracking.

Example:

Expand All @@ -2698,6 +2697,10 @@ def end_training(self):
for tracker in self.trackers:
tracker.finish()

if torch.distributed.is_initialized():
# needed when using torch.distributed.init_process_group
torch.distributed.destroy_process_group()

def save(self, obj, f, safe_serialization=False):
"""
Save the object passed to disk once per machine. Use in place of `torch.save`.
Expand Down
Loading