CrossEntropyLoss return single value when reduction is "none"

[leng-yue]

🐛 Describe the bug

When using LigerCrossEntropyLoss with reduction="none", it returns a single value, instead of an array.

Reproduce

import torch
from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss


def test_cross_entropy_matches_torch():
    # Test cases with various shapes and values
    batch_size = 3
    num_classes = 5
    seq_length = 4

    # Generate random logits and labels using torch
    torch.manual_seed(42)
    logits = torch.randn(batch_size, seq_length, num_classes, device='cuda')
    labels = torch.randint(0, num_classes, (batch_size, seq_length), device='cuda')

    # Set some labels to -100 to ignore them
    labels[0, -1] = -100  # Ignore last token of first sequence
    labels[1, -2:] = -100  # Ignore last two tokens of second sequence

    # Liger cross entropy
    liger_ce = LigerCrossEntropyLoss(reduction="none")
    liger_loss = liger_ce(logits.view(-1, num_classes), labels.view(-1))
    print(liger_loss)

    # PyTorch cross entropy
    torch_ce = torch.nn.CrossEntropyLoss(reduction="none")
    torch_loss = torch_ce(logits.view(-1, num_classes), labels.view(-1))

    # Check if losses match
    torch.testing.assert_close(
        liger_loss,
        torch_loss,
        rtol=1e-5,
        atol=1e-5,
        msg="Liger and PyTorch losses don't match",
    )

    print("All tests passed! LigerCrossEntropyLoss matches PyTorch's implementation")


if __name__ == "__main__":
    test_cross_entropy_matches_torch()

Versions

Environment Report:

Operating System: Linux-5.15.0-88-generic-x86_64-with-glibc2.35
Python version: 3.10.15
PyTorch version: 2.5.1+cu124
CUDA version: 12.4
Triton version: 3.1.0
Transformers version: 4.46.2

[ByronHsu]

Sorry. We only support "mean" and "sum" for now. Are you willing to contribute by adding an assertion or implementing "none" reduction? Thanks!

[leng-yue]

I think I can make a PR to add an assert to block "none" first, then let's see how to implement the "none" reduction one.

[Tcc0403]

The triton kernel of liger ce is actually giving none reduction loss, just need to make a reduction=="none" condition to output loss1d directly without torch.sum()

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

Line 288 in 79b940f

loss = torch.sum(loss_1d)

[hebiao064]

#take I made a PR about this issue, please take a look, thanks @ByronHsu @leng-yue @Tcc0403

[leng-yue]

Looks good to me