How to convert submodules to a different dtype simply

username: divinho
cooked: <p>I want to train a model where some modules are in bf16 and others in float32.</p>
<p>If the model is in bf16, then it seems to me using <code>torch.autocast(enabled=False)</code> for the relevant modules is not enough as the weights will still be in bf16.</p>
<p>But I can’t just iterate through the named_modules() and change the dtype because <code>.to()</code> creates a copy.</p>
<p>Ideally I could use an inplace version like <code>.to_()</code> but that doesn’t seem to exist.</p>
<p>Is there a simple way to do this? (that doesn’t involve something like keeping track of the parent modules and then assigning)</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="divinho" data-post="1" data-topic="202825">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/9dc877/48.png" class="avatar"> divinho:</div>
<blockquote>
<p>But I can’t just iterate through the named_modules() and change the dtype because <code>.to()</code> creates a copy.</p>
</blockquote>
</aside>
<p>That’s not the case for <code>nn.Module</code>s as the <code>to()</code> operation will be applied recursively to all registered submodules, parameters, and buffers.</p>
<pre data-code-wrap="python"><code class="lang-python">model = models.resnet18()

for name, module in model.named_modules():
    if "bn" in name:
        module.double()
        
print([(name, p.dtype) for name, p in model.named_parameters()])
# [('conv1.weight', torch.float32), ('bn1.weight', torch.float64), ('bn1.bias', torch.float64), ('layer1.0.conv1.weight', torch.float32), ('layer1.0.bn1.weight', torch.float64), ('layer1.0.bn1.bias', torch.float64), ('layer1.0.conv2.weight', torch.float32), ('layer1.0.bn2.weight', torch.float64), ('layer1.0.bn2.bias', torch.float64), ('layer1.1.conv1.weight', torch.float32), ('layer1.1.bn1.weight', torch.float64), ('layer1.1.bn1.bias', torch.float64), ('layer1.1.conv2.weight', torch.float32), ('layer1.1.bn2.weight', torch.float64), ('layer1.1.bn2.bias', torch.float64), ('layer2.0.conv1.weight', torch.float32), ('layer2.0.bn1.weight', torch.float64), ('layer2.0.bn1.bias', torch.float64), ('layer2.0.conv2.weight', torch.float32), ('layer2.0.bn2.weight', torch.float64), ('layer2.0.bn2.bias', torch.float64), ('layer2.0.downsample.0.weight', torch.float32), ('layer2.0.downsample.1.weight', torch.float32), ('layer2.0.downsample.1.bias', torch.float32), ('layer2.1.conv1.weight', torch.float32), ('layer2.1.bn1.weight', torch.float64), ('layer2.1.bn1.bias', torch.float64), ('layer2.1.conv2.weight', torch.float32), ('layer2.1.bn2.weight', torch.float64), ('layer2.1.bn2.bias', torch.float64), ('layer3.0.conv1.weight', torch.float32), ('layer3.0.bn1.weight', torch.float64), ('layer3.0.bn1.bias', torch.float64), ('layer3.0.conv2.weight', torch.float32), ('layer3.0.bn2.weight', torch.float64), ('layer3.0.bn2.bias', torch.float64), ('layer3.0.downsample.0.weight', torch.float32), ('layer3.0.downsample.1.weight', torch.float32), ('layer3.0.downsample.1.bias', torch.float32), ('layer3.1.conv1.weight', torch.float32), ('layer3.1.bn1.weight', torch.float64), ('layer3.1.bn1.bias', torch.float64), ('layer3.1.conv2.weight', torch.float32), ('layer3.1.bn2.weight', torch.float64), ('layer3.1.bn2.bias', torch.float64), ('layer4.0.conv1.weight', torch.float32), ('layer4.0.bn1.weight', torch.float64), ('layer4.0.bn1.bias', torch.float64), ('layer4.0.conv2.weight', torch.float32), ('layer4.0.bn2.weight', torch.float64), ('layer4.0.bn2.bias', torch.float64), ('layer4.0.downsample.0.weight', torch.float32), ('layer4.0.downsample.1.weight', torch.float32), ('layer4.0.downsample.1.bias', torch.float32), ('layer4.1.conv1.weight', torch.float32), ('layer4.1.bn1.weight', torch.float64), ('layer4.1.bn1.bias', torch.float64), ('layer4.1.conv2.weight', torch.float32), ('layer4.1.bn2.weight', torch.float64), ('layer4.1.bn2.bias', torch.float64), ('fc.weight', torch.float32), ('fc.bias', torch.float32)]
</code></pre>

username: divinho
cooked: <p>Ah cool! Thank you!</p>
<p>What if I wanted to convert specific parameters? When I try and assign like</p>
<pre><code class="lang-auto">module.weight = module.weight.to(torch.bfloat16)
</code></pre>
<p>I get</p>
<pre><code class="lang-auto">TypeError: cannot assign 'torch.cuda.BFloat16Tensor' as parameter 'weight' (torch.nn.Parameter or 
None expected)
</code></pre>

username: ptrblck
cooked: <p>Create a new <code>nn.Parameter</code> as the error message explains.</p>

username: divinho
cooked: <p>Might it make more sense to apply <code>.to()</code> to <code>weight.data</code> ? Would that allow me to skip recreating the parameter?</p>

username: divinho
cooked: <pre><code class="lang-auto">module.weight = nn.Parameter(module.weight.to(torch.bfloat16), requires_grad=False)
</code></pre>
<p>this doesn’t seem to work as expected, i get OOM while doing this (which implies this is not overwriting the old parameter weights)</p>

username: divinho
cooked: <p>using <code>.data</code> seems to do the trick</p>