`zero_grad` before `step` causes gradient explosion?

retsuh: <p>I have this simplified code snippet, which loads an image and feed to a model of 1 CNN layer.</p>
<pre><code class="lang-auto">import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 3, 3, padding='same').half()

    def forward(self, x):
        return self.conv(x)


def main(cfg):
    model = Model().cuda()
    dataset = cfg.dataset
    optimizer = optim.AdamW(model.parameters(), lr=cfg.learning_rate)

    train_dataloader = DataLoader(
        dataset, 
        batch_size=cfg.batch_size, 
        num_workers=cfg.num_workers,
        shuffle=False,
        pin_memory=True    
    )
    
    p = next(model.parameters())
    for epoch in range(cfg.max_epochs):
        for image, _ in train_dataloader:

            print(p[0, 0, 0, 0])

            image = image.to('cuda').half()
            image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)

            output = model(image)
            loss = F.mse_loss(output, image)

            loss.backward()
            optimizer.zero_grad()
            optimizer.step()


class TestConfig:
    max_epochs = 10000
    root = "./temp/"
    batch_size = 1
    num_workers = 0
    transform = transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
     )
    dataset = torchvision.datasets.CIFAR10(
        root='/mnt/HDD3/khanh/temp/',
        train=True,
        download=True,
        transform=transform    
    )
    dtype = torch.float16
    device = 'cuda'
    learning_rate = 1


if __name__ == '__main__':
    cfg = TestConfig()
    main(cfg)
</code></pre>
<p>I noticed that when calling <code>optimizer.step()</code> then <code>optimizer.zero_grad()</code>, the code works properly (loss decreases and the model converge).</p>
<p>But when I call <code>zero_grad()</code> then <code>step()</code>, then <code>p.grad</code> will be 0 after <code>zero_grad()</code> (which is expected), but <code>p[0, 0, 0, 0]</code> will become nan after <code>step()</code>.</p>
<p>Is this an expected behaviour? Since to my understanding, calling <code>zero_grad()</code> before <code>step()</code> should have the effect of not updating the weights at all.</p>

ptrblck: <p>Could you post a minimal and executable code snippet reproducing the issue?</p>

retsuh: <p>Hi, I have edited the full, executable code snippet.</p>

retsuh: <p>It seems that the problem is a combination of multiple reasons. The issue disappears in these 2 cases:</p>
<ol>
<li>The <code>dtype</code> is set to <code>torch.float32</code> instead of <code>torch.float16</code>.</li>
<li><code>AdamW</code> is replaced with <code>SGD</code> or is set with a large <code>eps</code> value (currently it works with <code>eps=0.1</code>).</li>
</ol>

ptrblck: <p>I cannot reproduce the issue using your code, but I also executed it for 4 epochs only instead of the specified 10k.</p>
<aside class="quote no-group" data-username="retsuh" data-post="4" data-topic="201574">
<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/retsuh/48/66213_2.png" class="avatar"> retsuh:</div>
<blockquote>
<p>The <code>dtype</code> is set to <code>torch.float32</code> instead of <code>torch.float16</code>.</p>
</blockquote>
</aside>
<p>Manually using <code>float16</code> can easily cause overflows and is thus not recommended. Use <code>torch.cuda.amp</code> instead for mixed-precision training.</p>
<aside class="quote no-group" data-username="retsuh" data-post="4" data-topic="201574">
<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/retsuh/48/66213_2.png" class="avatar"> retsuh:</div>
<blockquote>
<p><code>AdamW</code> is replaced with <code>SGD</code> or is set with a large <code>eps</code> value (currently it works with <code>eps=0.1</code>).</p>
</blockquote>
</aside>
<p>Stateful optimizers (such as <code>AdamW</code>) will use running stats to update parameters even if their gradient is zero so you might want to check if these stats are overflowing for some reason.</p>