Revert "mul: remove opmath cast sequence (pytorch#9663)"

[rajkthakur]

Commit 2a9138a removed .use_opmathtype_for_compute() from element-wise 'mul' operation, this breaks mixed-precision accumulation behavior expected by the Neuron compiler that traces/compile on CPU and later execute the binary on neuron hardwares, causing significant accuracy degradation in:

• Llama 3.1 70B models (16.7% throughput drop, accuracy failures)
• Mixtral 8x22B models (accuracy test failures)
• Other transformer models using mixed-precision compilation

Reverts: commit 2a9138a, other changes are result of rebase from r2.9
Fixes: Model accuracy failures with mixed-precision accumulation #9699

[rajkthakur]

Is the torchax test failure expected?

[ysiraichi]

• The TorchAX CI could be fixed by this change
• Instead of reverting the commit, I think it's better to make it backend specific, and land it on master and cherry-pick it to r2.9?

What do you think?

[rajkthakur]

While I understand the suggestion from #8545 to make this backend-specific, I believe a full revert is more appropriate atleast for release branch r2.9:

1. The change affects the fundamental numerical behavior of torch.mul, which is called extensively throughout any model. Even small precision differences compound in:
   • Attention mechanisms with large sequence lengths
   • Gradient accumulation/reduction over many steps and likely XLA based compilers need to explicitly handle this scenario.

2. Risk vs. benefit:
   • Risk: Our testing shows concrete accuracy regressions and are blocking for production workloads(mixtral, llama) on current compiler, which I guess might impact other users too.
   • Benefit: The original issue [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 was a question about the seeming unnecessary upcast/downcast that also appears in PyTorch CUDA. The upcast/downcast can be already removed by using PyTorch autocast so there's no need to fix [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 for r2.9. We can keep this in master branch for further investigation.

Let me know your thoughts?

[ysiraichi]

Thank you for your detailed analysis.

I'm sorry, but I still think that the best way to go about this is to make it backend specific. Actually, now that I'm thinking about it, I think it would make more sense to do it in the lowering step.

That said, since we don't really want to introduce a lot of changes here, and the infrastructure is already in place, I would say that we should make that use_opmathtype_for_compute() call backend specific.

It's a simple change (something like this condition for TPU) that fixes this problem for the Neuron backend, while leaving the other backends with the "backend-independent" choice of not doing that.

2. The original issue [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 was a question about the seeming unnecessary upcast/downcast that also appears in PyTorch CUDA. The upcast/downcast can be already removed by using PyTorch autocast so there's no need to fix [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 for r2.9. We can keep this in master branch for further investigation.

That "seeming unnecessary upcast/downcast" is exactly what we are talking about, here. And, no, PyTorch autocast does not solve this issue (see this comment).

[ysiraichi]

Thank you for the PR.
Let me know if you have any questions.

[ysiraichi]

Could you rebase this PR, so that there are only your commits?

[rajkthakur]

Could you rebase this PR, so that there are only your commits?

updated

[ysiraichi]

LGTM.
Let's just wait for the CI.

[rajkthakur]

@ysiraichi CI is completed.

[jeffhataws]

Thanks @rajkthakur @ysiraichi . It is merged now and ready to go. @bhavya01 will make a final 2.9 release candidate.