Fix mean subtraction in layer norm kernels

[nhamanasu]

Summary

This PR is the possible proposal to the Issue: #518

I've confirmed neither that the issue is really a issue nor that this PR can pass tests, so made this as a draft PR.

I'd like someone interested in this to check them please... 🙇

[Tcc0403]

Let's increase hidden_size in the unit test of layernorm to see whether the issue occurs.

[nhamanasu]

Thank you for the comment!

I guess this might be a problem when hidden_size is not the power of 2 (e.g., 1536) rather than in case of just a huge hidden_size, since in such cases mask will have False elements.

[Tcc0403]

Yes, you're right. The scenario you mentioned would probably happen when n_cols is less than BLOCKSIZE (not power of 2). Feel free to add weird shape to layernorm unit test like other kernels have.

[Tcc0403]

I've tested the main branch with an additional weird shape test case, can confirm that the issue exists.

❯ python3 -m pytest test/transformers/test_layer_norm.py
================================================= test session starts ==================================================
platform linux -- Python 3.10.12, pytest-8.3.4, pluggy-1.5.0
rootdir: /home/tcc/Liger-Kernel
configfile: pyproject.toml
plugins: rerunfailures-15.0, xdist-3.6.1
collected 5 items

test/transformers/test_layer_norm.py::test_liger_layer_norm[dtype0-1e-05-1e-05-2-8-64] PASSED                    [ 20%]
test/transformers/test_layer_norm.py::test_liger_layer_norm[dtype0-1e-05-1e-05-4-16-128] PASSED                  [ 40%]
test/transformers/test_layer_norm.py::test_liger_layer_norm[dtype0-1e-05-1e-05-4-256-1234] FAILED                [ 60%]
test/transformers/test_layer_norm.py::test_liger_layer_norm_functional[dtype0-1e-05-1e-05-2-8-64] PASSED         [ 80%]
test/transformers/test_layer_norm.py::test_liger_layer_norm_functional[dtype0-1e-05-1e-05-4-16-128] PASSED       [100%]

[Tcc0403]

Enhance unit test by adding a weird shape to test cases, then we can merge it!

[nhamanasu]

I really appreciate your elaborating on this topic! I'd also check this and will commit the weird-shape test cases. Please wait for a while 🙇