NaN on longer training sequence

[coconutruben]

Hi there,

We were trying the kernel(s) for a longer training sequence and observe the following behavior after a few 10ks of steps:

• when everything is in fp32 (NATTEN, and the rest of the model, including normalization, etc) everything is fine
• when NATTEN is in bf16, and everything else is in fp32, then the system runs, but the metrics start diverging and getting worse
• when everything is in bf16, we run into NaN issues

Have you observed anything like this before? do you know whether there might be a bug in the backward/gradient accumulation or handling, or something in the dtype versions/accumulation?

Thank you!

[alihassanijr]

Hello!

Not really; I've trained with some pretty large problem shapes in bf16 specifically, using our older Ampere FNA, Hopper FNA, and recently Blackwell FNA.

Can you share more information on your use case and the specific backend you're using?

If you're on Hopper or Blackwell DC, can you try using our Ampere kernel to see if the issue goes away? That would help narrow it down or rule some options out.

You could also switch to full self attention (if you use max kernel size and no causal masking, natten will just reroute to self attention kernels) as well to rule some things out.

We generally do large sweeps over randomly generated use cases, so it could be an integration bug and not a kernel bug.

If you're typecasting the tensors manually, it might be an autograd issue.