@teng-li, @csarofeen has a simple distributed module that supports mixed precision https://github.com/csarofeen/examples/blob/dist_fp16/imagenet/distributed.py, he was planning to push it to pytorch core soon. It is just a few lines of code, much easier to follow than DistributedDataParallel. Like yours here, it all-reduces all the parameters at the end of the iteration, and unlike DistributedDataParallel it does not have the limitation that all the parameters have to be updated at each iteration. It does not sync buffers and you are right, it is a good idea to have an option to sync or not sync buffers. We also tested it with resnet on DGX1.
And thank you for your work with nccl backend!

@ngimel Thanks for the previous code review. Yes, I did see @csarofeen 's DDP code and it was a lot simpler. But our current DistributedDataParallel model has also covered other important use cases that can be beneficial with other backends such as gloo. And with the current DDP implementation, if we use single GPU binding and launch a DDP process per GPU, I didn't see much perf degradation by binding a single GPU to our current DDP implementation with multiple processes compared to Christian's DDP either. Please see

#4870

for the perf comparisons with Christian's simpler version of DDP. Also the current DDP implementation also covers the single process multiple GPU case, which is another use case for users other than multi-process use case. So I guess it would probably be Ok to also add this mixed-precision support into the current DDP so that other use cases (such as gloo on ethernet with perf advantage due to the multi-threading reduction) can be covered as well.

CC: @apaszke @soumith

Since the single GPU binding will go to a different small code path as I planned. I think this PR should be OK to proceed forward. @apaszke

@apaszke Comments addressed.

Fixed merge conflict. @apaszke Please also take a look at this PR.

@apaszke refactored