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.

[jpata]

@hahahannes I wanted to circle back on this.
Did you get multi-node torchrun to work properly?
I see you used gloo instead of nccl, wouldn't that be very slow / inoptimal? Was that choice specifically for the CERN infra?

Why I'm asking is because I'm now interested to do multi-node training myself on AMD GPUs at the LUMI HPC, and Ray is not that well supported by AMD apparently.

[hahahannes]

Hi @jpata, yes I got it running but making a few adjustments. This was already some time ago. I made some commits in my fork at https://github.com/hahahannes/particleflow/tree/pytorch-ddp. If you want, I can open a cleaned up PR. I think I ran into a problem with the network interfaces with nccl. But this might have been very specific to the hardware I was using back then.

[jpata]

I saw your fork, I think I can try the relevant code from there, if I know it worked. I want to do some refactoring on my end as well for this. Do you have any set of slides summarizing the results? Was it what you showed at the OpenLab workshop in the winter?