[WIP] Fix confusion on Gemma

[yundai424]

Summary

why

1. From config point of view, Gemma1 is doing exact GeLU [ref] but gemma 1.1 and 2 are doing approximate gelu (gemma2, gemma1)
2. also, gemma uses hidden_activation config field and hidden_act is ignored (gemma1 code, gemma2 code)

That being said, we should be fine to claim that all of gemma 1, 1.1 and 2 are supported. But for safety i think we can first go with 1.1 and 2 first

what

1. adjust the monkey patch so it works properly with gemma2 code base too
2. checkstyle
3. [TODO] add final logit softcapping for gemma2 https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modeling_gemma2.py#L1054

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[qingquansong]

LGTM! Thanks for the fix! Wondering if we could relax the checking here so we can also use Gemma1 ?

[DocShotgun]

It looks like the same recently added lce_forward is being used for both Gemma and Gemma2 (#111).

It appears that there is a slight difference in the forward between Gemma and Gemma2 in the modeling code in transformers, specifically regarding logit softcapping (https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modeling_gemma2.py#L1054-L1057), which doesn't seem to be accounted for in the new lce code as far as I can tell. Would this potentially lead to incompatibility?

[yundai424]

yeah exactly, just figured out we need to add softcapping :/ i'll have a follow up PR to clarify that right now this only works for gamma 1 and 1.1 and we'll work on changing the kernel correspondingly asap