Can you share a colab link for this? It would be easier to debug.
Some additional questions:

Do you use FP16?
How did you tuned the learning rate?
Do you use gradient clipping?
What is the batch size?