Fused mixed precision updates with PyTorch amp

username: cbcase
cooked: <p>I am trying to get up to speed on the latest support for mixed precision in PyTorch. One thing that is not clear to me right now is the interaction among <code>autocast</code>, <code>DistributedDataParallel</code> and optimizer weight updates.</p>
<p>Suppose I am using bfloat16 (so we can ignore grad scaling for the moment) and I simply run:</p>
<pre><code class="lang-auto">with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
  loss = model(...)
loss.backward()
</code></pre>
<p>At that point, it looks like the <code>.grad</code> attributes on the model are in fp32 (as with the weights), since the autocast in forward of each weight was backprop’d through. Ideally I could:</p>
<ul>
<li>Produce bfloat grads from backprop</li>
<li>(Optionally) allreduce those grads (with DDP)</li>
<li>Run a single weight update kernel that reads the bfloat grads and then updates the fp32 weights (and optimizer state)</li>
</ul>
<p>There is clearly some support for doing things like this (eg, I see that the Adam optimizer has a bunch of logic around supporting passing fp16 scale values into the step function, which only makes sense if you are operating on fp16 grads), but it’s not clear how to reason about what values are in what types and when. Thanks!</p>

username: ptrblck
cooked: <ol>
<li>
<p>I don’t think this is directly possible since the gradients would be computed using the <code>dtype</code>, <code>device</code>, and memory layout of the corresponding parameter.</p>
</li>
<li>
<p>You might try to use <a href="https://pytorch.org/docs/stable/ddp_comm_hooks.html">DDP comm hooks</a> to transform the gradients before communicating them.</p>
</li>
<li>
<p>Same as in 1: I don’t believe a <code>dtype</code> mismatch is supported but I let others correct me in case there are some (experimental) implementations in the wild.</p>
</li>
</ol>

username: cbcase
cooked: <p>Thanks for the reply, Piotr! So it sounds like if we want to (eg) perform half-precision allreduce, it’s still the case that we need to do something like the (now deprecated, I see) FP16_Optimizer – ie, explicitly put the model dtype to half precision and manually manage model and master weights at each weight update?</p>

username: ptrblck
cooked: <p>Hi Carl! Yes, if you want to apply the now deprecated <code>apex.amp</code> <code>O2</code>-style mixed-precision training you would still need to handle the master parameters and gradients manually.<br>
However, you could check the experimental <a href="https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.MixedPrecision"><code>fsdp.MixedPrecision</code></a> class which allows you to specify the <code>reduce_dtype</code> and e.g. <code>keep_low_precision_grads</code> and might come close to your use case.<br>
Of course you would need to use <code>FSDP</code> so unsure if this would work for you.</p>

username: cbcase
cooked: <p>Thanks, Piotr – that all makes sense. And thanks or the pointer to FSDP – interested in exactly that, actually, so I’ll play around with the MixedPrecision interface. I think a combination of FSDP along with apex’s FusedAdam implementation (which supports some degree of explicit master weight management) may get the job done.</p>