Cow_woC:

detect_anomaly is already enabled which is how I discovered this problem to begin with.


What’s the output from this? It should give a bit more detail as to why it crashes? I assume it’s related to the question? LayerNorm’s grads become NaN? If so, which LayerNorm as you have multiple within your code.



 Cow_woC:

I already checked the input and output and neither contain NaN or Inf.


if there are no NaNs / Infs within your code is this just on the output of given Layers? You could try using register_full_backward_hook and see if any NaNs / Infs get passed back through that way. Given the NaN is in the gradient it’s most likely happening when computing your loss. You can create a function like
def _save_output(self, module, grad_input, grad_output):
  print(module, grad_output)

and then apply that function to all LayerNorm module within your network like,
for module in model.modules():
  module.register_full_backward_hook(_save_output)
#or you can manually place them of the LayerNorm modules yourself (in the loop doesn't work)

Then when you backprop your loss it’ll print out the gradient of the loss with respect to each layer’s output. That might help debug what layer (more specifically which LayerNorm in your case) is causing the NaN issue. Granted the gradient of your loss with respect to the parameters of a layer differs slightly to the grad_output variable, it’s still using in computing the gradient and if it has a NaN it’ll show you what Layer’s failing.



 Cow_woC:

The only thing that comes to mind is that I’ve got AMP mode enabled. After the first iteration the output gets mapped by a loss function to a value around ~100. During the backward pass AMP scales this loss value up to ~600,000.


Just to check you’re only using AMP within the forward pass? (Like the tutorial: Automatic Mixed Precision package - torch.cuda.amp — PyTorch 1.9.1 documentation). Backproping with AMP enabled might give rise to your NaNs? (Could perhaps try without using AMP in the backward pass and see if the LayerNorm.grad is NaN goes away?)