mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
fc330e197e
## Summary `LigerFusedLinearDistillationBase` computes the student's float32 log-softmax over the vocabulary, and the NLL (the hard loss), on every chunk even when `weight_hard_loss == 0`. In that case the hard loss is multiplied by zero and discarded. This happens in pure distillation, for example `LigerFusedLinearJSDLoss(weight_hard_loss=0.0, weight_soft_loss=1.0)`. This PR computes the hard loss, and the log-softmax it needs, only when `weight_hard_loss != 0` or when `return_soft_hard_loss=True` asks for it. Loss and gradients are unchanged: the existing correctness tests already cover `weight_hard_loss=0` against the torch reference. The hard loss is still reported whenever it is requested. ## Details Forward + backward in bf16, `LigerFusedLinearJSDLoss(weight_hard_loss=0.0, weight_soft_loss=1.0, beta=0.5)`, trainable student head, B·T = 4096, H = 4096. Each figure is the median of 10 timed iterations after 3 warm-ups, then the median over separate processes, with main and this branch run alternately. | V | compiled | main | this PR | |---|---|---|---| | 128,256 | yes (default) | 42.0 ms, 6.51 GiB | 41.9 ms, 6.51 GiB | | 151,936 | yes (default) | 71.1 ms, 7.98 GiB | **66.1 ms, 7.69 GiB (−7%)** | | 128,256 | no | 85.8 ms, 8.75 GiB | **75.5 ms, 8.26 GiB (−12%)** | | 151,936 | no | 101.1 ms, 10.34 GiB | **89.0 ms, 9.75 GiB (−12%)** | This touches the same function as #1487 (frozen-parameter gradients in the same base). The two are independent, but whichever lands second needs a small rebase. Under `torch.compile`, the compiler already absorbs some of the redundant work, so the saving depends on the shape. With `weight_hard_loss=0.5` the timings are unchanged (42.0 vs 42.1 ms at V = 128,256; 71.0 vs 71.1 ms at V = 151,936). ## Testing Done - Hardware Type: NVIDIA H100 80GB HBM3 (torch 2.14.0+cu130, triton 3.8.0) - [ ] run `make test` to ensure correctness: I ran the test files for the changed base rather than the full suite. `test/chunked_loss/test_jsd_loss.py` and `test/chunked_loss/test_cosine_loss.py`: 110 passed on this branch (104 on main, plus 6 new). They also pass on CPU. - [x] run `make checkstyle` to ensure code style (ruff 0.15.22 `check` and `format --check` on the changed files) - [ ] run `make test-convergence` to ensure convergence: not run; the convergence tests do not exercise the chunked losses. New tests, for both losses: - The CE is computed exactly when it is weighted or returned: skipped at weight 0 without a return, and computed at weight 0 with a return and at weight 0.5. - With `weight_hard_loss=0` and `return_soft_hard_loss=True`, the returned hard loss equals the cross-entropy. Co-authored-by: Jari Kolehmainen <[email protected]> Co-authored-by: Arup De <[email protected]>