Adrian Wälchli
866457e746
Fix pydocstyle errors in fully_sharded_data_parallel.py, api.py, graph_utils.py, distribute.py, iter_graph_module.py, comm_tensor.py, experimental_ops.py, batch_dim_utils.py, data_parallel.py, graph_optimization.py ( #113216 )
...
Fixes #113191
```
pydocstyle torch/distributed/fsdp/fully_sharded_data_parallel.py --count
```
On master: 80
After my changes on this PR: 3
```
pydocstyle torch/distributed/_spmd/comm_tensor.py --count
```
On master: 5
After my changes on this PR: 3
```
pydocstyle torch/distributed/_spmd/experimental_ops.py --count
```
On master: 3
After my changes on this PR: 1
```
pydocstyle torch/distributed/_spmd/iter_graph_module.py --count
```
On master: 39
After my changes on this PR: 27
```
pydocstyle torch/distributed/_spmd/graph_utils.py --count
```
On master: 16
After my changes on this PR: 4
```
pydocstyle torch/distributed/_spmd/distribute.py --count
```
On master: 19
After my changes on this PR: 10
```
pydocstyle torch/distributed/_spmd/api.py --count
```
On master: 10
After my changes on this PR: 3
```
pydocstyle torch/distributed/_spmd/batch_dim_utils.py --count
```
On master: 14
After my changes on this PR: 3
```
pydocstyle torch/distributed/_spmd/data_parallel.py --count
```
On master: 34
After my changes on this PR: 2
```
pydocstyle torch/distributed/_spmd/graph_optimization.py --count
```
On master: 35
After my changes on this PR: 13
Pull Request resolved: https://github.com/pytorch/pytorch/pull/113216
Approved by: https://github.com/ezyang
2023-11-10 03:08:32 +00:00
Peter Bell
66c32d099a
Use pytree.arg_tree_leaves everywhere ( #112394 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/112394
Approved by: https://github.com/lezcano
ghstack dependencies: #112391 , #112392 , #112393
2023-10-31 15:57:06 +00:00
Peter Bell
bbd5b935e4
Use pytree.tree_leaves everywhere ( #112324 )
...
This changes all the instances I could find of `tree_flatten(...)[0]` or
`x, _ = tree_flatten` to use `tree_leaves`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/112324
Approved by: https://github.com/lezcano
ghstack dependencies: #112327 , #112323
2023-10-30 03:39:04 +00:00
Chien-Chin Huang
41d7969590
[SPMD] Upstream iter_move_grads_and_optimizers ( #98785 )
...
This PR upstreams `iter_move_grads_and_optimizer` which delay some of the gradients and the corresponding optimizer to the next iteration. D44512863(credit to @lessw2020 ) is the internal implementation, which is only good for the old _SPMD expansion. This PR changes the implmentation to use the new APIs.
Differential Revision: [D44836486](https://our.internmc.facebook.com/intern/diff/D44836486/ )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98785
Approved by: https://github.com/mrshenli
2023-04-19 06:40:33 +00:00
Chien-Chin Huang
d863876545
[SPMD] Remove the unused code ( #99075 )
...
Remove the unused code
Differential Revision: [D44973692](https://our.internmc.facebook.com/intern/diff/D44973692/ )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/99075
Approved by: https://github.com/lessw2020
2023-04-14 14:35:55 +00:00
Chien-Chin Huang
2de67eaaee
[SPMD] Add a dump_graphs_to_files utils to facilitate graph transformation debug ( #98284 )
...
Throughout the compilation, there are multiple graphs that will be generated. This PR add an utils to dump the result graphs to a folder.
Differential Revision: [D44661599](https://our.internmc.facebook.com/intern/diff/D44661599/ )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98284
Approved by: https://github.com/mrshenli
2023-04-11 23:14:12 +00:00
Chien-Chin Huang
e0b958f975
[SPMD] Allow IterGraph support a more general subgraph movement ( #98360 )
...
Resubmit D44444398 due to the merging conflict.
The original assumption of IterGraph is very restrictive and only allow users to move a subgraph that only one node has the input from external nodes. This PR fixes the limitation.
Differential Revision: [D44689730](https://our.internmc.facebook.com/intern/diff/D44689730/ )
**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D44689730/ )!
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98360
Approved by: https://github.com/lessw2020
2023-04-06 19:13:37 +00:00
Shen Li
02179827cb
[Easy] Include SPMD and DTensor files in UFMT checks ( #98148 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98148
Approved by: https://github.com/fegin
2023-04-02 15:34:49 +00:00
Kazuaki Ishizaki
35fd5c548e
Fix typos under torch/distributed directory ( #95638 )
...
This PR fixes typos in comments and messages of `.py` files under torch/distributed directory
Pull Request resolved: https://github.com/pytorch/pytorch/pull/95638
Approved by: https://github.com/usamah1 , https://github.com/H-Huang , https://github.com/kit1980
2023-03-27 21:13:44 +00:00
Shen Li
021de486ff
[Easy] Apply black to format _spmd files ( #97534 )
...
No real changes. Format code to prepare for the PR on top.
Differential Revision: [D44376380](https://our.internmc.facebook.com/intern/diff/D44376380 )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/97534
Approved by: https://github.com/wanchaol
2023-03-25 01:09:41 +00:00
Chien-Chin Huang
250c054bdd
[SPMD] Pull the minimal working distribute API and SPMD module to PyTorch ( #94802 )
...
Pull the minimal working distribute API and SPMD module to PyTorch. The original code is on https://github.com/pytorch/tau/tree/main/spmd/compiler .
Other main contributors to the original code base: @anj-s, @lessw2020, @wanchaol @aazzolini
Differential Revision: [D43197230](https://our.internmc.facebook.com/intern/diff/D43197230/ )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94802
Approved by: https://github.com/anj-s , https://github.com/wanchaol
2023-02-16 00:36:16 +00:00
