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

[TPU] For XLA Strategy, added function arg to control broadcast_master_param() #17522

Merged
merged 10 commits into from
May 5, 2023
Merged

[TPU] For XLA Strategy, added function arg to control broadcast_master_param() #17522

merged 10 commits into from
May 5, 2023

Conversation

gkroiz
Copy link
Contributor

@gkroiz gkroiz commented Apr 28, 2023 •
edited
Loading

What does this PR do?

Adds boolean sync_module_states to XLAStrategy to control whether or not to call broadcast_master_params(). Broadcasting master parameters is not always needed. For example, when initializing random weights, if the seed is defined and the same on all devices, then broadcast_master_params() is not needed.

The motivation for this change is that broadcast_master_params() adds additional time to training and if not needed should be skipped.

Before submitting
  • Was this discussed/agreed via a GitHub issue? (not for typos and docs)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you make sure to update the documentation with your changes? (if necessary)
  • Did you write any new necessary tests? (not for typos and docs)
  • Did you verify new and existing tests pass locally with your changes?
  • Did you list all the breaking changes introduced by this pull request?
  • Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist
  • Is this pull request ready for review? (if not, please submit in draft mode)
  • Check that all items from Before submitting are resolved
  • Make sure the title is self-explanatory and the description concisely explains the PR
  • Add labels and milestones (and optionally projects) to the PR so it can be classified

@gkroiz
[TPU] Add xla strat broadcast_master_params arg
cdb79d9 
Added function argument to XLA strategy controlling whether to call
broadcast_master_params(). Broadcasting master parameters is not always
needed. For example, when initializing random weights, if the seed is
defined and the same on all devices, then broadcast_master_params() is
not needed.
@github-actions github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels Apr 28, 2023
carmocca
carmocca reviewed May 3, 2023
View reviewed changes
Copy link
Contributor

@carmocca carmocca left a comment

Choose a reason for hiding this comment

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

We should add a test like this one: https://github.com/Lightning-AI/lightning/pull/17370/files#diff-664848c0cba1d81c012e047af0b592829a44870a55e75c24884569e81e3e62beR41

@gkroiz
Copy link
Contributor Author

gkroiz commented May 3, 2023

We should add a test like this one: https://github.com/Lightning-AI/lightning/pull/17370/files#diff-664848c0cba1d81c012e047af0b592829a44870a55e75c24884569e81e3e62beR41

I added a test for the Fabric XLAStrategy.

carmocca
carmocca approved these changes May 4, 2023
View reviewed changes
Copy link
Contributor

@carmocca carmocca left a comment

Choose a reason for hiding this comment

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

LGTM! Last bit is to add a CHANGELOG entry for both Fabric and PyTorch

tests/tests_fabric/strategies/test_xla.py Outdated Show resolved Hide resolved
tests/tests_fabric/strategies/test_xla.py Outdated Show resolved Hide resolved
@mergify mergify bot added the ready PRs ready to be merged label May 4, 2023
@gkroiz gkroiz changed the title [TPU] For XLA Strategy, added function arg to control broadcast_master_params() [TPU] For XLA Strategy, added function arg to control broadcast_master_param() May 5, 2023
@carmocca carmocca enabled auto-merge (squash) May 5, 2023 17:38
@carmocca carmocca merged commit 8e6f24b into Lightning-AI:master May 5, 2023
This was referenced May 5, 2023
Merged
Merged
@gkroiz gkroiz deleted the tpu_broadcast_params branch June 13, 2023 15:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@carmocca carmocca carmocca approved these changes

@awaelchli awaelchli awaelchli left review comments

@justusschock justusschock justusschock approved these changes

@Borda Borda Borda approved these changes

@williamFalcon williamFalcon Awaiting requested review from williamFalcon

Assignees
No one assigned
Labels
fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package ready PRs ready to be merged
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@gkroiz @awaelchli @Borda @carmocca @justusschock