[torch-xla 2.9] Accuracy regressions for several large models

[jeffhataws]

🐛 Bug

We are seeing mixed-precision accuracy regressions for several large models, including mixtral and llama4.

We have narrowed the regression down to the change #9663.

To Reproduce

Will try the public tutorial: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-inference/tutorials/trn1-llama3.1-70b-instruct-accuracy-eval-tutorial.html

Expected behavior

No regression in accuracy compared to 2.8

Environment

• Reproducible on XLA backend [CPU/TPU]: Neuron
• torch_xla version: 2.9

Additional context

[ysiraichi]

This is actually expected due to the nature of the change.
I introduced this change in #7130 so as to make PyTorch CUDA and PyTorch/XLA more comparable (operations-wise).
More discussion can be seen in #8545.
A simple fix would be to promote the inputs to a higher precision manually.

[jeffhataws]

Thanks @ysiraichi . Could you explain what do you mean by "promote the inputs to a higher precision manually"? Is that a user change?

[ysiraichi]

Yes. Adding, for example, tensor.to(torch.float32) and converting back.

[jeffhataws]

@ysiraichi @qihqi Let's also revert for 2.9 since it's late, even though it's for a bug that we filed (#8545), and compiler is not able to fix this in time. See #9701 .

[ysiraichi]

Instead of reverting it, why don't we make it backend-specific, as suggested in #8545?

[jeffhataws]

We can check. Do you know what use_opmathtype_for_compute does and where the code for that is?

[jeffhataws]

Found it https://github.com/pytorch/xla/blob/d0b43ccdc0fa789ace63d66e7b4cdaf7c2904a1a/torch_xla/csrc/aten_xla_type.cpp#L181C41-L181C51 . Is this for all math ops or just torch.mul?