GeGLU kernel numerical issue

[Tcc0403]

🐛 Describe the bug

The current Liger geglu kernel is not close to torch's implementation, even in fp32. test_geglu.py can't pass with lower tolerances.

Some models couldn't pass the convergence test with low tolerance due to this issue. We need to explore whether there is a better implementation that is closer to torch's.

Reproduce

No response

Versions

Liger Kernel version: 0.6.4
PyTorch version: 2.7.1+cu126

[konstantinos-p]

Having a look at this, in tandem with chatgpt. Fyi if anyone else is interested:

• The pi value is hardcoded to float64 while the tests are 32 and 16: I tried hardcoding to other precisions but that didn't seem to work
• The order of operations could be wrong, in that multiplication is not commutative and could differ in the internal torch implementation.
• The tanh implementation could be different, since the current code uses the triton version.

[Tcc0403]

Thanks for initiating investigations!

The pi value is hardcoded to float64 while the tests are 32 and 16: I tried hardcoding to other precisions but that didn't seem to work

Liger hardcodes sqrt(2 / pi) intead of pi, could it be the reason? The goal is to pass fp32 test with lower precision, we can focus on fp32 geglu first.

The order of operations could be wrong, in that multiplication is not commutative and could differ in the internal torch implementation.

Yes, we can try using terms and contants in torch's geglu fashion.