DDP does not work with custom gradient (backward) computations

username: mseeger
cooked: <p>Normally, gradients of a model are computed by</p>
<pre><code class="lang-auto">outputs = model(inputs)
loss = loss_func(outputs, targets)
loss.backward()
</code></pre>
<p>All the customization is in <code>model</code>, and <code>loss.backward()</code> is just using <code>autograd</code> once on the full model graph. For such a standard setup, DDP works fine.</p>
<p>But my setup is different, and for good reason. I need to compute gradients on a model with very long context. Doing the above does not work at all, runs OOM immediately. I implemented a complex gradient computation code, which calls <code>autograd</code> several times for parts of the model, using checkpointing and stuff in between.</p>
<p>This works fine now on a single GPU, but I am banging my head against the wall to get DDP to work. I use <code>static_graph=True</code>, otherwise nothing works. But now, I see devices just hanging, and other ranks happily moving on to do the next gradient without syncing (all reduce) at all. Eventually, everything grinds to a halt.</p>
<p>This may all be related to this one: <a href="https://github.com/pytorch/pytorch/issues/24005" class="inline-onebox" rel="noopener nofollow ugc">Using `torch.utils.checkpoint.checkpoint_sequential` and `torch.autograd.grad` breaks when used in combination with `DistributedDataParallel` · Issue #24005 · pytorch/pytorch · GitHub</a>, which is open since 2019.</p>
<p>What am I missing here? In which sense is synchronization between devices dependent on how often <code>autograd</code> is called? How can I possibly get around this behavior?</p>
<p>Do I really need to implement DDP from scratch, calling broadcast or all reduce myself, just because the official implementation somehow links weights/gradient sync with whether autograd is called multiple times?</p>
<p>I’d love to understand this, and why this issue is still open. Is PyTorch not interested at all in people computing gradients in a non-standard way?</p>

username: mseeger
cooked: <p>When I look at <code>DistributedDataParallel</code>, I see options <strong><code>delay_all_reduce_named_params</code>, param_to_hook_all_reduce,</strong> and <strong>skip_all_reduce_unused_params</strong>.</p>
<p>I have the vague feeling these could be helpful in my case. Huge cudos to somebody who could point me to an example what these are doing, and whether they’d help getting DDP to work in my non-standard case.</p>
<p>I’d like to be able to use <code>autograd</code> several times, on parts of the model, until <code>x.grad</code> is being written for all params, and I’d not like DDP to issue broadcasts or whatever until all of that is done.</p>

username: H-Huang
cooked: <p>Hi! Thanks for the question. I think that issue is just so old it got lost and was never closed, activation checkpointing has changed quite a bit so not sure if it is still valid.</p>
<p>For your use case, a simple model / train loop example would really help in order to make changes for non-conventional use cases.</p>
<p>One thing is that the core performance improvement of DDP is bucketing gradients so that it communicates buckets while the other backward computation is happening, thus overlapping compute and communication. If you run autograd several times, it is not obvious how that bucketing should be done. If you are saying you would like to delay the communication until the very end when all grads have been computed, it makes sense just to call all_reduce manually on the gradients rather than use DDP. Also it was my understanding that <code>autograd.grad</code> does not trigger the DDP communication hooks and only <code>autograd.backward</code> does.</p>

username: mseeger
cooked: <p>Hello, thanks for the help. Isn’t there a way to switch off this bucketing optimization?</p>
<p>I am currently trying to use PyTorch Lightning, which has all sorts of nice structure, hooks, support for logging, checkpointing, etc. While I could maybe write all that from scratch, I’d prefer to just switch this bucketing off.</p>
<p>BTW: My algorithm does compute gradients of the model in the same order than a single <code>autograd.backward</code>, from top layer down to bottom. Since I have very long sequences broken into chunks, I need to accumulate gradients several times, adding contributions from each chunk.</p>
<p>I’d be perfectly able to signal your bucketing when each <code>x.grad</code> is finished and can be communicated. It is just what I cannot do is call <code>loss.backward()</code> on a single <code>loss</code>, encoding the whole compute graph over all layers and all chunks. Well, maybe I could do it with some very complicated autograd saved tensor hook packing and unpacking. But that would make a method that is already quite complex, even more complex.</p>