Unusual scaling in per-example gradient computation

jusjusjus: <p>Hi pytorch friends!<br>
I want to optimize per-example gradient computation, and found an interesting behavior when increasing the batch size.   My findings can be reproduced in <a href="https://colab.research.google.com/gist/jusjusjus/7c04e2645ba335061f96f45bab79a16b/scaling-in-per-example-computation.ipynb" rel="nofollow noopener">this notebook</a>.  Essentially I compute per-example losses then to call <code>backward</code> on each one:</p>
<pre><code class="lang-python">t0 = time()
losses = loss_fn(model(samples), labels)
model.zero_grad()
for i, l in enumerate(losses):
    l.backward(retain_graph=True)
    # do something with the individual gradients
sampling_rate = batch_size / (time() - t0)
print(f"{sampling_rate:.0f} Hz")
</code></pre>
<p><strong>I find that the computed sampling rate has a maximum for 20 to 60 samples dependent on the used GPU.</strong><br>
How can you (a) explain this behavior, and (b) possibly fix it so that I can get the most of my computations?<br>
Thanks for considering my problem, please AMA <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

albanD: <p>Hey,</p>
<p>This is most likely because of the large number of calls to <code>.backward()</code>.<br>
Out of curiosity, how do you “# do something with the individual gradients” given that you don’t zero_grad() before? You do something with all the gradients accumulated up to now?</p>

jusjusjus: <p>interesting; what type of overhead is generated by multiple calls to <code>backward</code>?  Regarding curiosity, I omitted most of that (including <code>zero_grad</code>).  It’s related to differential privacy.</p>

albanD: <blockquote>
<p>What type of overhead does multiple calls to <code>backward</code> generate?</p>
</blockquote>
<p>We need to go in cpp, discover the work that needs to be done, start GPU kernels for everything.<br>
That takes some time.</p>
<blockquote>
<p>omitted most of that (including <code>zero_grad</code> ). It’s related to differential privacy.</p>
</blockquote>
<p>One thing I would say is that if the final transformation that you’re doing on the gradient is linear, then you can apply that same (similar) transformation to the loss before accumulation.<br>
Otherwise, you can try to use something like <a href="https://github.com/cybertronai/autograd-hacks">this</a> to get a per-sample gradients in one backward pass (more or less) and then do all the post processing after that. That might be faster.</p>

jusjusjus: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71145">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>We need to go in cpp, discover the work that needs to be done, start GPU kernels for everything.<br>
That takes some time.</p>
</blockquote>
</aside>
<p>I modified the code like so to investigate your explanantion:</p>
<pre data-code-wrap="python"><code class="lang-python">for i, l in enumerate(losses):
    t = time()
    # just to do something with the gradients
    model.zero_grad()
    l.backward(retain_graph=True)
    print(i, time() - t)
</code></pre>
<p>What I get from this is</p>
<pre><code class="lang-auto">...
24 0.0007534027099609375
25 0.0007722377777099609
26 0.0007855892181396484
27 0.0024886131286621094
28 0.0009698867797851562
29 0.0012204647064208984
30 0.0010230541229248047
31 0.0008389949798583984
32 0.0024178028106689453
33 0.0007109642028808594
34 0.0007128715515136719
35 0.0006940364837646484
36 0.0009226799011230469
37 0.0046350955963134766
38 0.004754304885864258
39 0.004813194274902344
40 0.004805564880371094
41 0.004813432693481445
...
</code></pre>
<p>Why do you think does it take four times as long to go to cpp in iteration 37 compared with iteration 36?</p>
<p>Here’s an infographics over a bunch of runs on my dedicated nv 2080 (sampling rate ~ <code>1 / (time() - t)</code>, iteration: that’s the index into the batch in the loop across losses):<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/1/8/1800a1b081e397007fcd8c032822baa9378b78ba.png" alt="image" data-base62-sha1="3qkOncieyGH0rVQHf9iEMcANxJo" width="395" height="266"></p>
<p>However, thanks for the pointer to autograd hacks; I’ll try that out.</p>

albanD: <p>Hi,</p>
<p>I would bet on you having filled the GPU task queue and so you now have to wait for stuff to be finished on the GPU before queuing more work.</p>
<p>The GPU api is asynchronous, so here, you only meansure the time to launch the jobs.<br>
You can add <code>torch.cuda.synchronize()</code> before your timing if you want to measure the time it actually takes to compute the result.</p>

