Add type hints to distributed Adam optimizer

NVIDIA / apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch

BSD 3-Clause "New" or "Revised" License

8.2k stars 1.36k forks source link

Closed timmoon10 closed 11 months ago

timmoon10 commented 12 months ago

This PR makes some minor stylistic changes to the distributed Adam optimizer:

Add type hints for functions and member variables.
Remove unused functions, mainly the deprecated synchronization logic for async collectives in GradientBucket.sync_wait and ParameterBucket.sync_wait,
Make the ParameterFragment class a dataclass.

None of these changes should affect functionality.