How to preserve backward grad_fn after distributed operations

Andreas_Georgiou: <p>I am trying to implement model parallelism in a distributed cluster setting.</p>
<p>Let’s say I have a tensor <code>tensor</code> in each process and a number of operations have been performed on it (in each process independently). The tensor has a <code>.grad_fn</code> attached to it. Now I want to perform an <code>all_gather</code>. so that I create a list <code>[tensor_1, tensor_2...tensor_n]</code>. Then I can concatenate all those tensors using <code>torch.cat</code>. All the tensors in the list will lose the <code>grad_fn</code> property. My expectation is that process i will maintain the <code>grad_fn</code> for <code>tensor_i</code> in the list. It’s ok if all the others are lost. I want to be able to backward() after <code>torch.cat</code> in each process i through <code>tensor_i</code>. How can I achieve that? Any help is appreciated!</p>
<p>EDIT: I think I can just do <code>tensor_list[dist.get_rank()] = tensor</code> after the all_gather operation but I am not sure if there is a better way. Help?</p>

pietern: <p>Would it be possible to manually assign the <code>grad_fn</code> back to <code>tensor_i</code>?</p>
<p>I don’t think it’s a good idea to retain gradient functions on output tensors of collective functions. If anything, it would give an expectation of this working well out of the box, which is not the case. I think a better solution would be to stitch things together with <code>torch.autograd.grad</code> yourself, before and after the collectives.</p>

11123: <p>Do you have any idea? Do you want to calculate tensor_i in different process but accumulate between the processes so the loss will be attained by all the tensor_i?</p>

Andreas_Georgiou: <p>I’ve built this package that does this automatically now: <a href="https://github.com/ag14774/diffdist" rel="nofollow noopener">https://github.com/ag14774/diffdist</a>. So this question can be marked as solved</p>

11123: <p>That’s really cool! Thanks for your sharing!<br>
While I am not sure how the package works and whether it can be applied to such problem:</p>
<pre><code class="lang-auto">for iteration, data0, data1 in enumerate(data_loader, start_iter):
    tensor = model(data0)
    synchronize()
    tensors = dist.all_gather(tensor)
    loss = model(data1, tensors)
</code></pre>
<p>So in each process different data0 will generates a tensor, and the gathered tensors will be used for further training. Since ‘all_gather’ cannot preserve the ‘grad_fn’, can you give me some advice to solve it?<br>
Thanks a lot.</p>

Andreas_Georgiou: <p>Yes the package can do that. However, <code>tensor</code> needs to be of same shape and size in all processes. Then you can do something like:</p>
<pre><code class="lang-auto">for iteration, data0, data1 in enumerate(data_loader, start_iter):
    tensor = model(data0)
    synchronize()  # You probably do not need this since all_gather will force a sync
    gather_list = [torch.empty_like(tensor) for i in range(dist.get_world_size())]
    gather_list = diffdist.functional.all_gather(gather_list, tensor)
    loss = model(data1, gather_list)
</code></pre>
<p>Keep in mind though that <code>all_gather</code> is not very fast because its backprop involves running <code>dist.reduce</code> multiple times. When pytorch adds support for <code>reduce_scatter</code>, I will update the package to speed up the backprop.</p>

11123: <p>Thank you so much for your help. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"><br>
I tried the code, but the gather_list after <code>diffdist.functional.all_gather(gather_list, tensor)</code> also doesn’t contain each tensor’s grad_fn.</p>
<p>I found there is a parameter <code>self.next_backprop</code> in your code, do I need to set it? Sorry to bother you again.</p>

Andreas_Georgiou: <aside class="quote no-group" data-username="11123" data-post="7" data-topic="49343">
<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/11123/48/15404_2.png" class="avatar"> 11123:</div>
<blockquote>
<p>I found there is a parameter <code>self.next_backprop</code> in your code, do I need to set it? Sorry to bother you again.</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="Andreas_Georgiou" data-post="6" data-topic="49343">
<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/andreas_georgiou/48/11871_2.png" class="avatar"> Andreas_Georgiou:</div>
<blockquote>
<p>diffdist.functional.all_gather(gather_list, tensor)</p>
</blockquote>
</aside>
<p>Apologies, the line should be</p>
<pre><code class="lang-auto">gather_list = diffdist.functional.all_gather(gather_list, tensor)`
</code></pre>
<p>If you get any errors try setting <code>inplace=False</code>. No need to use <code>next_backprop</code></p>

11123: <p>Thank you again. one final question, I write a simple example  to see how the <code>grad_fn</code> works:</p>
<pre><code class="lang-auto">    # in each process:
    a = torch.tensor([1.0, 3.0], requires_grad=True).cuda()
    b = a + 2 * dist.get_rank()
    # gather
    bs = [torch.empty_like(b) for i in range(dist.get_world_size())]
    bs = diffdist.functional.all_gather(bs, b)
    # loss backward
    loss = (torch.cat(bs) * torch.cat(bs)).mean()
    loss.backward()
    print(a.grad)
</code></pre>
<p>I think <code>a</code> should has its gradient?  But currently it is <code>None</code>. I am a little bit lost.</p>

Andreas_Georgiou: <p>You are right it seems to be working for CPU but not for CUDA for some reason. I will investigate a bit more. Feel free to open a pull request if you find the problem</p>

Andreas_Georgiou: <p>I found the problem. The package is working fine. The problem is that when you set <code>requires_grad=True</code> you set it on the CPU version of <code>a</code>. Then you called <code>cuda()</code> which created another node in the graph. Gradient will pass through the GPU tensor <code>a</code> and then be accumulated to the CPU version of the tensor since that is the one that has <code>requires_grad</code> set to true. What you should do is <code>torch.tensor([1.0, 3.0], requires_grad=True, device='cuda')</code>. In a realistic scenario with normal training this won’t be a problem.</p>

11123: <p>Sorry for my late reply.</p>
<p>I tried your advice and then applied to my own model, it works! Thank you for your help. Actually I don’t know how do you implement your model parallelism, here I use <code>distributeddataparallel</code> in pytorch to distribute the model to different gpus of one node. So based on my experiment, I think maybe your work can also solve the distributed gpu grad_fn gathering problem? like in <a href="https://discuss.pytorch.org/t/will-dist-all-gather-break-the-auto-gradient-graph/47350" class="inline-onebox">Will "dist.all_gather" break the auto gradient graph?</a>. Thank you again.</p>

Andreas_Georgiou: <p>Glad it works!</p>
<aside class="quote no-group quote-modified" data-username="11123" data-post="12" data-topic="49343">
<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/11123/48/15404_2.png" class="avatar"> 11123:</div>
<blockquote>
<p>So based on my experiment, I think maybe your work can also solve the distributed gpu grad_fn gathering problem? like in <a href="https://discuss.pytorch.org/t/will-dist-all-gather-break-the-auto-gradient-graph/47350">Will “dist.all_gather” break the auto gradient graph?</a>.</p>
</blockquote>
</aside>
<p>Yes it seems that <code>diffdist</code> can handle that case. Of course different processes will have different computational graphs but with <code>diffdist</code> some nodes are inserted in the graph that will cause them to sync and communicate with each other. For example, doing a <code>Send</code> operation will cause a <code>Recv</code> to be called during backward in order to receive the gradient.</p>