[FSDP] Allow MixedPrecision to skip inputs

[mrshenli]

Stack from ghstack (oldest at bottom):

• [FSDP] Fix _pre_forward type annotation #90621
• -> [FSDP] Allow MixedPrecision to skip inputs #90620
• Add global registry to composable API contract #90579
• [FSDP] Allow nested FSDP wrapper to use different mixed precision #90523

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 our office hours

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

❌ 1 Failures

As of commit 85a007b:

The following jobs have failed:

• linux-focal-rocm5.2-py3.8 / test (default, 1, 2, linux.rocm.gpu)

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

[rohan-varma]

LGTM, but curious what's the use case to skip inputs?

[mrshenli]

Thanks @rohan-varma, one model I am working on has one forward argument, which is sensitive to precision. So have to keep it in fp32, instead of converting it to bfloat16 and convert back.

[awgu]

Overall, this looks good to me. Sorry for leaving all of the comments in separate reviews instead of just one.

[awgu]

My bad for suggesting to un-default cast_forward_inputs. I did not see that it already existed in past tests.

I think you need to add a value for cast_forward_inputs in this line:

pytorch/test/distributed/fsdp/test_fsdp_mixed_precision.py

Line 821 in a5d34ad

model = SaveForwardInputsModel(forward_inputs).cuda()

[facebook-github-bot]

This pull request has been merged in 80542ad.