Optim.step() is significantly SLOW on MPS

HaoChris_Chen: <p>The first screen shot shows the simple training loop with code details, which is running on CPU. It makes sense that loss.back() is low on CPU. The second one is running on MPS, and loss.back() is faster than the first one, which makes sense too.</p>
<p>However, if we look at optim.step(), which should be a very simple step, it is way too slow, especially on MPS, which is nonsense.<br>
I tried :<br>
export PYTORCH_ENABLE_MPS_FALLBACK=0<br>
export PYTORCH_ENABLE_MPS_FALLBACK=1<br>
torch.mps.synchronize()<br>
None of the above helps the situation.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/4/84d5272ff18359aaf6d45d3c15eac8ef1d69cbb3.png" data-download-href="https://discuss.pytorch.org/uploads/default/84d5272ff18359aaf6d45d3c15eac8ef1d69cbb3" title="Screenshot 2023-11-04 at 10.24.19 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/84d5272ff18359aaf6d45d3c15eac8ef1d69cbb3_2_407x500.png" alt="Screenshot 2023-11-04 at 10.24.19 AM" data-base62-sha1="iX5HFVOwf1ZvmFLfahKmZ6BBWld" width="407" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/84d5272ff18359aaf6d45d3c15eac8ef1d69cbb3_2_407x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/84d5272ff18359aaf6d45d3c15eac8ef1d69cbb3_2_610x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/4/84d5272ff18359aaf6d45d3c15eac8ef1d69cbb3.png 2x" data-dominant-color="212121"><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">Screenshot 2023-11-04 at 10.24.19 AM</span><span class="informations">693×850 80.4 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><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/1/b140dd94be407bdb7ddd01764a18889f89871e1a.png" data-download-href="https://discuss.pytorch.org/uploads/default/b140dd94be407bdb7ddd01764a18889f89871e1a" title="Screenshot 2023-11-04 at 10.23.59 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b140dd94be407bdb7ddd01764a18889f89871e1a_2_492x375.png" alt="Screenshot 2023-11-04 at 10.23.59 AM" data-base62-sha1="pi3u8spdPMwruNFwoOaVqm1ufqO" width="492" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b140dd94be407bdb7ddd01764a18889f89871e1a_2_492x375.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b140dd94be407bdb7ddd01764a18889f89871e1a_2_738x562.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/1/b140dd94be407bdb7ddd01764a18889f89871e1a.png 2x" data-dominant-color="212223"><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">Screenshot 2023-11-04 at 10.23.59 AM</span><span class="informations">745×567 42.9 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>