torchrun support

[hahahannes]

I am trying to run MLPF with Torchrun for multi node multi GPU training and I needed to make some changes. I am following the guide at https://pytorch.org/tutorials/beginner/ddp_series_fault_tolerance.html.

Let me know if that might interesting for you. I am tracking my changes in my fork. There is one commit with the changes so far. I added an argument use-torchrun but I think torchrun could also replace the mp.spawn part

[jpata]

Thanks, we didn't really investigate multi-node running so far AFAIK, as we don't generally have such nodes available for testing.
Maybe @erwulff did with Ray.
But it's good to support it if it doesn't break other things or make things too complicated.

[hahahannes]

Alright! Just for some background: I am using it to validate our ml platfom on ml.cern.ch. We have an abstraction for running torchrun to run multi node pytorch workloads using KubeFlow PytorchJob and I also use it to run the training with Ray Train on top of KubeFlow.

[erwulff]

I have indeed done some basic tests of multi-node training using Ray Train. It works, but has not been used extensively.