Support CE after grad acc fix

[ByronHsu]

Summary

Based on #374, but make it leaner

1. The use of cross entropy in model code has changed after grad fix
2. It changed from module CrossEntropy to functional cross_entropy
3. Our monkey patching needs to change accordingly
4. While also make sure backward compatibility by adding a condition for different versions

Notable Changes

1. Add a functional api for CE to take keyword args
2. Add back conv test with logits to test CE convergence
3. Add back comp test for transformers 4.44

Testing Done

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

[hongpeng-guo]

Thanks for providing this exemplary PR! While I am working on enabling kwargs for all other operators, I meet a few questions and would like to hear your suggestions. 😄

For the function signature, should we try to follow their counter-parts' signatures in torch.nn.Function? I found there are a few cases that make it hard to make the signature compatible with torch, i.e:

1. rms_norm and layer_norm in torch ask for non-optional arg normalized_shape that is not required in liger;
2. group_norm in liger ask for non-optional num_channels which is not required in torch;
3. gelu in liger takes two args a, b as inputs, but for torch, it only asks for one input

There are also ops that is not available in torch.nn.Functional, i.e., the fused ops and jsd ops.

What would be a good strategy here to redefine the function signature here? cc @ByronHsu

[ByronHsu]

Thanks @hongpeng-guo! I think for now we don't need to be restricted by torch since some of the layers like rmsnorm is not actually taken from torch. Let's just put whatever arg liger has as kwargs and don't worry about torch. We can do some adjustment once we receive more community feedback.