Using Torch.utils.checkpoint.checkpoint with DataParallel

JWarlock: <p>Hi,<br>
I’m trying to figure out how to use <code>Torch.utils.checkpoint.checkpoint</code> to save memory consumption with multi-gpu training (<code>nn.DataParallel</code> or <code>nn.parallel.DistributedDataParallel</code>)</p>
<p>But it seem that using checkpointing and multi-gpu training simultaneously greatly slows down training speed. See:</p><aside class="quote" data-post="1" data-topic="22837">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/f475e1/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/torch-utils-checkpoint-checkpoint-with-multi-gpus/22837">Torch.utils.checkpoint.checkpoint with multi-gpus</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi, 
The program is blocked when I use Torch.utils.checkpoint.checkpoint with multi-GPUs. And when I change to single GPU training, the program works! I also try to remove the checkpoint function and the program also works with multi-GPUs. Does checkpoint function not support multi-GPUs training? 
Thanks!
  </blockquote>
</aside>
<p>
and</p><aside class="onebox githubissue" data-onebox-src="https://github.com/gpleiss/efficient_densenet_pytorch/issues/36">
  <header class="source">

      <a href="https://github.com/gpleiss/efficient_densenet_pytorch/issues/36" target="_blank" rel="noopener nofollow ugc">github.com/gpleiss/efficient_densenet_pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/gpleiss/efficient_densenet_pytorch/issues/36" target="_blank" rel="noopener nofollow ugc">MultiGPU efficient densenets are slow</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2018-04-30" data-time="17:05:23" data-timezone="UTC">05:05PM - 30 Apr 18 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/wandering007" target="_blank" rel="noopener nofollow ugc">
          <img alt="wandering007" src="https://avatars.githubusercontent.com/u/6906411?v=4" class="onebox-avatar-inline" width="20" height="20">
          wandering007
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">I just want to benchmark the new implementation of efficient densenet with the c<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">ode [here](https://gist.github.com/wandering007/d7f3db11285f382fefa39a9cda99ad0d). However, it seems that the used checkpointed modules are not broadcast to multiple GPUs as I got the following errors:

```
  File "/home/changmao/efficient_densenet_pytorch/models/densenet.py", line 16, in bn_function
    bottleneck_output = conv(relu(norm(concated_features)))
  File "/home/changmao/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/changmao/anaconda3/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py", line 49, in forward
    self.training or not self.track_running_stats, self.momentum, self.eps)
  File "/home/changmao/anaconda3/lib/python3.6/site-packages/torch/nn/functional.py", line 1194, in batch_norm
    training, momentum, eps, torch.backends.cudnn.enabled
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_batch_norm)
```
I think that the checkpoint feature provides weak support for `nn.DataParallel`.</span></p>
  </div>

  </article>

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

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

<p>I’ve found an unofficial implementation of gradient checkpointing for DenseNet which works fine for multi-gpu: <a href="https://github.com/csrhddlam/pytorch-checkpoint" class="inline-onebox" rel="noopener nofollow ugc">GitHub - csrhddlam/pytorch-checkpoint</a>. But I wonder is there any efficient way to use checkpointing and multi-gpu training simultaneously in the official pytorch implementation.</p>
<p>Thanks!</p>

mrshenli: <p><code>DistributedDataParallel</code> currently does not work with checkpoint, see this discussion: <a href="https://github.com/pytorch/pytorch/issues/24005" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/24005</a></p>
<p>This PR (<a href="https://github.com/pytorch/pytorch/pull/35083" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/35083</a>) fixes the error, but as under the new <code>delay_allreduce</code> mode, communication and computation no longer overlap, it is expected to be slower.</p>
<p><code>DataParallel</code> might work with checkpoint, have you tried that?</p>