pytorch
dfdf4e79 - Added serialization to postlocal_SGD. (#80435)

Commit
2 years ago
Added serialization to postlocal_SGD. (#80435) Fixes #75666 Current PR adds the functionality for `PostLocalSGD` communication hook and tests that communication hook can be properly saved and restored. Similar to https://github.com/pytorch/pytorch/pull/79334, where serialization was added to `PowerSGD`. ``__getstate__`` Returns: ``` ``Dict[str, Any]`` which will be pickled and saved. ``process_group`` and ``subgroup`` are not serializable and excluded from a returned state. ``` ``__setstate__`` ``` Takes provided ``state`` and retrieves ``PostLocalSGDState``. ``process_group`` and ``subgroup`` are set to default process_group and subgroup respectively. Default subgroup is equivalent to the subgroup on each node. ``` Small adjustment to `PowerSGD`'s warning message. Refactored unittest, i.e. separated parity and log checks. Pull Request resolved: https://github.com/pytorch/pytorch/pull/80435 Approved by: https://github.com/awgu
Author
Olga Andreeva
Committer
Parents
Loading