Is there a way to selectively load model params into GPU during forward pass?

jerinphilip: <p>I came across some way to change the GPU for component modules of a larger model in the following link:</p>
<p><a href="https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html</a></p>
<p>However, my model runs in a distributed environment (ddp) and have several sub-modules holding all of which for a forward-backward cycle each of which will decrease the size of batch I can fit for an update. But for each forward/backward, not all of these are active and therefore are not required to stay in GPU (as in I can load these for a single-batch) and otherwise leave it in the non-gpu memory.</p>
<p>In other words, I have parameters <code>Theta + theta[t]</code> (for t=1…T), where t is a particular task. I want to only load a single <code>theta[t]</code> for a forward and backward pass into the GPU and fit larger batches. Currently I’m holding all <code>theta[t]</code> in the GPU.</p>
<p>Is it possible to use the same semantics if it’s the same (sub)-module (<code>theta[t]</code>) to achieve the intention described above?</p>

mrshenli: <p>Hey <a class="mention" href="/u/jerinphilip">@jerinphilip</a>, I believe this is possible. You can use <code>Tensor.to(device)</code> to move the parameters to the GPUs in the forward pass, and the <code>to</code> (i.e., copy) operator should be added into the autograd graph, so that the backward pass will compute gradients for the original on-CPU parameters properly. Let me know if it didn’t work.</p>
<p>Note that, although this can reduce the footprint on GPU memory, DDP would still need to communicate the same amount of parameters, as that is determined at DDP construction time. And as those parameters are on CPU, you won’t be able to use NCCL which might cause considerable slow down.</p>

jerinphilip: <p>Where can I read more on the DDP communications setup? Thanks in advance.</p>

mrshenli: <p>Hey <a class="mention" href="/u/jerinphilip">@jerinphilip</a>, this page briefly describes DDP: <a href="https://pytorch.org/docs/master/notes/ddp.html" rel="nofollow noopener">https://pytorch.org/docs/master/notes/ddp.html</a></p>
<p>We have a paper with more comprehensive details. Let me upload that to archive.</p>

jerinphilip: <aside class="quote no-group" data-username="mrshenli" data-post="2" data-topic="85501">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mrshenli/48/12220_2.png" class="avatar"> mrshenli:</div>
<blockquote>
<p>Note that, although this can reduce the footprint on GPU memory, DDP would still need to communicate the same amount of parameters, as that is determined at DDP construction time. And as those parameters are on CPU, you won’t be able to use NCCL which might cause considerable slow down.</p>
</blockquote>
</aside>
<p>Where do I obtain details corresponding to this particular information? Isn’t only <code>.grad</code> meant to be communicated and the workers applying the updates individually? If my parameters of  <code>theta[t]</code> has only gradients for the particular task, would this help the case? I’m reading the <strong>Forward Pass</strong> section of <a href="https://pytorch.org/docs/master/notes/ddp.html#internal-design" rel="noopener nofollow ugc">Internal Design</a>, with <code>find_unused_parameters</code>, it is possible to operate on a subgraph, correct(?). I already have this enabled.</p>

mrshenli: <blockquote>
<p>Where do I obtain details corresponding to this particular information?</p>
</blockquote>
<p>We need to go through some internal approval process to publicly disclose that paper. It will take some time. For now  <a href="https://pytorch.org/docs/master/notes/ddp.html" class="inline-onebox" rel="noopener nofollow ugc">Distributed Data Parallel — PyTorch master documentation</a> is the best place for overall intro. The implementation of DDP is linked below:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/nn/parallel/distributed.py">
  <header class="source">

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

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


      <pre><code class="lang-py">from contextlib import contextmanager
import copy
import itertools

import torch

import torch.cuda.comm
import torch.distributed as dist

if dist.is_available():
    from torch.distributed.distributed_c10d import _get_default_group

from ..modules import Module
from .replicate import replicate
from .scatter_gather import scatter_kwargs, gather
from .parallel_apply import parallel_apply
from torch.cuda._utils import _get_device_index


def _find_tensors(obj):
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/nn/parallel/distributed.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/csrc/distributed/c10d/reducer.cpp">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/csrc/distributed/c10d/reducer.cpp" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/csrc/distributed/c10d/reducer.cpp" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/csrc/distributed/c10d/reducer.cpp</a></h4>


      <pre><code class="lang-cpp">#include &lt;torch/csrc/distributed/c10d/reducer.h&gt;

#include &lt;functional&gt;

#include &lt;c10/core/DeviceGuard.h&gt;
#include &lt;c10/util/Exception.h&gt;
#include &lt;torch/csrc/autograd/engine.h&gt;
#include &lt;torch/csrc/autograd/function_hook.h&gt;
#include &lt;torch/csrc/autograd/functions/accumulate_grad.h&gt;
#include &lt;torch/csrc/autograd/profiler.h&gt;
#include &lt;torch/csrc/autograd/utils/lambda_post_hook.h&gt;
#include &lt;torch/csrc/distributed/c10d/comm.h&gt;
#include &lt;torch/csrc/utils/hash.h&gt;
#include &lt;torch/csrc/utils/memory.h&gt;

namespace c10d {
namespace {

inline int64_t current_time_in_nanos() {
  return torch::autograd::profiler::getTime();
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/ebd869153c6adb37507d2ecb6a9fe3fd495fbb6e/torch/csrc/distributed/c10d/reducer.cpp" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<blockquote>
<p>Isn’t only  <code>.grad</code>  meant to be communicated and the workers applying the updates individually?</p>
</blockquote>
<p>No. Currently at construction time, DDP creates a mapping from parameters to buckets, and always communicate all buckets even if some gradients are not used in one iteration. The reason for doing so is that it is possible process 1 only computes grad A and process 2 only computes grad B. However, AllReduce operation requires all processes to provide the same set of input tensors. So in this case, both process 1 and 2 need to communicate grad A and B. DDP can use another communication to first figure out which grads are used globally. However, if block waiting for this signal, there will be no overlap between communication and computation, which could result in &gt;30% slowdown in some cases.</p>
<blockquote>
<p>If my parameters of  <code>theta[t]</code>  has only gradients for the particular task, would this help the case?</p>
</blockquote>
<p>It helps to skip computation but not communication. DDP always communicates all parameters in the model you passed to DDP constructor.</p>
<blockquote>
<p>I’m reading the  <strong>Forward Pass</strong>  section of <a href="https://pytorch.org/docs/master/notes/ddp.html#internal-design" rel="noopener nofollow ugc">Internal Design </a>, with  <code>find_unused_parameters</code> , it is possible to operate on a subgraph, correct(?)</p>
</blockquote>
<p>That flag only allows DDP to skip waiting for grads of those parameters. The communication phase is the same regardless the value of <code>find_unused_parameters</code>.</p>

jerinphilip: <aside class="onebox pdf">
  <header class="source">
      <a href="https://arxiv.org/pdf/2006.15704.pdf" target="_blank" rel="nofollow noopener">arxiv.org</a>
  </header>
  <article class="onebox-body">
    <a href="https://arxiv.org/pdf/2006.15704.pdf" target="_blank" rel="nofollow noopener"><span class="pdf-onebox-logo"></span></a>
<h3><a href="https://arxiv.org/pdf/2006.15704.pdf" target="_blank" rel="nofollow noopener">2006.15704.pdf</a></h3>

<p class="filesize">1161.94 KB</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>This the relevant paper?</p>

mrshenli: <p>Yep, it is the paper. Sorry, I forgot to update it here today.</p>