I would like to train A3C or distributed DQN on GPU with the new torch.distributed API. These algorithms boil down to writing GPU Hogwild training correctly. Papers like Elastic SGD 17 also needs async GPU code to reproduce.
I used to work with Tensorflow distributed mode, which has a whole collection of abstractions and wrappers to implement async training. https://www.tensorflow.org/deploy/distributed 34
A decent implementation includes parameter servers and high-level managers to take care of gradient communication, parameter syncing, and shared Adam/Adagrad optimizers, for example.
Unfortunately, I cannot find any official tutorials or example code that show how to write a basic GPU Hogwild with parameter servers. The torch.distributed primitives are too low-level to use correctly. I looked at the source code of torch.nn.parallel.DistributedDataParallel, hoping to get some inspiration. It’s also too involved for me to understand and rewrite for my use case.
I understand that the distributed mode is very new. I’d really appreciate it if anyone can give some guidance on how to emulate TF’s distributed semantics in pytorch. For example, what are the main steps and which MPI primitives should be used in each step? Ideally, I’d love to see some skeleton code. I can figure out the rest of the details by myself, but I need something to start with. Thanks in advance!