I created a wrapper to work like executing python -m torch.distributed.launch --nproc_per_node 2 ./my_script.py but from my script. I do call trains.init in the subprocesses, I the actually difference between the subproceses supposed to be, in terms or arguments, local_rank that's all.It may be possible and that I'm not distributing the model between the GPUs in an optimal way or at least in a way that matches your framework.
If you have an example it would be great.