Torch1.10.2 is slower than torch1.9.1

Hukongtao: <p>I recently upgraded my pytorch version from 1.9.1 to 1.10.2 due to project needs. But found that the training speed dropped by 20% to 30% for the same multi-tasks model.</p>
<p>After profiling, I found the reason is that the code of torch1.10 has the following logic:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/v1.10.2/torch/nn/parallel/distributed.py#L1378">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/v1.10.2/torch/nn/parallel/distributed.py#L1378" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/v1.10.2/torch/nn/parallel/distributed.py#L1378" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/v1.10.2/torch/nn/parallel/distributed.py#L1378</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="1368" style="counter-reset: li-counter 1367 ;">
          <li>            # already have been joined and have stale module buffers.</li>
          <li>            if self._join_config.enable:</li>
          <li>                authoritative_rank = self._find_common_rank(</li>
          <li>                    self._distributed_rank, True</li>
          <li>                )</li>
          <li>            else:</li>
          <li>                # The process with rank 0 is considered the authoritative copy.</li>
          <li>                authoritative_rank = 0</li>
          <li>            # Update self.modules_buffers incase any buffers were</li>
          <li>            # reassigned.</li>
          <li class="selected">            self._assign_modules_buffers()</li>
          <li>            self._distributed_broadcast_coalesced(</li>
          <li>                self.modules_buffers,</li>
          <li>                self.broadcast_bucket_size,</li>
          <li>                authoritative_rank,</li>
          <li>            )</li>
          <li>
          </li>
<li>def _passing_sync_batchnorm_handle(self, module):</li>
          <li>    for layer in module.modules():</li>
          <li>        if isinstance(layer, torch.nn.modules.SyncBatchNorm):</li>
          <li>            if self.device_type == "cpu":</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>When using DDP, the model will assign module buffers before each forward.Time will be spent in the following sections:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/71f889c7d265b9636b93ede9d651c0a9c4bee191/torch/nn/parallel/distributed.py#L752">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/71f889c7d265b9636b93ede9d651c0a9c4bee191/torch/nn/parallel/distributed.py#L752" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/71f889c7d265b9636b93ede9d651c0a9c4bee191/torch/nn/parallel/distributed.py#L752" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/71f889c7d265b9636b93ede9d651c0a9c4bee191/torch/nn/parallel/distributed.py#L752</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="742" style="counter-reset: li-counter 741 ;">
          <li>    """</li>
          <li>    Assigns module buffers to self.modules_buffers which are then used to</li>
          <li>    broadcast across ranks when broadcast_buffers=True. Note that this</li>
          <li>    must be called every time buffers need to be synced because buffers can</li>
          <li>    be reassigned by user module,</li>
          <li>    see https://github.com/pytorch/pytorch/issues/63916.</li>
          <li>    """</li>
          <li>    # Collect buffers for modules, filtering out buffers that should be ignored.</li>
          <li>    named_module_buffers = [</li>
          <li>        (buffer, buffer_name)</li>
          <li class="selected">        for buffer_name, buffer in self.module.named_buffers()</li>
          <li>    ]</li>
          <li>    self.modules_buffers = [</li>
          <li>        buffer</li>
          <li>        for (buffer, buffer_name) in named_module_buffers</li>
          <li>        if buffer_name not in self.parameters_to_ignore</li>
          <li>    ]</li>
          <li>
          </li>
<li>def _build_param_to_name_mapping(self, parameters):</li>
          <li>    param_to_param_index = {parameters[0][i]: i for i in range(len(parameters[0]))}</li>
          <li>    param_set = set(parameters[0])</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>I would like to know how I can avoid this assign of buffers?</p>