title: How SLURM and Pytorch handle multi-node multi-gpu training together
body: <p>I would like to ask how the gradients aggregate when being trained with multi-node multi-gpu in a cluster using Slurm to manage workload. The code is written using Pytorch.</p>
  
  <p>For example when launching a script <code>train.py</code> in Slurm to train a model on 4 nodes with 4GPUs per node as below, what do the <code>srun</code> command do exactly? </p>
  
  <pre><code>srun python train.py
  </code></pre>
  
  <p>It is written in the documentation that <code>srun</code> is used to run a parallel job on cluster managed by Slurm. However, I would like to know what steps it performs. I imagine that it copies the script to 16 GPUs and each GPU will run the script at the same time. Is that correct?</p>
  
  <p>If I use <code>torch.distributed.init_process_group()</code> in my script to handle multi-gpu training, how Slurm will handle the gradients collected from each GPU together with Pytorch? I assume that there is a master port (which is a GPU device in a node assigned by Slurm) that gathers the gradients. And all 16 GPUs in 4 nodes will transfer the gradients to this master port thanks to the <code>torch.distributed</code> package? Slurm helps to allocate resources but it does not have any role in the aggregation of the gradients?</p>
  
  <p>I'm also not sure if I should launch the script using just <code>srun</code> as above or should I specify the <code>torch.distributed.launch</code> in my command as below. I want to make sure the gradients are collected correctly.</p>
  
  <pre><code># NGPU equals to number of GPUs/node
  export NGPU=4
  srun python -m torch.distributed.launch --nproc_per_node=$NGPU train.py
  </code></pre>
  
  <p>Thank you in advance for your help!</p>