pytorch
125e9256 - [FSDP] Add back `forward_prefetch` (#85177)

Commit
3 years ago
[FSDP] Add back `forward_prefetch` (#85177) - This implements explicit forward prefetching following the static 1st iteration's pre-forward order when `forward_prefetch=True` in the FSDP constructor. - This has the same unit test coverage as the original `forward_prefetch`. - I checked via print statements that the prefetches are happening, but since I cannot get a good CPU bound workload, it is hard to tell via traces that the prefetch is working. Pull Request resolved: https://github.com/pytorch/pytorch/pull/85177 Approved by: https://github.com/zhaojuanmao
Author
Andrew Gu
Committer
Parents
Loading