pytorch
a95889ba - [FSDP] Add initial `summon_full_params(with_grads=True)` (#85738)

Commit
2 years ago
[FSDP] Add initial `summon_full_params(with_grads=True)` (#85738) This adds `summon_full_params(with_grads=True)` for `use_orig_params=True` and `offload_to_cpu=False`. Filling in the `use_orig_params=False` case requires some already-planned refactoring, and the `offload_to_cpu=True` case needs some additional work as well. Adding this is helpful for debugging `use_orig_params=True` to make sure gradients are being updated correctly. Pull Request resolved: https://github.com/pytorch/pytorch/pull/85738 Approved by: https://github.com/rohan-varma
Author
Committer
Parents
Loading