Add weight support for LigerCrossEntropy

[Tcc0403]

Summary

Resolve #404.

TODO:

• (RFC) Expose weight paramter at LigerFusedLinearCrossEntropyLoss, but we need to consider renaming some variables to distinguish weight of linear layer and weight for ce.
• Add unit test for FLCE after exposing weight

Testing Done

It hasn't fully tested with other params.

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

[pramodith]

Thanks for taking care of this! Had a few minor suggestions.

Another TODO is based on the original paper linked in the original issue for this feature. We also need to support a sample level weight. i.e. a weight that can be applied to each element of the batch if we have logits in the shape (B, S, V). We'd have sample level weights of shape (B, ). This is what's proposed in the C-RLFT paper. https://arxiv.org/abs/2309.11235

[Tcc0403]

Feel free to push to this branch or even take over it and open a new PR, I won't be able to update that often in the next few months. Just trying to make the first step when I got time.

[Tcc0403]

This test couldn't pass somehow. I might miss something.

[pramodith]

Feel free to push to this branch or even take over it and open a new PR, I won't be able to update that often in the next few months. Just trying to make the first step when I got time.

Gotcha! I'll try wrapping it up, you've done most of the heavy lifting already.