Andrew Gu
19c7df89fb
[FSDP()][20/N][Easy] Move functions in file ( #87932 )
...
This PR is easy. I just wanted to group functions in the file according to the same logical order.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87932
Approved by: https://github.com/mrshenli
2022-11-02 11:32:48 +00:00
Andrew Gu
4635f56da1
[FSDP()][18/N] Refactor pre_forward_unshard() ( #87931 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87931
Approved by: https://github.com/mrshenli
2022-11-02 11:32:45 +00:00
Andrew Gu
0a752688bd
[FSDP()][17/N] Refactor _fsdp_root_pre_forward() ( #87930 )
...
This PR moves `_fsdp_root_pre_forward()` to `_runtime_utils.py`.
Note: This PR includes a (temporary) fix for `NO_SHARD` + `CPUOffload(offload_params=True)`, where we set `non_blocking=False` when copying the gradient from device to host. It is only included in this PR since the test was **flaky** (but not consistently failing) on this PR , so I needed to fix to unblock land.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87930
Approved by: https://github.com/mrshenli
2022-11-02 11:32:42 +00:00
Andrew Gu
1f34067e9d
[FSDP()][16/N] Refactor post-forward/pre-backward ( #87929 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87929
Approved by: https://github.com/mrshenli
2022-11-01 17:26:03 +00:00
Andrew Gu
90c5f856b2
[FSDP()][14/N] Refactor pre-forward/post-backward ( #87927 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87927
Approved by: https://github.com/mrshenli
2022-11-01 17:25:59 +00:00
Andrew Gu
b1750d0440
[FSDP()][13/N] Refactor unshard/reshard/grads ( #87926 )
...
This PR is not too complicated. We just move unshard/reshard/grads out to `_runtime_utils.py` and make them take `state: _State` instead of `self`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87926
Approved by: https://github.com/mrshenli
2022-11-01 13:37:31 +00:00
Andrew Gu
cbc9faebfe
[FSDP()][1/N] Start refactoring FSDP root pre-forward ( #87915 )
...
Welcome! This PR starts the refactoring journey.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87915
Approved by: https://github.com/mrshenli
2022-10-29 06:50:30 +00:00
