linkedin / linkedin/Liger-Kernel
JSD kernel returns NaN for interior beta when a vocabulary entry underflows under both distributions
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
### 🐛 Describe the bug
The interior-`beta` branch of `_jsd_kernel` forms the mixture in probability space and then takes its log:
https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/jsd.py
```python
Q = tl.exp(X_shifted) * exp_max # = exp(X)
P = tl.exp(Y_shifted) * exp_max # = exp(Y)
beta_P = beta * P
one_minus_beta_Q = (1 - beta) * Q
M = beta_P + one_minus_beta_Q
log_M = tl.log(M)
loss = beta_P * Y + one_minus_beta_Q * X - M * log_M
```
When a vocabulary entry's probability underflows to zero under **both** distributions, `M == 0.0`, so `log_M == -inf` and `M * log_M` evaluates `0.0 * -inf == NaN`. That `NaN` then propagates through the row sum and poisons the whole loss.
Empirically the threshold is a log-probability of roughly `-104` (i.e. below the smallest fp32 subnormal) under both sides. That is reachable in ordinary distillation with a wide logit spread, and the temperature divisor in `LigerFusedLinearJSD` pushes it closer for `temperature < 1`.
The `beta == 0.0` and `beta == 1.0` branches are unaffected — they compute `Y_prob * (Y - X)` and never take the log of a probability-space sum.
### Reproducer
```python
import math
import torch
from liger_kernel.ops.jsd import LigerJSDFunction
device = "cuda"
student_logits = torch.tensor([[0.0, -120.0, -1.0, -2.0]], device=device)
teacher_logits = torch.tensor([[0.5, -130.0, -2.0, -1.0]], device=device)
X = student_logits.log_softmax(-1) # log Q
Y = teacher_logits.log_softmax(-1) # log P
for beta in (0.0, 0.5, 1.0):
print(f"beta={beta}: liger={LigerJSDFunction.apply(X, Y, None, beta, -100).item()}")
# Reference with the mixture formed in log space.
beta = 0.5
log_M = torch.logaddexp(Y + math.log(beta), X + math.log1p(-beta))
ref = beta * (Y.exp() * (Y - log_M)).sum() + (1 - beta) * (X.exp() * (X - log_M)).sum()
print(f"beta=0.5: reference={ref.item()}")
```
Output:
```
beta=0.0: liger=0.13237953186035156
beta=0.5: liger=nan
beta=1.0: liger=0.18083927035331726
beta=0.5: reference=0.03686588257551193
```
The entry that underflows contributes ~0 to the true divergence, so the expected value is the finite `0.0369`.
### Suggested fix
Form the mixture in log space so the underflowing entry contributes an exact zero instead of `NaN`:
```python
log_beta_P = Y + tl.log(beta)
log_one_minus_beta_Q = X + tl.log(1 - beta)
max_log = tl.maximum(log_beta_P, log_one_minus_beta_Q)
log_M = max_log + tl.log(tl.exp(log_beta_P - max_log) + tl.exp(log_one_minus_beta_Q - max_log))
M = tl.exp(log_M)
loss = beta_P * Y + one_minus_beta_Q * X - M * log_M
dX = one_minus_beta_Q * (X - log_M)
```
`log_M` stays finite whenever either side is finite, and `M * log_M` is then a well-defined `0 * finite == 0` for the underflowing entries. A guard as simple as `log_M = tl.where(M > 0, tl.log(M), 0.0)` would also remove the `NaN`, though it is less accurate for entries that are small but not fully underflowed.
Happy to send a PR with the log-space version plus a regression test if that direction looks right.
### Versions
- `liger-kernel` 0.8.2 (the same code is on `main` as of today)
- `torch` 2.13.0+cu130, `triton` 3.7.1
- NVIDIA RTX A6000, CUDA 13.0
Contributor guide
Research direction
Start in src/liger_kernel/ops/jsd.py at the interior-beta branch of _jsd_kernel, then run the supplied CUDA reproducer to confirm the NaN. Form the mixture in log space and add a regression test for entries that underflow under both distributions; done means the interior-beta result is finite and the beta 0.0 and 1.0 branches remain unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100