DistributedDataParallel modify gradient before averaging

username: anxu
cooked: <p>Hi all! I think the “DistributedDataParallel” automatically average the gradient when calling “loss.backward()”. But is it possible to first compute the local gradient of the parameters, then do some modification to the local gradient, and finally average the gradient among the workers?</p>
<p>Thanks!</p>

username: Yanli_Zhao
cooked: <p><a class="mention" href="/u/anxu">@anxu</a> tensor.register_hook(customHook) may work for your case, you need to write customHook to modify grad of the tensor</p>

username: anxu
cooked: <p>Hi Yanli,</p>
<p>I am not sure whether tensor.register_hook will work, but the documentation mentioned that,</p>
<blockquote>
<p>Forward and backward hooks defined on  <code>module</code>  and its submodules won’t be invoked anymore, unless the hooks are initialized in the  <code>forward()</code>  method.</p>
</blockquote>
<p>Besides I need to first collect the whole gradient and then do some modification. Now I am turning to torch.distributed.all_reduce, but it will be easier if there is a way to do this via DistributedDataParallel.</p>

username: meilu_zhu
cooked: <p>Hi, anxu <a class="mention" href="/u/anxu">@anxu</a> , the “DistributedDataParallel” automatically average the gradient when calling “loss.backward()”,<br>
But I didn’t find the corresponding script in pytorch source code, Do you know where it is ?</p>

username: meilu_zhu
cooked: <p>Hi, Yanli <a class="mention" href="/u/yanli_zhao">@Yanli_Zhao</a> , the “DistributedDataParallel” automatically average the gradient when calling “loss.backward()”,<br>
But I didn’t find the corresponding script in pytorch source code, Do you know where it is ?</p>

username: rvarm1
cooked: <p>DDP averages gradients by all-reducing them across participating processes (see <a href="https://pytorch.org/docs/stable/_modules/torch/distributed/distributed_c10d.html#all_reduce" rel="nofollow noopener">https://pytorch.org/docs/stable/_modules/torch/distributed/distributed_c10d.html#all_reduce</a>). Some specific bits that include gradient averaging can be found in the <code>allReduce</code> calls here: <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp</a></p>

username: andyljones
cooked: <p><a class="mention" href="/u/yanli_zhao">@Yanli_Zhao</a>’s solution works great. You can  register the hook either before or after DDP’ing the model.  Though the docs say that hooks are removed, that’s either not actually the case or it doesn’t apply to hooks on the tensors themselves.</p>
<p>Here’s some demo code:</p>
<pre><code class="lang-auto">from torch.nn.parallel import DistributedDataParallel as DDP
from torch import nn
import torch
import os
import torch.distributed as dist
import torch.multiprocessing as mp

def setup(rank, world_size):
    """Setup code comes directly from the docs:  

    https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
    """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'

    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.manual_seed(42)

def cleanup():
    dist.destroy_process_group()

def pre_average(g):
    print(f'Pre-DDP hook ({g.device}): {g[0, 0]}')

def post_average(g):
    print(f'Post-DDP hook ({g.device}): {g[0, 0]}')

def worker(rank, world_size):
    # Set up multiprocessing stuff
    setup(rank, world_size)

    # Create a trivial model 
    model = nn.Linear(1, 1, bias=False).to(rank)
    torch.nn.init.constant_(model.weight, 1.)

    # Create some trivial data. 
    # Gradients for x = (1, 2) should be (2, 8)
    x = torch.tensor([rank+1]).float().to(rank)

    # Register a hook before and after DDP'ing the model
    model.weight.register_hook(pre_average)
    model = DDP(model, device_ids=[rank])
    model.module.weight.register_hook(post_average)

    # Backprop!
    l = model(x).pow(2).sum()
    l.backward()

    # Check what's left in the gradient tensors
    print(f'Final ({x.device}): {model.module.weight.grad[0, 0]}')

    cleanup()

if  __name__ == '__main__':
    world_size = 2
    mp.spawn(worker,
                args=(world_size,),
                nprocs=world_size,
                join=True) 
</code></pre>
<p>Run from the terminal, this should print</p>
<pre><code class="lang-nohighlight">Pre-DDP hook  (cuda:0): 2.0
Post-DDP hook (cuda:0): 2.0
Pre-DDP hook  (cuda:1): 8.0
Post-DDP hook (cuda:1): 8.0
Final value   (cuda:0): 5.0
Final value   (cuda:1): 5.0
</code></pre>

username: feiyuhuahuo
cooked: <p>Hi <a class="mention" href="/u/anxu">@anxu</a>， I did some test.  <code>Tensor.register_hook()</code> does work in DDP. I first wrap my model with DDP, then I modify the grad manually with a custom function in the <code>forward</code> code where I construct my model:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/b/dbcd77f7cab66827035e4a3420e7686f145faf88.png" alt="image" data-base62-sha1="vmsMollJwpa5GfrlIqRCE6pwJNC" width="258" height="55"><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b514c58e556b1af5fbf2b0607f2f22af0dbe9d9a.png" data-download-href="https://discuss.pytorch.org/uploads/default/b514c58e556b1af5fbf2b0607f2f22af0dbe9d9a" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b514c58e556b1af5fbf2b0607f2f22af0dbe9d9a.png" alt="image" data-base62-sha1="pPUVmVDLxXrA5lzXKGbYsLaF6cy" width="690" height="21" data-dominant-color="42464A"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">761×24 6.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Here’s the gradients after calling <code>loss.backward()</code>:<br>
If don’t modify the grad:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/6/d6bc543c8b9cdcfde91a06d2642731af650c9e1b.png" alt="image" data-base62-sha1="uDDFEpjtU5BQ77N02d0oSR8SN2H" width="595" height="35"><br>
And if modify the grad:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e4f1713aab2c843ceee8fb466dda5eef0dcd497c.png" alt="image" data-base62-sha1="wFk9mtktssxcr9o9lWbpZ2iPWPO" width="601" height="37"><br>
You can see the gradients are just 3 times after being modified, which means <code>.register_hook()</code>  works well in DDP.</p>

username: yuanze1024
cooked: <p>Hi, I just checked <a href="https://pytorch.org/docs/stable/_modules/torch/distributed/distributed_c10d.html#all_reduce" rel="noopener nofollow ugc">this url</a> you mentioned, and find that the function <code>all_reduce</code> has a default reduction_op <strong>SUM</strong>. So is the implicit operation <strong>SUM</strong> or <strong>Mean</strong>?</p>