最近在用PyTorch复现一个ViT做图像分类，训练到第30个epoch左右loss突然就变nan了，学习率用的是1e-4，AdamW，batch size 64。奇怪的是重新跑一遍同样的代码，有时候能在50个epoch都正常，有时候20个epoch就炸了。梯度裁剪也加了，数据也检查过没有脏数据。想问下这种随机性的nan一般是什么原因导致的？跟混合精度有关系吗，我开了amp但没做loss scaling。
最近在用PyTorch复现一个ViT做图像分类，训练到第30个epoch左右loss突然就变nan了，学习率用的是1e-4，AdamW，batch size 64。奇怪的是重新跑一遍同样的代码，有时候能在50个epoch都正常，有时候20个epoch就炸了。梯度裁剪也加了，数据也检查过没有脏数据。想问下这种随机性的nan一般是什么原因导致的？跟混合精度有关系吗，我开了amp但没做loss scaling。
amp没做loss scaling基本就是元凶了，fp16下梯度很容易underflow，偶尔炸一次太正常了。你加个GradScaler试试，顺便看看ViT里attention那块的softmax有没有溢出。我之前也遇到过类似情况，最后发现是某个attention层的logits太大，加了scaling之后稳定多了。