jusjusjus: <p>That’s a great idea.  I added that to the linked notebook (not saving it; you’ll have to do it yourself), and now I receive the following numbers on an nvidia T4:</p>
<pre><code class="lang-auto">at 3-batches, the sample rate is 614 +/- 122 Hz
at 24-batches, the sample rate is 883 +/- 132 Hz
at 64-batches, the sample rate is 380 +/- 5 Hz
at 96-batches, the sample rate is 265 +/- 2 Hz
</code></pre>
<p>I’m still a bit puzzled, though, that we observe a maximum sample processing rate at some 20 samples per batch. Wouldn’t one expect the runtime of the for-loop to scale linearly with the batch size, thus giving a nearly constant sample rate?</p>

albanD: <p>That is indeed interesting.<br>
I did a quick test and it has nothing to do with the backwards: if you only measure the forward time, you see the same behavior <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"><br>
So I would guess that the kernel for larger batch size are more optimized than the ones for small batch size. And this is why you see this.<br>
Note that if you use cudnn, it might even take completely different algorithms depending on the input shape!</p>

jusjusjus: <aside class="quote no-group" data-username="albanD" data-post="8" data-topic="71145">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>So I would guess that the kernel for larger batch size are more optimized than the ones for small batch size. And this is why you see this.</p>
</blockquote>
</aside>
<p>Yeah, though I think it’s the other way around when trying to backprop sample-by-sample.  Maybe it’s possible to somehow fix the algorithm to the version for small batch sizes.</p>

jusjusjus: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="71145">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Otherwise, you can try to use something like <a href="https://github.com/cybertronai/autograd-hacks" rel="noopener nofollow ugc">this </a> to get a per-sample gradients in one backward pass (more or less) and then do all the post processing after that. That might be faster.</p>
</blockquote>
</aside>
<p>I followed your advise, but to make it work for my GAN application I had to make some major adoptions of the code.  I include this link here for other people’s reference that might encounter the same problem.  Thanks again:</p>
<p><a href="https://github.com/jusjusjus/autograd-hacks" class="onebox" target="_blank" rel="noopener nofollow ugc">https://github.com/jusjusjus/autograd-hacks</a></p>
<p>Fyi, I also marked it as a solution, even though I think we still have some unresolved questions to the initial observation.</p>

albanD: <p>cc <a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a> the author of autograd-hacks if you want to discuss the necessary changes.</p>
<p>Did that solve your perf problems? Or at least remove any multi-backward fear?</p>

albanD: <p>Also this is a modified version of your code that prints runtime instead of throughput.<br>
<a href="https://colab.research.google.com/drive/1aMCMZaSBak0PssL8nt2ff3mgFRKeNWaP" class="onebox" target="_blank">https://colab.research.google.com/drive/1aMCMZaSBak0PssL8nt2ff3mgFRKeNWaP</a><br>
This makes it clear what happens: for small enough batch size, you get a sub-linear scaling as the GPU is under-utilized. At some point, you fully use the GPU and you start scaling linearly.</p>
<p>That would explain why the throughput looks this way.</p>

Yaroslav_Bulatov: <p><a class="mention" href="/u/jusjusjus">@jusjusjus</a> glad to see it was useful! I had a similar issue with GANs – saving forward quantities as module attribute means generator/discriminator passes got mixed together. Also, keeping activations around took too much memory, so I ended up using lower level <a href="https://github.com/cybertronai/autograd-lib" rel="nofollow noopener">autograd-lib</a> for my application. Instead of storing activations in global variable, store them in closure-local variable, this way the memory gets released as you are done with the layer – <a href="https://colab.research.google.com/drive/1mUWfuydS7KqoPcFnFtrg4uCLMO-4nHOj#scrollTo=WMDbKbLtrxwG" rel="nofollow noopener">example</a></p>
<p>The slowness has to do with batch size not saturating GPU compute units, so it could be called a hardware issue. However, autograd-hacks just restructures the computation to do a single large matmul instead of concatting together a bunch of small matmuls. Why couldn’t this be done automatically? In this sense it’s a software issue.</p>

jusjusjus: <aside class="quote no-group" data-username="albanD" data-post="12" data-topic="71145">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>That would explain why the throughput looks this way.</p>
</blockquote>
</aside>
<p>Are you sure about that?  In a linear scaling, one expects what I call <em>sample rate</em> above (<code>batch_size / (t1-t0)</code>) to saturate.  However, it starts to decrease with increasing batch sizes (see post above).  Larger batches, it seems, are computed less efficiently than running multiple computes on smaller batches.</p>

albanD: <p>But the fact that you do as many backwards as there are samples in the batch makes them slower when the batch size increases.<br>
In that case, I can buy that there is a sweet spot for a small number of samples.<br>
Also note that each call to backward uses the FULL graph. So the more sample, the more expensive the backward for each individual sample.</p>