[dist_optim] introduce distributed functional optimizer #45221
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
wanchaol
requested review from
apaszke,
mrshenli,
pietern,
pritamdamania87,
rohan-varma and
zhaojuanmao
as code owners
This was referenced Sep 23, 2020
wanchaol
added a commit
that referenced
this pull request
Sep 23, 2020
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training ghstack-source-id: 0c75d84 Pull Request resolved: #45221
💊 CI failures summary and remediationsAs of commit 500ac24 (more details on the Dr. CI page):
XLA failureJob pytorch_xla_linux_bionic_py3_6_clang9_test is failing. Please create an issue with title prefixed by This comment was automatically generated by Dr. CI (expand for details).Follow this link to opt-out of these comments for your Pull Requests.Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group. This comment has been revised 34 times. |
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
wanchaol
changed the title
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
wanchaol
added a commit
that referenced
this pull request
Sep 23, 2020
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training ghstack-source-id: 115adfe Pull Request resolved: #45221
pritamdamania87
suggested changes
Sep 23, 2020
Contributor
pritamdamania87
left a comment
pritamdamania87
left a comment
There was a problem hiding this comment.
Looks good overall, requesting changes to see if we can further optimize the local optimizer step to be completely in torchscript.
| @@ -0,0 +1,90 @@ | |||
| from typing import List, Dict, Optional | |||
Contributor
There was a problem hiding this comment.
nit: maybe rename this file to functional_adagrad.py
torch/distributed/optim/adagrad.py
Outdated
| # NOTE: This should be only used by distributed optimizer internals | ||
| # and not meant to expose to the user. | ||
| @torch.jit.script | ||
| class FunctionalAdagrad(object): |
Contributor
There was a problem hiding this comment.
nit: Should this be _FunctionalAdagrad for now to indicate its internal only?
|
|
||
| # TODO: use foreach API in optim.functional to do all the computation | ||
|
|
||
| def _make_sparse(grad, grad_indices, values): |
Contributor
There was a problem hiding this comment.
LGTM. What is the advantage of writing this in this way?
Collaborator
Author
There was a problem hiding this comment.
It's mainly because nested functions are not supported by TorchScript, so I moved it out as a separate function. Also I feel like define a nested function inside a for loop is not a good practice
torch/distributed/optim/optimizer.py
Outdated
| rpc_futs.append(rpc.rpc_async( | ||
| optim.owner(), | ||
| optimizer.owner(), | ||
| _local_optimizer_step, |
Contributor
There was a problem hiding this comment.
Can we have two versions of _local_optimizer_step? If we're using a torchscript functional optimizer, we can have a torchscript _local_optimizer_step which means RPC will call an torchscript function and there would be no GIL anywhere. If it is an optimizer that is not torchscript we call a regular python _local_optimizer_step.
In the current version, we still end up holding GIL since _local_optimizer_step is not torchscript.
Collaborator
Author
There was a problem hiding this comment.
hmm I see, this make sense, the problem of this is that _LocalOptimizer still contains Lock() objects to protect those who doesn't have a functional optimizer implementation. I could try adding a new _LocalScriptOptimizer to enable the use case you mentioned
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
wanchaol
added a commit
that referenced
this pull request
Sep 24, 2020
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training ghstack-source-id: 0b4ef3c Pull Request resolved: #45221
pritamdamania87
approved these changes
Sep 24, 2020
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
wanchaol
added a commit
that referenced
this pull request
Sep 24, 2020
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training ghstack-source-id: 5dd33d0 Pull Request resolved: #45221
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training [ghstack-poisoned]
wanchaol
added a commit
that referenced
this pull request
Sep 24, 2020
This PR introduces a distributed functional optimizer, so that distributed optimizer can reuse the functional optimizer APIs and maintain their own states. This could enable the torchscript compatible functional optimizer when using distributed optimizer, helps getting rid of GIL and improve overall performance of training, especially distributed model parallel training ghstack-source-id: 0e0a953 Pull Request resolved: #45221
Contributor
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack:
This PR introduces a distributed functional optimizer, so that
distributed optimizer can reuse the functional optimizer APIs and
maintain their own states. This could enable the torchscript compatible
functional optimizer when using distributed optimizer, helps getting rid
of GIL and improve overall performance of training, especially distributed
model parallel training
Differential Revision: D23935256