How to use torch.autograd.Function with AMP?

pallgeuer: <p>I implemented a custom function using <code>torch.autograd.Function</code>, but this gives me an error when I enable AMP. It looks conceptually something like this:</p>
<pre><code class="lang-auto">class CustomFunction(torch.autograd.Function):

	@staticmethod
	def forward(ctx, inp, dim):
		...
		ctx.save_for_backward(tensor_value)
		...
		return answer

	@staticmethod
	@torch.autograd.function.once_differentiable
	def backward(ctx, grad_answer):
		...
		tensor_value, = ctx.saved_tensors
		...
		foo_tensor = grad_answer.addcmul(some_tensor, tensor_value, value=2)
		...
</code></pre>
<p>This results in:</p>
<pre><code class="lang-auto">RuntimeError: !(has_different_input_dtypes &amp;&amp; !config.promote_inputs_to_common_dtype_ &amp;&amp; (has_undefined_outputs || config.enforce_safe_casting_to_output_ || config.cast_common_dtype_to_outputs_)) INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/TensorIterator.cpp":405, please report a bug to PyTorch.
</code></pre>
<p>The problem is that <code>tensor_value</code> is <code>float32</code>, while due to AMP <code>grad_answer</code> is <code>float16</code>, and <code>addcmul</code> doesn’t natively promote to the widest input. <strong>What is the intended / best / most general way to circumnavigate this issue?</strong></p>
<p>The forward pass is executed under AMP / <code>torch.autocast</code>, so calling <code>addcmul</code> there would promote to the widest input and not cause an issue - but seeing as <code>addcmul</code> is needed in the backward pass and the backward pass is <em>not</em> executed under AMP this is an issue. Other operations like <code>add</code> always promote as required, which is why you don’t see this issue for them.</p>
<p>Just manually hard-coding a cast of <code>grad_answer</code> to <code>float32</code> is not a good/general solution because it is brittle, relies on inherent knowledge of which tensor happens to end up as which type during AMP, and if anything changes might thereby force conversions to float32 where both inputs are actually float16. Also, the <code>backward()</code> function could theoretically contain tens of calls to methods that don’t auto-promote. Trying to fix all those up with manual casts is undesirable, inflexible, and might force certain operations to run at higher precision than actually intended, or invoke unnecessary copies etc.</p>
<p>How is this normally handled?<br>
Can <code>addcmul</code> be changed to auto-promote (implementation change)?<br>
Can <code>config.promote_inputs_to_common_dtype_</code> be temporarily enabled?<br>
What dtype should the return value of the backward pass actually have? Always same as grad_answer? Same as inp? No restriction?</p>

soulitzer: <p>There is a method to compute the promoted type based on the dtype of the inputs. <a href="https://pytorch.org/docs/stable/generated/torch.promote_types.html" class="inline-onebox" rel="noopener nofollow ugc">torch.promote_types — PyTorch 2.0 documentation</a>. Does this solve your issue?</p>

pallgeuer: <p>So if I understand correctly I would have to do something like:</p>
<pre><code class="lang-auto">dtype = torch.promote_types(some_tensor, tensor_value)
dtype = torch.promote_types(dtype, grad_answer)
foo_tensor = grad_answer.to(dtype=dtype).addcmul(some_tensor.to(dtype), tensor_value.to(dtype), value=2)
</code></pre>
<p>for each call to a non-auto-promoting function like <code>addcmul</code> in <code>backward()</code>?<br>
I guess for simple isolated cases this could be a workaround.</p>
<p><em>How about wrapping the <code>backward()</code> call in a <code>torch.autocast</code> context?</em></p>
<p>The documentation for <code>torch.autocast</code> says (<a href="https://pytorch.org/docs/stable/amp.html#torch.autocast" class="inline-onebox" rel="noopener nofollow ugc">Automatic Mixed Precision package - torch.amp — PyTorch 2.1 documentation</a>):</p>
<blockquote>
<p><a href="https://pytorch.org/docs/stable/amp.html#torch.autocast" rel="noopener nofollow ugc"><code>autocast</code></a> should wrap only the forward pass(es) of your network, including the loss computation(s). Backward passes under autocast are not recommended. Backward ops run in the same type that autocast used for corresponding forward ops.</p>
</blockquote>
<p>The first half of that seems to dissuade us from my proposition. Howcome?<br>
Then what does the second bit mean?</p>

soulitzer: <p>The second half: autograd records the types that the operations were performed in in forward, and does casting automatically so that gradients match the dtypes of the inputs.</p>
<p>Have you checked out <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#autocast-and-custom-autograd-functions" class="inline-onebox" rel="noopener nofollow ugc">CUDA Automatic Mixed Precision examples — PyTorch 2.0 documentation</a>? When using autocast you should decorate your custom Function forward and backward like so to prevent type mismatch errors:</p>
<pre><code class="lang-auto">class MyMM(torch.autograd.Function):
    @staticmethod
    @custom_fwd
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a.mm(b)
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        a, b = ctx.saved_tensors
        return grad.mm(b.t()), a.t().mm(grad)
</code></pre>

pallgeuer: <p>I hadn’t seen that yet.<br>
Digging into the code it seems to do pretty much what I naively suggested - it executes the backward method with <code>torch.autocast</code> if the forward part had it enabled. This is definitely the intended and clean solution then, thanks.</p>
<p>I guess the only thing left I’m unsure about is why in general the backward pass shouldn’t be run under <code>torch.autocast</code>, but running arbitrary backward methods under <code>torch.autocast</code> like custom_bwd does is fine.</p>

soulitzer: <blockquote>
<p>it executes the backward method with <code>torch.autocast</code> if the forward part had it enabled.</p>
</blockquote>
<p>Ahh interesting, that would make sense.</p>
<p>I think the difference is that whatever you do in custom Function forward is opaque to the autograd engine, so it wouldn’t be able to do that automatic casting I mentioned in this particular case.</p>