Can `autocast` handle networks with layers having different dtypes?

username: ronan.fruit
cooked: <p>Hi,</p>
<p><code>torch version = 2.5.0</code></p>
<p>I am wondering whether <code>torch.autocast</code> can handle neural networks with layers having different dtypes.</p>
<p>The following code suggests it cannot:</p>
<pre><code class="lang-auto">import torch
net = torch.nn.Sequential(
    torch.nn.Linear(2, 10, dtype=torch.float16),
    torch.nn.ReLU(),
    torch.nn.Linear(10, 10, dtype=torch.float32),
    torch.nn.ReLU(),
)
with torch.autocast("cuda"):
    net(torch.as_tensor([[1., 2.]], dtype=torch.float16))
</code></pre>
<p>=&gt; it raises <code>RuntimeError: mat1 and mat2 must have the same dtype, but got Half and Float</code></p>
<p>But at some point recently I had started to believe <code>autocast</code> was able to handle such cases.<br>
I loaded a pretrained model on HuggingFace and changed just the <code>dtype</code> of one layer. Without autocast, the same <code>RuntimeError</code> was raised, but with <code>autocast</code> the error disappeared.<br>
Here is a Minimum Reproducible Example:</p>
<pre><code class="lang-auto">from transformers import pipeline
pipe = pipeline(
    "text-classification", 
    model="Qwen/Qwen2.5-0.5B",
    torch_dtype=torch.float16,
    device_map="cuda"
)
pipe_model_named_parameters = {k: v for k, v in pipe.model.named_parameters()}
for p in pipe_model_named_parameters:
    if "score" in p:  # convert last trainable layer to float32 for stability during training
        pipe_model_named_parameters[p].data = pipe_model_named_parameters[p].data.to(dtype=torch.float32)
with torch.autocast("cuda"):  # context manager needed in case not all layers have the same dtype
    print(pipe(["a", "b", "z"]))
</code></pre>
<p>Any insight on what <code>autocast</code> allows and does not allow ?</p>

username: soulitzer
cooked: <p>Seems to work if you use cuda tensors with torch.autocast(“cuda”)</p>
<pre><code class="lang-auto">import torch
net = torch.nn.Sequential(
    torch.nn.Linear(2, 10, dtype=torch.float16),
    torch.nn.ReLU(),
    torch.nn.Linear(10, 10, dtype=torch.float32),
    torch.nn.ReLU(),
).to("cuda")
with torch.autocast("cuda"):
    net(torch.tensor([[1., 2.]], dtype=torch.float16, device="cuda"))
</code></pre>

username: Aknw_Fen
cooked: <p>I’m unsure that casting all layers manually <code>autocast</code> will then have any effect:</p>
<p>Autocast will respect <a href="https://pytorch.org/docs/stable/amp.html#autocast-op-reference" rel="noopener nofollow ugc">types assigned manually</a>:</p>
<blockquote>
<p>Ops called with an explicit <code>dtype=...</code> argument are not eligible, and will produce output that respects the <code>dtype</code> argument.</p>
</blockquote>

username: soulitzer
cooked: <p>The dtype argument here for nn.Linear should just be used to initialize the dtype of the weight<br>
<a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/linear.py" rel="noopener nofollow ugc">torch/nn/modules/linear.py</a></p>

username: ronan.fruit
cooked: <p>Thanks for your reply <a class="mention" href="/u/soulitzer">@soulitzer</a><br>
Indeed it seems the problem was caused by <code>net</code> not being on <code>cuda</code> device while <code>cuda</code> is specified in <code>torch.autocast</code>.<br>
Both options below work fine:</p>
<ol>
<li>adding <code>.to("cuda")</code> to <code>net</code> and specifying <code>device="cuda"</code> as <a class="mention" href="/u/soulitzer">@soulitzer</a> pointed out</li>
<li>replacing <code>torch.autocast("cuda")</code> by <code>torch.autocast("cpu")</code></li>
</ol>
<p>In case 1, the returned tensor has type <code>torch.float16</code> while in case 2 the returned tensor has type <code>torch.bfloat16</code>. Not sure why there is such a difference.</p>