PyTorch multi-GPU training, multi-node or single node, async or synced

[albertz]

For the PyTorch engine (#1120).

Multi-GPU training. Multi-node or single node. Async vs sync. It should also work well with consumer GPUs, also in the multi-node case. Currently distributed data parallel synchronized training, using DistributedDataParallel, and using our own dataset sampling logic, extending Dataset._get_default_random_seed_offset.

This issue is to track the process, discuss anything relevant.

[albertz]

So far we mostly follow the PyTorch DDP tutorial, which uses synchronized accumulated gradients.

[albertz]

One issue was that when wrapping the module with DistributedDataParallel, it was important that we would call that forward function. But this was not really compatible with our train_step function where we want to pass the original (unwrapped) module. However, we solve this by just setting up the right context, what DistributedDataParallel.forward does internally.