Loss Explodes to NaN in large batch-size

[IsNoobgrammer]

🐛 Bug

When using larger batch size on v4-8 , the loss after 2 optimizer step becomes nan ;
This issue primarily due to newer torch-xla and libtpu ;

To Reproduce

Here is the Notebook you can use for quick repro

Environment

• Reproducible on XLA backend [TPU (v4-8)]:
• torch-xla 2.7.0+gitc76e949 (can occur in any torch-xla after 2.5.0)
• libtpu 0.0.10.dev20250210+nightly

Additional context

As mentioned earlier #8591 and #8683 (comment)

However even after trying recommended choice we were not able to mitigate it

Tagging @tengyifei @zpcore !! Thanks for your efforts