mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Yulong Cao
f6385e909b
## Summary
Add `FusedLinearKLDivLoss`, which fuses the vocabulary-head linear
projection with
KL(Q || P) loss against an explicit target distribution, for knowledge
distillation
workflows where teacher probabilities are available directly (e.g.
cached / offline
soft labels), complementing `LigerKLDIVLoss` (element-wise, no fusion)
and
`LigerFusedLinearJSD` (teacher side needs its own projection).
The kernel follows the chunked gradient-in-forward structure of
`fused_linear_jsd.py`,
so the full `BT x V` logits tensor is never materialized, and adopts the
`_CHUNK_MEM_CONST = 16` token-chunk geometry of
`fused_linear_cross_entropy.py`.
## Details
- Supports `reduction` (`"batchmean"` / `"mean"` / `"sum"`, same
semantics as
`torch.nn.KLDivLoss`), `ignore_index` via `shift_labels`, softmax
`temperature`,
and an `eps` clamp so that `0 * log(0)` contributes 0.
- Supports `accum_dtype` for fp32 weight-gradient accumulation across
chunked
low-precision GEMMs (same semantics as `LigerFusedLinearJSD` /
`LigerFusedLinearCrossEntropyLoss`), including the
`torch.addmm(out_dtype=fp32)`
fast path on torch >= 2.8 / sm80+ and a fallback elsewhere.
- The log-softmax backward is computed in place on the saved log-probs
and the
`chunk_size x V` fp32 temporaries are freed eagerly, so the peak during
the
trailing GEMMs stays low.
- `reduction="none"` is intentionally unsupported (gradient-in-forward
requires a
scalar loss); a `ValueError` points users to `LigerKLDIVLoss` instead.
- bf16 + `reduction="sum"` test tolerances follow the existing
convention in
`test_fused_linear_cross_entropy.py`.
- Peak-memory improvement over the unfused PyTorch baseline (full
fwd+bwd,
`max_memory_allocated`, bf16, llama-3 head shape H=4096 / V=128256, RTX
4060
Laptop GPU): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB
(1.83x);
at BT=1024 the benchmark harness reports 4793 -> 4244 MB. Runtime is at
parity
with the baseline at BT=1024 (0.89x).
## Testing Done
- Hardware Type: NVIDIA GeForce RTX 4060 Laptop GPU (sm89); NVIDIA Tesla
T4 (sm75)
- [x] run `make test` to ensure correctness (32 pre-existing failures
unrelated to
this change: 18 `test_mlp.py` shared-memory OOR on sm89 and 14
`test_grpo_loss.py`
LUSPO numerics; both reproduce on `main` without this change)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence (does not
exercise the new
standalone loss)
<details>
<summary>Test logs (RTX 4060 Laptop, torch 2.11.0+cu128, triton
3.7.1)</summary>
```
make test: 32 failed, 4152 passed, 914 skipped, 15 xfailed in 1050.50s (0:17:30)
test/transformers/test_fused_linear_kl_div.py: 172 passed, 0 failed, 0 skipped
failures (all pre-existing on this GPU, reproduced on main without this change):
18 x test/transformers/test_mlp.py
triton.runtime.errors.OutOfResources: shared memory
Required: 131072, Hardware limit: 101376 (sm89 has 100 KB usable smem)
14 x test/chunked_loss/test_grpo_loss.py
chunked GRPO LUSPO numerics (also seen by other contributors, see #1397)
make checkstyle: ruff check and ruff format --check pass on all files in this PR
```
</details>
<details>
<summary>Test logs (Kaggle Tesla T4, torch 2.10.0+cu128, triton
3.6.0)</summary>
```
test/transformers/test_fused_linear_kl_div.py: 51 passed, 25 skipped in 10.29s
(bf16 cases skipped by design: T4 is sm75 without native bf16 support)
```
</details>