fixed_fused_linear_cross_entropy should pass through kwargs

[npuichigo]

Summary

Testing Done

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

[Tcc0403]

The reason why we don't want to pass kwargs is because **kwargs also contains FlashAttentionKwrags, resulting in issue #650.
More detailed explanation #651 (comment)

[npuichigo]

So maybe we should pass through ce_weight.

[Tcc0403]

Sure, it should be fine to pass through other arguments, such as ce_weight, label_smoothing, lse_square_scale, etc.
However, they are not natively supported in hf transformers, I'm not sure about the namings for these liger flce specified kwargs.
https://github.com/huggingface/transformers/blob/947a37e8f5bc50bc0e9a77c0d16b038adcb056d0/src/transformers/utils/generic.py#L858

[JuanPZuluaga]

Hello, I think, it is also necessary to add kwargs in:

Liger-Kernel/src/liger_kernel/transformers/functional.py

Line 70 in 24416a4

def liger_fused_linear_cross_entropy(

I'm hitting a compilation error with torch.compile.

Basically:

def liger_fused_linear_cross_entropy(
    input,
    weight,
    target,
    bias=None,
    ce_weight=None,
    ignore_index: int = -100,
    lse_square_scale: float = 0.0,
    label_smoothing: float = 0.0,
    reduction: str = "mean",
    softcap: Optional[float] = None,
    return_z_loss: bool = False,
    accum_dtype=None,
    use_token_scaling: bool = False,
    return_token_accuracy: bool = False,
    **kwargs
):

fixes it.