linkedin / linkedin/Liger-Kernel

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

Open Beginner friendly
#1,373 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

## 🐛 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.

https://github.com/linkedin/Liger-Kernel/blob/a5d795efd2c1436549e70118ef519134e9c27833/src/liger_kernel/ops/tvd.py#L88-L89

The kernel was added in #324.

## Reproduce

```python
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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at src/liger_kernel/ops/tvd.py lines 88-89 and compare the Triton gradient behavior with TorchTVDLoss in test/transformers/test_tvd.py. Reproduce the all-tied case from the issue, add coverage for exact and partial ties, and run the TVD tests. Done means tied elements have zero gradient while existing non-tied behavior and reductions remain correct.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
84/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.