Default bucket_cap_mb param cause training hang for large model?

lzhbrian: <p>I am not sure whether this is a bug of ddp training.<br>
When using the default value for bucket_cap_mb, my model runs fluently on 1 gpu, but got stuck (hanging at loss.backward()) when using multiple gpus.</p>
<p>However, when I set bucket_cap_mb=50, the training runs successfully on multiple gpus.</p>
<p>Another solution is keeping bucket_cap_mb unchanged, while using a smaller model. In this way, the training also succeeds.</p>
<p>Is this expected? It seems like when using a larger model, we need a larger bucket_cap_mb to make it train correctly?</p>

mrshenli: <blockquote>
<p>Is this expected? It seems like when using a larger model, we need a larger bucket_cap_mb to make it train correctly?</p>
</blockquote>
<p>hmm, I’d be surprised if this is the case. DDP makes sure that all ranks launch the same set of allreduce in the same order.</p>
<p>Besides DDP, does any part of your model/program also launches collective communications? Asking because the symptom looks like there is another collective fired somewhere during the backward, and caused desync problem when combined together with DDP’s allreduce.</p>

lzhbrian: <p>Thanks for the reply.<br>
As you mentioned, I did launch some communications elsewhere.</p>
<p>I am training a multi-task learning model, and I want to log each heads’ loss collected from all gpus. So before loss.backward(), I firstly dist.all_reduce() the losses and also some stats from dataloader.</p>
<p>My training step looks like this:</p>
<pre><code class="lang-python">def dist_reduce_sum(tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    return rt

def train_one_iter(inp, gt):

    optimizer.zero_grad()

    # forward model
    output = model(inp)

    # calc loss for each head
    l1 = loss_head1(output[head1], gt[head1])
    l2 = loss_head2(output[head2], gt[head2])
    l3 = loss_head3(output[head3], gt[head3])

    # barrier
    torch.distributed.barrier()

    # this is for logging (collect each head loss from all gpu, and also some stats from dataloader)
    stat1 = dist_reduce_sum(inp[head1][some_stat])
    stat2 = dist_reduce_sum(inp[head2][some_stat])
    stat3 = dist_reduce_sum(inp[head3][some_stat])
    l1_all = dist_reduce_sum(l1) / stat1
    l2_all = dist_reduce_sum(l2) / stat2
    l3_all = dist_reduce_sum(l3) / stat3
    logging.info(l1_all, l2_all, l3_all)

    # backward
    total_loss = l1 + l2 + l3
    total_loss.backward() # !! this is where it hangs !!

    # step
    optimizer.step()

</code></pre>
<p>Is my above training procedure problemetic for DDP?<br>
Thanks for the help again ! <a class="mention" href="/u/mrshenli">@mrshenli</a></p>

lzhbrian: <p>btw, I use timm Regnet800M model and convert model to syncbn using their func here.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/rwightman/pytorch-image-models/blob/4f72bae43be26d9764a08d83b88f8bd4ec3dbe43/timm/models/layers/norm_act.py#L130">
  <header class="source">

      <a href="https://github.com/rwightman/pytorch-image-models/blob/4f72bae43be26d9764a08d83b88f8bd4ec3dbe43/timm/models/layers/norm_act.py#L130" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/rwightman/pytorch-image-models/blob/4f72bae43be26d9764a08d83b88f8bd4ec3dbe43/timm/models/layers/norm_act.py#L130" target="_blank" rel="noopener nofollow ugc">rwightman/pytorch-image-models/blob/4f72bae43be26d9764a08d83b88f8bd4ec3dbe43/timm/models/layers/norm_act.py#L130</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="120" style="counter-reset: li-counter 119 ;">
          <li>    # Do not create this module directly or use the PyTorch conversion function.</li>
          <li>    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:</li>
          <li>        x = super().forward(x)  # SyncBN doesn't work with torchscript anyways, so this is fine</li>
          <li>        if hasattr(self, "drop"):</li>
          <li>            x = self.drop(x)</li>
          <li>        if hasattr(self, "act"):</li>
          <li>            x = self.act(x)</li>
          <li>        return x</li>
          <li>
          </li>
<li>
          </li>
<li class="selected">def convert_sync_batchnorm(module, process_group=None):</li>
          <li>    # convert both BatchNorm and BatchNormAct layers to Synchronized variants</li>
          <li>    module_output = module</li>
          <li>    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):</li>
          <li>        if isinstance(module, BatchNormAct2d):</li>
          <li>            # convert timm norm + act layer</li>
          <li>            module_output = SyncBatchNormAct(</li>
          <li>                module.num_features,</li>
          <li>                module.eps,</li>
          <li>                module.momentum,</li>
          <li>                module.affine,</li>
      </ol>
    </code></pre>



  </article>

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

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


mrshenli: <p>From <a class="mention" href="/u/yanli_zhao">@Yanli_Zhao</a>,</p>
<p>Hey <a class="mention" href="/u/lzhbrian">@lzhbrian</a>, which PyTorch version are you using? For some old PyTorch releases, this desync might be caused by DDP rebuild bucket logic.</p>

lzhbrian: <p>I am using:</p>
<p>torch==1.10.0+cu111<br>
torchvision==0.11.0+cu111<br>
python==3.8.12<br>
nvidia driver: 460.27.04<br>
cuda: cuda_11.1.TC455_06.29190527_0</p>
<p>NCCL is used as backend for DDP.</p>
<p>GPUs are 8xA100 (using only 4 of them) running in a docker container with Ubuntu 18.04.6 LTS.</p>
<hr>
<p>I tried another env, the training also hangs:<br>
pytorch                   1.12.1          py3.9_cuda11.3_cudnn8.3.2_0<br>
torchvision               0.13.1               py39_cu113    pytorch<br>
cudatoolkit               11.3.1               h2bc3f7f_2    defaults</p>

lzhbrian: <p>Another observation update:</p>
<p>If I disabled syncbn, the training also succeeded …</p>

lzhbrian: <p>Hi <a class="mention" href="/u/mrshenli">@mrshenli</a> , any update on this ?</p>
<p>Thanks!</p>