先调用zero_grad()再调用step()致权重变NaN，该现象是否正常？

我编写了一段简化代码，加载图像并输入至仅含1个CNN层的模型中：

def main(cfg):
    model = Model().cuda()
    dataset = Dataset(cfg)
    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 idx, (target) in enumerate(train_dataloader, start=1):
            to_pil_image(target.squeeze(0)).save('test.jpg')

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

            target = target.to('cuda')
            target = F.interpolate(target, (256, 256), mode='bilinear', align_corners=False)
            output = model(target)
            loss = F.mse_loss(target, output)
            loss.backward()
            # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0)
            optimizer.zero_grad()
            optimizer.step()

我发现当先调用optimizer.step()再调用optimizer.zero_grad()时，代码运行正常，损失下降且模型收敛。但当调整为先调用zero_grad()再调用step()时，p.grad变为0（符合预期），但p[0, 0, 0, 0]会变为NaN。我的理解是先调用zero_grad()再调用step()应完全不更新权重，请问该现象是否属于预期行为？

这不属于预期行为，背后的原因和AdamW的特性有关：

• AdamW是自适应优化器，会维护一阶矩、二阶矩这类累积状态变量，即便梯度被zero_grad()置为0，step()操作依然会更新这些累积状态。结合初始权重、学习率等参数，可能触发数值不稳定，最终导致权重变为NaN。
• 标准训练流程是反向传播计算梯度 → 执行optimizer.step()更新权重 → 调用optimizer.zero_grad()清空梯度，颠倒顺序的操作本身不符合规范。当梯度为0时执行step()，优化器内部的状态更新逻辑会出现异常，进而引发权重的数值问题。
• 你的理解存在偏差：zero_grad()后调用step()并不是“完全不更新权重”。AdamW的更新不仅依赖当前梯度，还依赖之前迭代累积的状态信息，梯度为0仅会让当前梯度项的贡献为0，但累积的矩依然会参与计算，最终可能导致异常数值。

内容的提问来源于stack exchange，提问作者K.XGai