Vocab-parallel cross entropy with label smoothing returns inf when a logit is far below the row max #1618

With label_smoothing > 0, LigerVocabParallelCEFunction.forward takes torch.log(exp_buf) to get the log-probabilities (ops/vocab_parallel_cross_entropy.py, line 304). exp_buf holds exp(x - max) in fp32, which is 0 once a logit is about 104 below the row max, so the loss is inf where F.cross_entropy is finite. A vocabulary masked with a large negative value gives such rows. The gradient is not affected.

import torch
import torch.nn.functional as F

from liger_kernel.transformers import LigerVocabParallelCrossEntropy

x = torch.tensor([[[0.0, -200.0, 1.0, 2.0]]], device="cuda")  # [seq, batch, vocab]
y = torch.tensor([[2]], device="cuda")

for label_smoothing in (0.0, 0.1):
    liger = LigerVocabParallelCrossEntropy(label_smoothing=label_smoothing)(x, y)
    ref = F.cross_entropy(x.view(-1, 4), y.view(-1), label_smoothing=label_smoothing)
    print(f"label_smoothing={label_smoothing}: liger {liger.item():.6f}  torch {ref.item():.6f}")

label_smoothing=0.0: liger 1.407606  torch 1.407606
label_smoothing=0.1: liger inf  torch 6.432606

Run on main (42dc23c) with an NVIDIA L40S, torch 2.15.0a0 and Triton 3.8.0. LigerMegatronCrossEntropy uses the same code. That line also builds two [BT, V_local] fp32 temporaries.