[FSDP2] Added APIs for explicit fwd/bwd prefetching

[awgu]

Stack from ghstack (oldest at bottom):

• -> [FSDP2] Added APIs for explicit fwd/bwd prefetching #128884

This PR adds two APIs set_modules_to_forward_prefetch and set_modules_to_backward_prefetch to enable explicit forward/backward all-gather prefetching, respectively.

def set_modules_to_forward_prefetch(self, modules: List[FSDPModule]): -> None
def set_modules_to_backward_prefetch(self, modules: List[FSDPModule]): -> None

Motivation
FSDP2 implements reasonable defaults for forward and backward prefetching. In forward, it uses implicit prefetching and allows two all-gather output tensors to be alive at once (so that the current all-gather copy-out can overlap with the next all-gather). In backward, it uses explicit prefetching based on the reverse post-forward order.

However, there may be cases where with expert knowledge, we can reduce communication bubbles by moving all-gathers manually. One way to expose such behavior is to expose prefetching limits, i.e. integers that configure how many outstanding all-gathers/all-gather output tensors can be alive at once. IMIHO, this leans toward easy, not simple (see PyTorch design principles).

The crux of the problem is that there may be special cases where manual intervention can give better performance. Exposing a prefetching limit and allowing users to pass a value >1 just smooths over the problem since such a limit would generally apply over the entire model even though it possibly should not. Then, expert users will see a specific all-gather that they want to deviate from this limit, and there is little we can do.

Thus, we instead choose to expose the most primitive extension point: namely, every FSDPModule gives an opportunity to prefetch other all-gathers in forward and in backward. How to leverage this extension point is fully up to the user. Implementing the prefetch limit can be done using this extension point (e.g. record the post-forward order yourself using forward hooks, iterate over that order, and call the set_modules_to_forward_prefetch / set_modules_to_backward_prefetch APIs).

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k

Differential Revision: D58700346

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128884

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 25fff88 with merge base 24443fe ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[awgu]

Example of prefetching 1 FSDP module explicitly in backward:

Example of prefetching 2 FSDP modules explicitly in backward:

[awgu]

Example of prefetching 1 FSDP module explicitly in forward:

Example of prefetching 2 FSDP module explicitly in forward:

Example of prefetching 4 FSDP module explicitly in forward:

[ckluk2]

This is great!

[awgu]

@awgu has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[sanketpurandare]

Curious to know the use case. This will only help in case of FSDPModules that have unbalanced compute. So that we can use the compute of a larger FSDPModule to overlap the prefetch of smaller ones.

Also, will your extension of wrapping List[nn.Module] eliminate the need for this?

[awgu]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here