pytorch

1c37119a - [FSDP] New fix for composing with other module wrappers (#87950)

[FSDP] New fix for composing with other module wrappers (#87950) We change `.module` to pass through `ActivationWrapper` directly to the inner wrapped module. This should fix the state dict issues. Given the invariant that `.module` always returns the inner wrapped module, FSDP always registers the `FlatParameter` on the inner wrapped module, regardless of if there is an intermediate `ActivationWrapper` or not. This avoids casing on whether `ActivationWrapper` is added before or after FSDP construction. This PR removes the added unit test in `test_fsdp_misc.py` for changing the wrapped module because I would rather not complicated `_lazy_init()` logic just to support that kind of adversarial behavior. The user should not be swapping out the wrapped module arbitrarily or deleting the `FlatParameter`. I mainly had those tests to make sure that all branches of the code I added was correct. Differential Revision: [D40799961](https://our.internmc.facebook.com/intern/diff/D40799961) Pull Request resolved: https://github.com/pytorch/pytorch/pull/87950 Approved by: https://github.com/zhaojuanmao