Fix Gemma RMSNorm

[davidgonmar]

Summary

Fixes #74. Allows Gemma's RMSNorm by adding a generic offset parameter to RMSNorm.

Testing Done

Parametrize the RMSNorm tests to test both Llama's and Gemma's versions (as per Hugging Face's transformers library).

• Hardware Type: tested on NVIDIA L4/L40
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

Test outputs

test/transformers/test_rms_norm.py ................................                                                                                                                                                                                            [100%]

========================================================================================================================= 32 passed in 3.74s =========================================================================================================================

test/convergence/test_mini_models.py ......                                                                                                                                                                                                                    [100%]

========================================================================================================================= 6 passed in 32.04s =========================================================================================================================

[ByronHsu]

Awesome! Our first kernel related PR :-) Left a comment

[ByronHsu]

@yundai424 I am curious if we need to do the upcasting too in our kernel? or we don't because anyway it gets casted to x.dtype. Our old impl doesn't have that but the result is still consistent

[davidgonmar]

@yundai424 I am curious if we need to do the upcasting too in our kernel? or we don't because anyway it gets casted to x.dtype. Our old impl doesn't have that but the result is still consistent

The test results seem to be consistent with the reference too. I think it's a tradeoff between complexity (because some models have slight inconsistencies) and exact reproduction. Let me know what you think and I can modify it (or even on another PR, to keep things clean)

[ByronHsu]

lgtm! waiting @yundai424 to final check

[ByronHsu]

can we add gemma to convergence test as well?

[davidgonmar]

can we add gemma to convergence test as well?

adding it

[davidgonmar]

can we add gemma to convergence test as well?

done. as mentioned in the discord channel, I made the tolerance for bfloat16 convergence test a bit bigger, since the casting difference seems to be slightly affecting it (for the regular tests it's fine, but it seems to add up in the end-to-end training situation). if we end up deciding to match it one-to-one in the future, we can make it stricter. let me know what you think

[yundai424]

my read is yes we need to change the dtype to x to fp32. Regardless of if it's llama or gemma they both do the norm part in fp32, only difference is gemma doing the scaling in fp32 too while llama do it in mixed precision. But of course we can do it in a separate PR 🤔

[ByronHsu]

May you join our discord and say hi? The PR looks very solid and we would love more contribution from you https://discord.gg/CX2YmNmn

[ByronHsu]

Also can you add your hardware type in the PR description? I just updated PR template

[yundai424]

LGTM, thanks for the contribution! Let's create a follow-up issue for the dtype thing