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]

So, the issue seems to be with combining label_smoothing with weighted loss. I've been staring at the code and equations for a while now but I can't pinpoint anything that's wrong. Simply multiplying the final loss with the weight of the label token seems like the right thing to do to me.

If not there can only be an issue with the:

scaled_x_sum term since all the other terms in smoothed loss are also a part of the plain ce loss which we know works correctly.

[Tcc0403]

Figuring out where it doesn't work is a big! I'll take a look on Saturday.

[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.

[Tcc0403]

I took a look at torch's impl, and here's how they compute smooth_loss
https://github.com/pytorch/pytorch/blob/2682e5e0d48a8200c1672b6a42250d3c8de44190/aten/src/ATen/native/LossNLL.cpp#L558

    if (weight.defined()) {
      // Expand weight to the correct number of dims for broadcasting with input / target
      auto weight_broadcast_shape = SmallBuffer<int64_t, 5>(input.dim());
      std::fill(weight_broadcast_shape.begin(), weight_broadcast_shape.end(), 1);
      weight_broadcast_shape[class_dim] = weight.size(0);
      Tensor weight_ = weight.view(weight_broadcast_shape);

      smooth_loss = -(input * weight_).sum(class_dim);

related code blocks in liger:

Liger-Kernel/src/liger_kernel/ops/cross_entropy.py

Line 119 in bd65c47

scaled_x_sum += tl.sum(tl.where(X_offsets < n_cols, -eps * X_block, 0.0))

Liger-Kernel/src/liger_kernel/ops/cross_entropy.py

Lines 194 to 196 in bd65c47

	if label_smoothing > 0:
	smooth_loss = scaled_x_sum + label_smoothing * lse
	loss = loss * (1 - label_smoothing) + smooth_loss

[Tcc0403]

we can rewrite it with torch.masked_select

sum_of_non_ignore_weight = (torch.gather(weight, dim=0, index=target.masked_select(target_mask))
            .sum()
            .item()
        )

Refer to torch's impl mentioned above

[winglian]

@pramodith anything I can do to help with this PR?

[pramodith]

@pramodith anything I can do to help with this PR?

Hey @winglian I won't be able to look into this any further, feel free to take over and see if you can figure out the source of discrepancy. The tests fail when combining smoothing loss with weighted ce.