Add image for better explanation to FSDP tutorial #2644
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/2644
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 077d1c0 with merge base f05f050 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
github-actions
bot
added
distributed
medium
docathon-h2-2023
and removed
cla signed
labels
| @@ -46,6 +46,15 @@ At a high level FSDP works as follow: | |||
| * Run reduce_scatter to sync gradients | |||
| * Discard parameters. | |||
|
|
|||
| The key insight behind full parameter sharding is that we can decompose the all-reduce operations in DDP into separate reduce-scatter and all-gather operations. | |||
There was a problem hiding this comment.
I am not sure that this is the correct statement.
Even though an all-reduce can be decomposed as a reduce-scatter and all-gather, the current phrasing might suggest that DDP's gradient all-reduce is being decomposed into a gradient reduce-scatter and gradient all-gather. However, FSDP actually all-gathers parameters.
Whether or not this decomposition of all-reduce into reduce-scatter and all-gather is the key insight is not obvious to me. If we show this decomposition, we probably want more exposition.
There was a problem hiding this comment.
I agree. Adding this picture would demand a clear explanation.
I am unsure what to write. If you can suggest something or direct me to where I can read about this topic, that'd be very helpful.
There was a problem hiding this comment.
Maybe something like the following:
One way to view FSDP's sharding is to decompose the DDP gradient all-reduce into reduce-scatter and all-gather. In particular, FSDP reduce-scatters gradients such that each rank has a shard of the gradients in backward, updates the corresponding shard of the parameters in the optimizer step, and all-gathers them in the next forward.
< Figure >
There was a problem hiding this comment.
Sounds good. Thanks
Co-authored-by: Andrew Gu <31054793+awgu@users.noreply.github.com>
NicolasHug
changed the title
Fixes #2613
Description
The tutorial lacked an explanation for what's going on behind parameter sharding
Checklist
cc @wconstab @osalpekar @H-Huang @kwen2501 @sekyondaMeta @svekars @carljparker @NicolasHug @kit1980 @subramen