linkedin / linkedin/Liger-Kernel
LigerTVDLoss returns a non-zero gradient where p == q
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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