[FSDP2] Used multi-grad hook when no inputs require grad

[awgu]

Stack from ghstack (oldest at bottom):

• -> [FSDP2] Used multi-grad hook when no inputs require grad #129259
• Added register_multi_post_accumulate_grad_hook #131949

This PR uses register_multi_post_accumuluate_grad_hook to run the post-backward logic when the module inputs do not require gradient.

For context, FSDP2 currently relies on a hook on the module inputs that require gradient to run the post-backward logic (like a module full backward hook except with pytree support). This means that if none of the module inputs require gradient, then the post-backward logic is deferred to the final callback that runs at the end of backward, which may not be timely in some hybrid parallelism cases (e.g. a sparse arch before the FSDP dense arch). To address this case, we use this new multi-post-accumulate-grad hook.

Since whether the module inputs requiring grad could be dynamic from iteration to iteration, we guard this on a flag to make sure that we use the existing logic if the module inputs do require gradient.

Common Case

In the common case, the only module whose forward inputs do not require gradient is the root FSDP module, where the root FSDP module is the overall model's root. In this case, the multi-grad hook simply moves the root's post-backward from the final callback to a preceding AccumulateGrad node.

New:

Old:

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

Differential Revision: D59012616

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

✅ You can merge normally! (2 Unrelated Failures)

As of commit a05c807 with merge base 2820e1d ():

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 2, 3, linux.8xlarge.nvidia.gpu) (gh) (related job)
  distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_nested_fully_shard_fullgraph_backend_aot_eager
• trunk / linux-focal-cuda11.8-py3.10-gcc9-experimental-split-build-test / test (distributed, 2, 3, linux.8xlarge.nvidia.gpu) (gh) (#129539)
  distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_nested_fully_shard_fullgraph_backend_aot_eager

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

[awgu]

For context, the existing multi-grad hook does not work so well with the existing post_backward because the last parameter's .grad is not assigned yet when the mulit-grad hook runs (expected behavior).

We actually prefer this mulit-post-accumulate-grad hook, where the last parameter's .grad has already been assigned.

[awgu]

The key idea is that before, we just return when there are no module inputs that require grad, falling back to the final callback that runs at the end of backward to run post-backward for this module.

This multi-post-acc-grad hook allows running the module's post-backward earlier (but still correctly).

[sanketpurandare]

This is neat!

[awgu]

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

[sanketpurandare]

This is very clean! Kudos!