Enhance Distributed Fused Adam

NVIDIA / apex

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

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

8.42k stars 1.4k forks source link

Closed alpha0422 closed 2 months ago

alpha0422 commented 3 months ago

This PR enhances distributed fused adam by:

Support NHWC layout (required by some Conv related models, e.g. Diffusion models);
Fix the gradient clipping bug;
Support CUDA graph;

@timmoon10 @crcrpar Please help review, thanks.