Gradient hooks with DataParallel and DDP

username: ekurtic
cooked: <p>Hi everyone,</p>
<p>I’d appreciate some help/suggestions on this issue:<br>
For simplicity, let’s say I have a model with only Conv2d layer in it. I would like to store gradient of the weight param of Conv2d layer over time, so I’ve registered a hook on it. When one batch of data is processed and loss computed, backward pass will start gradient computation and my hook will be triggered. In the hook function, I can store the gradient (for example in a list).</p>
<p>Now my problem is to do the same thing but with DataParallel model. If my hooks looks like this:</p>
<pre><code class="lang-python">def my_gradient_hook(grad):
    print(grad.device)
</code></pre>
<p>the only print that I see is: <code>cuda:0</code>. What happened to prints/hooks on other GPUs, because each GPU should run a backward pass in parallel and compute gradients?</p>
<p>If I try the same thing with DDP, I get the desired behavior. Prints that I see are: <code>cuda:0</code>, <code>cuda:1</code>, <code>cuda:2</code>…</p>
<p>I guess it has something to do with multiprocessing, since DDP spawns multiple processes and DataParallel only one process with multiple threads. I guess I should somehow gather gradients from different threads in the hook function?</p>

username: agolynski
cooked: <p>Hi,</p>
<p>Seems like this question had come up before. Seems like this limitation is intrinsic to DataParallel, e.g. see <a href="https://discuss.pytorch.org/u/Yanli_Zhao">Yanli_Zhao</a> answer here:</p><aside class="quote" data-post="1" data-topic="73354">
  <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/a/47e85d/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/when-i-run-model-on-multiple-gpus-register-hook-is-invalid/73354">When I run model on multiple GPUs,register_hook is invalid</a> <a class="badge-category__wrapper " href="/c/distributed/12"><span data-category-id="12" data-drop-close="true" class="badge-category "><span class="badge-category__name">distributed</span></span></a>
  </div>
  <blockquote>
    I want to save gradients of internal variables through register_hook() or retain_grad(). 
When I run model on single GPU, it works. 
But when I run model on multiple GPUs through wrapping model into nn.DataParallel, I find that it doesn’t work. 
Can anyone help me?
  </blockquote>
</aside>

<p>Also see warnings here</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/nn/parallel/data_parallel.py#L67">
  <header class="source">

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

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



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="57" style="counter-reset: li-counter 56 ;">
          <li></li>
          <li>.. warning::</li>
          <li>    It is recommended to use :class:`~torch.nn.parallel.DistributedDataParallel`,</li>
          <li>    instead of this class, to do multi-GPU training, even if there is only a single</li>
          <li>    node. See: :ref:`cuda-nn-ddp-instead` and :ref:`ddp`.</li>
          <li></li>
          <li>Arbitrary positional and keyword inputs are allowed to be passed into</li>
          <li>DataParallel but some types are specially handled. tensors will be</li>
          <li>**scattered** on dim specified (default 0). tuple, list and dict types will</li>
          <li>be shallow copied. The other types will be shared among different threads</li>
          <li class="selected">and can be corrupted if written to in the model's forward pass.</li>
          <li></li>
          <li>The parallelized :attr:`module` must have its parameters and buffers on</li>
          <li>``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`</li>
          <li>module.</li>
          <li></li>
          <li>.. warning::</li>
          <li>    In each forward, :attr:`module` is **replicated** on each device, so any</li>
          <li>    updates to the running module in ``forward`` will be lost. For example,</li>
          <li>    if :attr:`module` has a counter attribute that is incremented in each</li>
          <li>    ``forward``, it will always stay at the initial value because the update</li>
      </ol>
    </code></pre>



  </article>

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

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


username: ekurtic
cooked: <p>Hi, thanks a lot for the link. I have already seen this warning and discussion, and (if I understood correctly) they are focused on <em>updating</em> parameters/buffers on non-master GPU. I understand that all modifications will be discarded by destroying these replicas and I’m fine with that. My goal is not to update any states, it’s just to print gradients at each device.</p>