Slow Embedding backward

utsav: <p>Hi!</p>
<p>My network has a large embedding layer <code>[141713, 128]</code>. The forward pass takes about 0.01s but the backward is taking almost 0.47s which is 47x of the forward operation.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/1/11fa0a68fc35457c7387714063a66b547bbefae5.png" data-download-href="https://discuss.pytorch.org/uploads/default/11fa0a68fc35457c7387714063a66b547bbefae5" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/11fa0a68fc35457c7387714063a66b547bbefae5_2_689x230.png" alt="image" data-base62-sha1="2z1PfDYcIAsrrxYcs41I1wgXYr3" width="689" height="230" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/11fa0a68fc35457c7387714063a66b547bbefae5_2_689x230.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/11fa0a68fc35457c7387714063a66b547bbefae5_2_1033x345.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/11fa0a68fc35457c7387714063a66b547bbefae5_2_1378x460.png 2x" data-dominant-color="1E2227"><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">image</span><span class="informations">2536×848 237 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>Also, when I used <code>torch.autograd.profiler.profile(use_cuda=True)</code> then I saw a significant amount of time is taken by <code>embedding_dense_backward</code> on the CPU but the network is training on GPU. This could be the reason for the slow-down during backward.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/1/f196e793c3c0e9a5d9b3adaf769ca5192355551a.png" data-download-href="https://discuss.pytorch.org/uploads/default/f196e793c3c0e9a5d9b3adaf769ca5192355551a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f196e793c3c0e9a5d9b3adaf769ca5192355551a_2_690x458.png" alt="image" data-base62-sha1="ytcnGmKBRVW2xdMCRU0AyGm3JEK" width="690" height="458" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f196e793c3c0e9a5d9b3adaf769ca5192355551a_2_690x458.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f196e793c3c0e9a5d9b3adaf769ca5192355551a_2_1035x687.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/1/f196e793c3c0e9a5d9b3adaf769ca5192355551a_2_1380x916.png 2x" data-dominant-color="21272C"><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">image</span><span class="informations">2824×1876 1.9 MB</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>I also tried using <code>sparse=True</code> in the embedding layer, but that did not have any significant impact on the timing.</p>
<p>Could you please provide some insights into this and ways to overcome it?</p>

Bjorn_Lindqvist: <p>I’m seeing the same problem. But it is possible the slowdown seen in the profile trace is a mirage. See <a href="https://github.com/pytorch/pytorch/issues/30711" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/30711</a> The root cause is probably somewhere else.</p>