Add torch.nn.init.uniform_ operator to ShardedTensor. #63997
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
… to mimic [torch.nn.init.normal_, uniform_, kaiming_uniform_]
Summary:
Note _sharded_tensor module is a temporary place.
Ideally we want something like torch.nn.init(ShardedTensor, ...) to ensure consistent UX with Tensor.
To support that, we need either:
a) Update torch/nn/init.py with normal_(ShardedTensor, ), uniform_(ShardedTensor,...), and kaiming_uniform_(ShardedTensor, ...), or
b) Add torch.nn.init.{funcs} into __torch_function__ dispatchers (Currently __torch_function__ does not handle these funcs)
Test Plan:
(pytorch) ... $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestShardedTensorNNInit --v
Reviewers:
Subscribers:
Tasks:
Tags:
[ghstack-poisoned]
bowangbj
requested review from
H-Huang,
albanD,
cbalioglu,
jbschlosser,
mingzhe09088,
mrshenli,
pritamdamania87,
rohan-varma,
wayi1 and
zhaojuanmao
as code owners
facebook-github-bot
added
oncall: distributed
Contributor
🔗 Helpful links
💊 CI failures summary and remediationsAs of commit a570b13 (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 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 to the (internal) Dr. CI Users group. |
Contributor
Author
|
@bowangbj has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
bowangbj
changed the title
…ing_uniform_] utils to mimic torch.nn.init.[normal_, uniform_, kaiming_uniform_]"
Summary:
Note _sharded_tensor module is a temporary place.
Ideally we want something like torch.nn.init(ShardedTensor, ...) to ensure consistent UX with Tensor.
To support that, we need either:
a) Update torch/nn/init.py with normal_(ShardedTensor, ), uniform_(ShardedTensor,...), and kaiming_uniform_(ShardedTensor, ...), or
b) Add torch.nn.init.{funcs} into __torch_function__ dispatchers (Currently __torch_function__ does not handle these funcs)
Test Plan:
(pytorch) ... $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestShardedTensorNNInit --v
Reviewers:
Subscribers:
Tasks:
Tags:
tmp
Differential Revision: [D30563017](https://our.internmc.facebook.com/intern/diff/D30563017)
[ghstack-poisoned]
…rm_, kaiming_uniform_] utils to mimic torch.nn.init.[normal_, uniform_, kaiming_uniform_]"
Summary:
Note _sharded_tensor module is a temporary place.
Ideally we want something like torch.nn.init(ShardedTensor, ...) to ensure consistent UX with Tensor.
To support that, we need either:
a) Update torch/nn/init.py with normal_(ShardedTensor, ), uniform_(ShardedTensor,...), and kaiming_uniform_(ShardedTensor, ...), or
b) Add torch.nn.init.{funcs} into __torch_function__ dispatchers (Currently __torch_function__ does not handle these funcs)
Test Plan:
(pytorch) ... $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestShardedTensorNNInit --v
Reviewers:
Subscribers:
Tasks:
Tags:
Differential Revision: [D30563017](https://our.internmc.facebook.com/intern/diff/D30563017)
[ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Aug 26, 2021
… to mimic torch.nn.init.[normal_, uniform_, kaiming_uniform_]
Summary:
Note _sharded_tensor module is a temporary place.
Ideally we want something like torch.nn.init(ShardedTensor, ...) to ensure consistent UX with Tensor.
To support that, we need either:
a) Update torch/nn/init.py with normal_(ShardedTensor, ), uniform_(ShardedTensor,...), and kaiming_uniform_(ShardedTensor, ...), or
b) Add torch.nn.init.{funcs} into __torch_function__ dispatchers (Currently __torch_function__ does not handle these funcs)
Test Plan:
(pytorch) ... $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestShardedTensorNNInit --v
Reviewers:
Subscribers:
Tasks:
Tags:
ghstack-source-id: c7712e4
Pull Request resolved: #63997
Contributor
Author
|
@bowangbj has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
bowangbj
removed request for
H-Huang,
albanD,
cbalioglu,
mingzhe09088,
mrshenli,
rohan-varma,
wayi1 and
zhaojuanmao
CI Flow Status⚛️ CI FlowRuleset - Version:
You can add a comment to the PR and tag @pytorchbot with the following commands: # ciflow rerun, "ciflow/default" will always be added automatically
@pytorchbot ciflow rerun
# ciflow rerun with additional labels "-l <ciflow/label_name>", which is equivalent to adding these labels manually and trigger the rerun
@pytorchbot ciflow rerun -l ciflow/scheduled -l ciflow/slowFor more information, please take a look at the CI Flow Wiki. |
1 similar comment
CI Flow Status⚛️ CI FlowRuleset - Version:
You can add a comment to the PR and tag @pytorchbot with the following commands: # ciflow rerun, "ciflow/default" will always be added automatically
@pytorchbot ciflow rerun
# ciflow rerun with additional labels "-l <ciflow/label_name>", which is equivalent to adding these labels manually and trigger the rerun
@pytorchbot ciflow rerun -l ciflow/scheduled -l ciflow/slowFor more information, please take a look at the CI Flow Wiki. |
Contributor
Author
|
@bowangbj has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
Contributor
Author
|
Thanks Wanchao and Pritam, resolved all the cmts, ready to submit. Will follow up the args / kargs issue in follow up PR. |
Summary: Use torch_function to extend torch.nn.init.uniform_ The Init is done in SPMD fashion. Note that ideally we want to aggregate sharded tensors into a global tensor, init it and reshard. It's fine to run it SPMD since uniform is I.I.D indepenent and identifically distributed. Also enable unit test for test_linear.py for OSS test Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_linear.py --v (before runs this command is no-op) or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D30563017](https://our.internmc.facebook.com/intern/diff/D30563017) [ghstack-poisoned]
CI Flow Status⚛️ CI FlowRuleset - Version:
You can add a comment to the PR and tag @pytorchbot with the following commands: # ciflow rerun, "ciflow/default" will always be added automatically
@pytorchbot ciflow rerun
# ciflow rerun with additional labels "-l <ciflow/label_name>", which is equivalent to adding these labels manually and trigger the rerun
@pytorchbot ciflow rerun -l ciflow/scheduled -l ciflow/slowFor more information, please take a look at the CI Flow Wiki. |
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
Summary: Use torch_function to extend torch.nn.init.uniform_ The Init is done in SPMD fashion. Note that ideally we want to aggregate sharded tensors into a global tensor, init it and reshard. It's fine to run it SPMD since uniform is I.I.D indepenent and identifically distributed. Also enable unit test for test_linear.py for OSS test Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_linear.py --v (before runs this command is no-op) or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 99e45f7 Pull Request resolved: #63997
CI Flow Status⚛️ CI FlowRuleset - Version:
You can add a comment to the PR and tag @pytorchbot with the following commands: # ciflow rerun, "ciflow/default" will always be added automatically
@pytorchbot ciflow rerun
# ciflow rerun with additional labels "-l <ciflow/label_name>", which is equivalent to adding these labels manually and trigger the rerun
@pytorchbot ciflow rerun -l ciflow/scheduled -l ciflow/slowFor more information, please take a look at the CI Flow Wiki. |
Contributor
Author
|
@bowangbj has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
Contributor
|
This pull request has been merged in b6df043. |
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
…hardedTensor Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
…kaiming_uniform_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
… ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
…t.kaiming_uniform_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
…m_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 21, 2021
…hardedTensor Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 012f671 Pull Request resolved: #67057
bowangbj
added a commit
that referenced
this pull request
Oct 22, 2021
…torch.nn.init.kaiming_uniform_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 22, 2021
…iming_uniform_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 22, 2021
…hardedTensor Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 54ce4ba Pull Request resolved: #67057
bowangbj
added a commit
that referenced
this pull request
Oct 22, 2021
… and torch.nn.init.kaiming_uniform_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
bowangbj
added a commit
that referenced
this pull request
Oct 22, 2021
…it.kaiming_uniform_ ops to ShardedTensor" Summary: Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D31845654](https://our.internmc.facebook.com/intern/diff/D31845654) [ghstack-poisoned]
facebook-github-bot
pushed a commit
that referenced
this pull request
Oct 26, 2021
…hardedTensor (#67057) Summary: Pull Request resolved: #67057 Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from #63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Imported from OSS Reviewed By: pritamdamania87 Differential Revision: D31845654 fbshipit-source-id: e7aedc0972539da59f7b84bbbf617caf6b206d52
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:
Summary:
Use torch_function to extend torch.nn.init.uniform_
The Init is done in SPMD fashion. Note that ideally we want to aggregate sharded tensors into a global tensor, init it and reshard. It's fine to run it SPMD since uniform is I.I.D indepenent and identifically distributed.
Also enable unit test for test_linear.py for OSS test
Test Plan:
a) Unit Test
(pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v
(pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_linear.py --v (before runs this command is no-op)
or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit#
Reviewers:
Subscribers:
Tasks:
Tags:
Differential Revision: D30563017
cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang