Seeking suggestions for embedding into ddp

Hi, PyTorch 1.8 have this new hook torch.nn.parallel.DistributedDataParallel.register_comm_hook(), any advices on how to integrate grace into ddp using the dist examples?

[hangxu0304]

dist.all_reduce simply sums up the dense inputs across the nodes. It doesn’t support value-index pair sparse tensor format. To perform allreduce for sparse tensors, you need to use allgather to collect all sparse tensors and cast them into dense format locally, then sum them up. In grace, we have developed allgather for tensors with different length, please check here.

[hangxu0304]

Thank you very much. Your implementation is really good. I do have the following 2 suggestions for your optimization:

• In DgcCompressor.decompress(), there is a For Loop to cast the sparse tensor into dense format, which can be very expensive in case of large gradients. You may want to use this api scatter_add to gain some speed.
• Index and value allgather operations can be merged since you already concat them after DGC compression.

Regarding the poor performance, are you comparing with GRACE DGC or the no compression baseline? And also please note that, gradient compression is not always beneficial due to various model architectures, network conditions, and different number of nodes. Could you please specify your testing environment?