Fused Adam slow, request for fused AdamW

miraodasilva: <p>Hi, I tried using Adam(…,fused=True) and it does yield a 10-15% speedup on 1 GPU, as expected. However, when I try to use it with ddp (via lightning) with 2 GPUs, it yields no speedup at all. Is this expected behavior? I realize that this may be an issue with lightning and not torch, in which case this would not be the appropriate forum, but would just like to check if anyone has experienced similar behavior. (I’m running the latest torch and lightning on a machine with 2 2080tis).</p>
<p>Also, I was wondering if there are any plans to implement fused AdamW? Using AdamW’s weight decay is quite important for transformer training and, assuming there are speedups to be had, it would be great to have a fused version of it.</p>
<p>Thank you so much in advance!</p>

ptrblck: <aside class="quote no-group" data-username="miraodasilva" data-post="1" data-topic="174942">
<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/miraodasilva/48/58390_2.png" class="avatar"> miraodasilva:</div>
<blockquote>
<p>However, when I try to use it with ddp (via lightning) with 2 GPUs, it yields no speedup at all. Is this expected behavior?</p>
</blockquote>
</aside>
<p>You might need to profile the actual DDP workload in your Lightning setup and check the kernel execution, where the bottleneck of your training is etc. Based on your observation, the speedup from <code>Adam(... fused=True)</code> might be wiped out by another slowdown.</p>
<aside class="quote no-group" data-username="miraodasilva" data-post="1" data-topic="174942">
<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/miraodasilva/48/58390_2.png" class="avatar"> miraodasilva:</div>
<blockquote>
<p>Also, I was wondering if there are any plans to implement fused AdamW?</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/crcrpar">@crcrpar</a> has implemented fused <code>AdamW</code> in <a href="https://github.com/pytorch/pytorch/pull/88015">this PR</a> and it should already be available if I’m not mistaken.</p>

miraodasilva: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="174942">
<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>You might need to profile the actual DDP workload in your Lightning setup and check the kernel execution, where the bottleneck of your training is etc. Based on your observation, the speedup from <code>Adam(... fused=True)</code> might be wiped out by another slowdown.</p>
</blockquote>
</aside>
<p>Right, that makes sense, I will look into it.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="174942">
<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><a class="mention" href="/u/crcrpar">@crcrpar</a> has implemented fused <code>AdamW</code> in <a href="https://github.com/pytorch/pytorch/pull/88015" rel="noopener nofollow ugc">this PR </a> and it should already be available if I’m not mistaken.</p>
</blockquote>
</aside>
<p>I checked now and indeed it is available! Apologies for this, I recently upgraded to torch 2, I suppose the previous version didn’t have fused AdamW yet.</p>
<p>Thanks a lot for the help and kudos to you for all the hard work you do for the community!</p>