Should the parameter nproc_per_node be equal on two different GPU nodes

slyviacassell: <p>I have two GPU nodes. One has two GPUs and the other has only one GPU. I want to use them for distributed training and I run with this bash code:<br>
Node 1</p>
<pre><code class="lang-bash">python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr="0.0.0.0" --master_port=3338 train_dist.py  --restore 0 --config-file configs/vgg16_nddr_additive_4_unpool_aug_shortcut_sing_cosine_dist.yaml
</code></pre>
<p>Node 2</p>
<pre><code class="lang-bash">python -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr="0.0.0.0" --master_port=3338 train_dist.py  --restore 0 --config-file configs/vgg16_nddr_additive_4_unpool_aug_shortcut_sing_cosine_dist.yaml
</code></pre>
<p>But the program is stuck. And I have no idea about it. But when the <code>nproc_per_node</code> is set to 1 on both of them there is no problem. So, should the num of GPU on each distributed node always be the same? Does there have any other solutions to run the unbalanced distributed training?</p>

pritamdamania87: <p>This is currently a limitation of <code>torch.distributed.launch</code> where it assumes all nodes are symmetric. Basically, on each node it assumes the world_size is <code>nproc_per_node * nnodes</code> and as a result you see the hang since this is not consistent across all nodes.</p>

slyviacassell: <p>Thanks a lot! I successfully solve it with <code>mutliporcessing.spawn</code>. Also, there is another way which means we need to re-implement <code>launch.py</code>.</p>

logicShu: <p>Perhaps this answer should be revised?<br>
I read the source code and found that this has now been changed.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/master/torch/distributed/elastic/agent/server/api.py#L580">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/master/torch/distributed/elastic/agent/server/api.py#L580" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/distributed/elastic/agent/server/api.py#L580" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/distributed/elastic/agent/server/api.py#L580</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="570" style="counter-reset: li-counter 569 ;">
          <li>    start_idx: int = 0,</li>
          <li>    end_idx: int = -1,</li>
          <li>) -&gt; Tuple[int, List[int]]:</li>
          <li>    if end_idx == -1:</li>
          <li>        end_idx = len(role_infos)</li>
          <li>    prefix_sum = 0</li>
          <li>    total_sum = 0</li>
          <li>    for idx in range(start_idx, end_idx):</li>
          <li>        if role_idx &gt; idx:</li>
          <li>            prefix_sum += role_infos[idx].local_world_size</li>
          <li class="selected">        total_sum += role_infos[idx].local_world_size</li>
          <li>    return (</li>
          <li>        total_sum,</li>
          <li>        list(range(prefix_sum, prefix_sum + role_infos[role_idx].local_world_size)),</li>
          <li>    )</li>
          <li>
          </li>
<li># pyre-fixme[56]: Pyre was not able to infer the type of the decorator</li>
          <li>#  `torch.distributed.elastic.metrics.prof`.</li>
          <li>@prof</li>
          <li>def _assign_worker_ranks(</li>
          <li>    self, store, group_rank: int, group_world_size: int, spec: WorkerSpec</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>