Flat model grads, CUDA graph capture, and grad-copy-elision

username: cbcase
cooked: <p>This question is a bundle of a few things I’ve struggled to answer from the docs.</p>
<p>At the top level, I am interested in approaches to maintaining “flat grads” – ie, the model gradients are a single contiguous buffer (of the same dtype as model params), and the <code>.grad</code> attribute for each param is a pointer into that contiguous buffer. Afaict, this is still possible (if not exactly encouraged) by manually mutating <code>.grad</code> to be a view tensor like so:</p>
<pre><code class="lang-auto">m = torch.nn.Linear(hidden, hidden).cuda()
buf = torch.zeros(m.weight.numel(), dtype=m.weight.dtype, device=m.weight.device)
m.weight.grad = buf[0:m.weight.numel()].view_as(m.weight)
# And so on for all the the parameters
# After running fwd + bwd, we have:
assert buf.untyped_storage().data_ptr() == m.weight.grad.untyped_storage().data_ptr()
</code></pre>
<p>Two partially-related questions related to this technique:</p>
<ol>
<li>
<p>How does this interact with cuda graph capture – and more generally, is it possible to reason at all about the pointers in a captured cuda graph? By definition, the output <code>.grad</code> of each gradient-computing kernel needs to be stable across graph executions – so if I do something like the above, will it capture the pointer that is a reference-to-buf on the first run and use it in subsequent runs? Or does the graph capture need to use its own internal memory pool to allocate the grad attributes?</p>
</li>
<li>
<p>There is (afaik) a small optimization when calling <code>zero_grad(set_to_none=True)</code> where the computed gradient doesn’t need to be <em>accumulated</em> during backprop but can instead be <em>stored</em> (saving an extra memory round trip) – I heard this referred to as “grad copy elision”. Is there any way to convince autograd to perform this same optimization in the flat grads case defined above, where I want to ensure the grad pointer remains unchanged, but I don’t care about accumulating it?</p>
</li>
</ol>
<p>Thanks!</p>

username: ptrblck
cooked: <p>In the <a href="https://pytorch.org/docs/stable/notes/cuda.html#whole-network-capture">whole-network capture example</a> the gradients are explicitly deleted via <code>set_to_none=True</code> with the comment:</p>
<blockquote>
<p>Sets grads to None before capture, so backward() will create .grad attributes with allocations from the graph’s private pool</p>
</blockquote>
<p>which points towards your last statement. A bit further down you are then seeing the <code>g.replay</code> which might look a bit weird since the <code>zero_grad</code> call wasn’t captures. However, again the comment mentions:</p>
<blockquote>
<p>replay() includes forward, backward, and step.<br>
You don’t even need to call optimizer.zero_grad() between iterations<br>
because the captured backward refills static .grad tensors in place.</p>
</blockquote>
<p>Based on this I would assume deleting the gradients could break the capture. However, it would still be interesting to see if you could  “bake in” the flattened gradient buffers without deleting them in the capture or if this would directly break and raise errors.</p>