mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
1c37119a1f
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 |
||
|---|---|---|
| .. | ||
| _checkpoint | ||
| _comm_hooks | ||
| _optimizer_overlap | ||
| _quantization | ||
| ddp_comm_hooks | ||
| model_averaging | ||
| __init__.py | ||
| join.py | ||