Optimizer.step() is slow than backprop

kbtorcher: <p>I am using Adam optimizer to train TD3 for offline RL setting. I saw significant increase in training time. Hence, I ran performance profiling for my code. What’s weird is backprop takes less time than update step. I ensured that I am properly using <code>torch.cuda.synchronize</code>. I’ve attached my code snippet for the same below.</p>
<p>Currently, using batch size of 256 for Hopper env and using pytorch 1.7.1.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/e/be3b013fab413b9b03cd3180f24fb86b59f80520.png" data-download-href="https://discuss.pytorch.org/uploads/default/be3b013fab413b9b03cd3180f24fb86b59f80520" title="actor_backprop_time"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be3b013fab413b9b03cd3180f24fb86b59f80520_2_690x227.png" alt="actor_backprop_time" data-base62-sha1="r8R8tG2H7CbfuRxJqDZIMc9VToI" width="690" height="227" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be3b013fab413b9b03cd3180f24fb86b59f80520_2_690x227.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be3b013fab413b9b03cd3180f24fb86b59f80520_2_1035x340.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/e/be3b013fab413b9b03cd3180f24fb86b59f80520.png 2x" data-dominant-color="FCFAFA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">actor_backprop_time</span><span class="informations">1351×445 23.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f92bd4b7cc19c04cf0090324830eaedaccf111a6.png" data-download-href="https://discuss.pytorch.org/uploads/default/f92bd4b7cc19c04cf0090324830eaedaccf111a6" title="actor_step_time"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f92bd4b7cc19c04cf0090324830eaedaccf111a6_2_690x227.png" alt="actor_step_time" data-base62-sha1="zygNruFQjIq6E4UsoHtpUzJJAiO" width="690" height="227" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f92bd4b7cc19c04cf0090324830eaedaccf111a6_2_690x227.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f92bd4b7cc19c04cf0090324830eaedaccf111a6_2_1035x340.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/9/f92bd4b7cc19c04cf0090324830eaedaccf111a6.png 2x" data-dominant-color="F8E9E5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">actor_step_time</span><span class="informations">1340×442 35.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/4/44cb32af072e800b97b412c494f41beaacfadcb0.png" data-download-href="https://discuss.pytorch.org/uploads/default/44cb32af072e800b97b412c494f41beaacfadcb0" title="critic_backprop_time"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44cb32af072e800b97b412c494f41beaacfadcb0_2_690x228.png" alt="critic_backprop_time" data-base62-sha1="9OzOH5GYqOBeCibMaUAnF6TiGd2" width="690" height="228" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44cb32af072e800b97b412c494f41beaacfadcb0_2_690x228.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44cb32af072e800b97b412c494f41beaacfadcb0_2_1035x342.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/4/44cb32af072e800b97b412c494f41beaacfadcb0.png 2x" data-dominant-color="FBF6F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">critic_backprop_time</span><span class="informations">1362×452 33.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/7/a76f0b4680cfa0d0e8646b459605c5ad59687bda.png" data-download-href="https://discuss.pytorch.org/uploads/default/a76f0b4680cfa0d0e8646b459605c5ad59687bda" title="critic_step_time"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a76f0b4680cfa0d0e8646b459605c5ad59687bda_2_690x230.png" alt="critic_step_time" data-base62-sha1="nTbEcbX5GoVpboo0hYFB1tAyQjw" width="690" height="230" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a76f0b4680cfa0d0e8646b459605c5ad59687bda_2_690x230.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a76f0b4680cfa0d0e8646b459605c5ad59687bda_2_1035x345.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/7/a76f0b4680cfa0d0e8646b459605c5ad59687bda.png 2x" data-dominant-color="F8E9E5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">critic_step_time</span><span class="informations">1350×451 41.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Below is my code snippet <strong>for critic</strong> (temp1, temp2, temp3, temp4 are cuda events):</p>
<pre><code class="lang-auto">temp2.record()

critic_loss.backward()

temp3.record()
torch.cuda.synchronize()
critic_backprop_time = temp2.elapsed_time(temp3)
#print("critic_backprop_time", critic_backprop_time)
writer.add_scalar("critic_backprop_time", critic_backprop_time, self.total_it)

temp2.record()

self.critic_optimizer.step()

temp3.record()
torch.cuda.synchronize()
critic_step_time = temp2.elapsed_time(temp3)
#print("critic_step_time", critic_step_time)
writer.add_scalar("critic_step_time", critic_step_time, self.total_it)
</code></pre>
<p>Code snippet <strong>for actor</strong>:</p>
<pre><code class="lang-auto">temp4.record()

actor_loss.backward()

temp3.record()
torch.cuda.synchronize()
actor_backprop_time = temp4.elapsed_time(temp3)
#print("actor_backprop_time", actor_backprop_time)
writer.add_scalar("actor_backprop_time", actor_backprop_time, self.total_it)

temp4.record()

self.actor_optimizer.step()

temp3.record()
torch.cuda.synchronize()
actor_step_time = temp4.elapsed_time(temp3)
#print("actor_step_time", actor_step_time)
writer.add_scalar("actor_step_time", actor_step_time, self.total_it)

temp4.record()
</code></pre>

ptrblck: <p>Try to update to the latest PyTorch release with the latest CUDA runtime and check the profiles again.<br>
In newer PyTorch releases a few optimizers accept the <code>foreach</code> argument, which could speed up the <code>step()</code> call.<br>
Also, I don’t understand what “for Hopper env” means.</p>

kbtorcher: <p>I updated to the latest pytorch version. After toggling <code>foreach</code> argument to <code>True</code> for Adam optimizer, it doesn’t speed up the <code>step()</code> call. The timing profile for <code>step()</code> is almost the same. For <code>backward()</code>, it performs slightly worse, leading to an increase of 2 minutes in total time for 10k iterations.</p>
<p><code>Hopper env</code> stands for hopper mujoco environment.</p>