LigerTVDLoss returns a non-zero gradient where p == q

[truong-v]

🐛 Describe the bug

Total variation distance is 0.5 * sum |p - q|, so its derivative with respect to p is 0.5 * sgn(p - q), which is 0 wherever the two distributions agree. The Triton kernel computes the gradient with a two-way tl.where(p > q, ...) that has no zero branch, so p == q falls into the p < q branch and every tied element gets -0.5 * scale instead of 0. The loss itself is correct — only the gradient is wrong, and only on ties.

Two configurations hit it with default settings: self-distillation, where student and teacher are identical at the start of training (the loss is 0.0 but every element of the gradient is -0.5 / batch), and padded vocabulary slots, which softmax to exactly 0.0 in both distributions and then get a spurious gradient pushing them up. The reference the test suite compares against, TorchTVDLoss in test/transformers/test_tvd.py (torch.abs(p - q) / 2.0), has autograd gradient 0 at a tie — so the suite's own reference already disagrees with the kernel, and the tests never catch it only because both distributions are drawn at random and never tie.

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

Lines 88 to 89 in a5d795e

	# Fuse reduction scaling into gradient computation (eliminates separate Python division)
	grad_res = tl.where(p > q, 0.5 * scale, -0.5 * scale)

The kernel was added in #324.

Reproduce

import torch

from liger_kernel.transformers.tvd import LigerTVDLoss

torch.manual_seed(0)
p = torch.randn(8, 512, device="cuda").softmax(dim=-1)  # student
q = p.clone()  # teacher: identical at step 0 of self-distillation

x_liger = p.clone().requires_grad_(True)
x_torch = p.clone().requires_grad_(True)

LigerTVDLoss(reduction="batchmean")(x_liger, q).backward()
(torch.abs(x_torch - q) / 2.0).sum().div(x_torch.size(0)).backward()  # TorchTVDLoss from test_tvd.py

print("loss                :", LigerTVDLoss()(p, q).item())
print("liger grad [0, :4]  :", x_liger.grad[0, :4].tolist())
print("torch grad [0, :4]  :", x_torch.grad[0, :4].tolist())
print("elements wrong      :", int((x_liger.grad != x_torch.grad).sum()), "/", x_liger.grad.numel())

loss                : 0.0
liger grad [0, :4]  : [-0.0625, -0.0625, -0.0625, -0.0625]
torch grad [0, :4]  : [0.0, 0.0, 0.0, 0.0]
elements wrong      : 4096 / 4096

The same happens with reduction="sum" and reduction="mean", and with a partial tie (padded vocabulary): only the tied columns are wrong.

Versions

• Liger-Kernel commit a5d795efd2c1436549e70118ef519134e9c27833 (main), editable install
• GPU: NVIDIA H100 NVL
• Liger Kernel version: 0.8.1
• PyTorch version: 2.6.0+cu124, CUDA 12.4
• Triton version: 3.2.0
• Transformers version: 5.14.1
• Python 3.10.20, Linux 5.15.0