mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Michaël DefferrardandVaibhav Jindal
5b828eef2b
## Summary Fix NaNs in generalized JSD when both student and teacher probabilities underflow to zero in FP32. The previous implementation constructed the mixture distribution in probability space. For log probabilities around `-104` or lower, both probabilities could underflow to zero, resulting in `log(0)` and NaN loss terms. This change: - computes `log(M)` using an element-wise log-sum-exp formulation - evaluates the loss as weighted KL terms to avoid `0 * -inf` - applies the fix consistently to the Triton, Ascend, and cuTile backends - updates the PyTorch reference to use the equivalent stable formulation `torch.logaddexp` - adds standalone and fused-linear regression tests Both new regression tests return NaN on `main` and pass with this change. ## Testing Done Hardware Type: A100 - [x] run `make test` to ensure correctness (18 unrelated chunked GRPO LUSPO failures) - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence (does not exercise JSD) Ascend and cuTile received the equivalent backend change but were not runtime-tested as I don't have the hardware. Co-authored-by: Vaibhav Jindal <[email protected]>