Optimizing for-loops in PyTorch

username: neel_g
cooked: <p>Hi, I have a training loop where the number of iterations (i.e the number of times a tensor goes through the <strong>same</strong> block of layers) is dynamic - its bounded at some maximum value, but durig train time can change dynamically.</p>
<p>This repo demonstrates it well: <a href="https://github.com/aks2203/deep-thinking/blob/6b8aa99bfc3456541cc842d67f93a3fa41aa9055/deepthinking/models/dt_net_1d.py#L68-L87" rel="noopener nofollow ugc">Github Repo</a></p>
<p>(relevant lines are <code>76, 80-81</code> rest can be ignored)</p>
<p>Here’s a small snippet to further illustrate what I’m using:</p>
<pre><code class="lang-py">iters = randrange(1,10)
input : torch.tensor = ...

outputs  = ... # preallocate `iters` lenght tensor, to be filled at each iteration
for i in range(iters):
    interim_tensor = myblock(input) # myblock is a standard torch nn.Module
    outputs[i] = interim_tensor
</code></pre>
<p>In this case, the overhead from compilation is pretty large. I’m using <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"> <code>Accelerate</code> and GPU utilization spikes anywhere from 2% to 98% several times a second. Clearly, this is really inefficient.</p>
<p>Any ideas how I can optimize it? Perhaps the only way I can think of is having a <em>static</em> graph for the maximum amount of iterations, and then exiting the graph at some point, letting AutoGrad Automagically calculate and backprop accordingly.</p>
<p>Thanks,<br>
N</p>

username: neel_g
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> sorry for the ping, but would you have any idea how I can achieve this?</p>

username: neel_g
cooked: <p><img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"> Does anyone have any idea at all? <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_frowning_face.png?v=12" title=":slightly_frowning_face:" class="emoji" alt=":slightly_frowning_face:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>My idea would also be to try to script or compile the model/function, but it seems you have already tried it and the overhead is not acceptable:</p>
<aside class="quote no-group" data-username="neel_g" data-post="1" data-topic="182201">
<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/neel_g/48/42759_2.png" class="avatar"> neel_g:</div>
<blockquote>
<p>In this case, the overhead from compilation is pretty large.</p>
</blockquote>
</aside>
<p>How many different values for <code>iters</code> do you have? Could the overhead be amortized by a longer runtime assuming the number of different values for <code>iters</code> is not huge?</p>

username: neel_g
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="182201">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>compile the model/function,</p>
</blockquote>
</aside>
<p><code>torch.compile</code> didn’t give too much of a speedup; I reckon it might be about 10%?</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="182201">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>How many different values for <code>iters</code> do you have? Could the overhead be amortized by a longer runtime assuming the number of different values for <code>iters</code> is not huge?</p>
</blockquote>
</aside>
<p>well, <code>iters</code> is bounded in <code>[1, max_iters]</code> (both incl.); I don’t get what you mean by a longer runtime though…</p>

username: ptrblck
cooked: <p>Well, if <code>max_iters</code> is e.g. 10, but you are training your model for 2 weeks, the initial overhead might be noise compared to the long runtime of your training.</p>

username: neel_g
cooked: <p>ah, no my training times are usually &lt;20h on an RTX 3090.<br>
I’m not <em>fully</em> sure where the bottleneck here is, but I don’t think changing iterations like that provides any help in the training process.</p>
<p>Is there a way I can have a static graph for <code>max_iters</code>, traverse the <code>iters</code> part of it, and then exit it? so the graph would still be static - just that autograd won’t consider other iterations as part of the computation graph for that specific training step…</p>

username: ptrblck
cooked: <p>I’m unsure if <code>torch.compile</code> is able to do so, but <a class="mention" href="/u/marksaroufim">@marksaroufim</a> might know.</p>

username: marksaroufim
cooked: <p>Yeah I don’t think <code>torch.compile</code> would help much here, you’re probably better off rewriting your code in a batched way</p>

username: neel_g
cooked: <p>Thanks! could you elaborate on this a little bit? wdym by a batched way? also doesn’t <code>torch.compile</code> break the dynamic graph into pieces of static graphs?</p>

username: marksaroufim
cooked: <p>whoops my b, I just reread your code and noticed that your loops are dependent on the previous input so my batching suggestion won’t work</p>
<p>OK I think you have a few options</p>
<ol>
<li>Compile your model once for each size up front, inductor has a code cache which should speedup but not totally eliminate compilation times</li>
<li>Use <code>dynamic=True</code> when compiling the model</li>
<li>Have a static graph for the max number of iterations</li>
</ol>
<p>Also keep in mind on consumer cards <code>torch.compile()</code> may or not give you the best performance relative to server GPUs like A100</p>

username: neel_g
cooked: <aside class="quote no-group" data-username="marksaroufim" data-post="11" data-topic="182201">
<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/marksaroufim/48/36747_2.png" class="avatar"> marksaroufim:</div>
<blockquote>
<p>Have a static graph for the max number of iterations</p>
</blockquote>
</aside>
<p>Would that be something like what I mentioned above?</p>
<blockquote>
<p>Is there a way I can have a static graph for <code>max_iters</code>, traverse the <code>iters</code> part of it, and then exit it? so the graph would still be static - just that autograd won’t consider other iterations as part of the computation graph for that specific training step…</p>
</blockquote>
<p>If so, how I accomplish this?</p>