mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Daniel LentonandJari Kolehmainen
ec10acbbe6
## Summary `LigerFusedLinearPPOBase`, which backs `LigerFusedLinearGRPOLoss` and its variants (GRPO, DAPO, Dr. GRPO, BNPO, CISPO, SAPO, ...), always makes the `lm_head` weight (and bias) differentiable (`weight.detach().requires_grad_(True)`). As a result, `_ChunkedSelectiveLogProbFunction.backward` allocates a float32 `[V, H]` gradient buffer and runs a `grad_logits.T @ hidden` GEMM on every sequence × vocab chunk, even when the weight does not require grad. With LoRA/PEFT the `lm_head` is frozen and that gradient is discarded. This PR differentiates the weight and bias only when they require grad, as #1441 did for the preference losses. `_ChunkedSelectiveLogProbFunction` now computes only the gradients autograd asks for (`ctx.needs_input_grad`). The loss and the input gradients are unchanged, and trainable parameters take the same path as before. ## Details Forward + backward in bf16, `LigerFusedLinearGRPOLoss(beta=0.0, use_ref_model=False)` with its other defaults (compiled), B = 8, T = 1024, 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 | lm_head | main | this PR | |---|---|---|---| | 128,256 | frozen | 471.6 ms, 6.14 GiB | **263.7 ms, 3.25 GiB (−44%)** | | 151,936 | frozen | 557.4 ms, 7.22 GiB | **311.7 ms, 3.79 GiB (−44%)** | | 128,256 | trainable | 472.2 ms | 471.6 ms | | 151,936 | trainable | 558.8 ms | 558.0 ms | #1294 also edits `fused_linear_ppo.py`, changing the GEMM operand dtypes in the same backward. The two changes are independent, but whichever lands second needs a rebase. ## 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 file for the changed base rather than the full suite. `test/chunked_loss/test_grpo_loss.py`: 994 passed and 614 skipped on this branch (976 passed and 614 skipped on main, plus 18 new). - [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: - A frozen vs a trainable `lm_head`, with and without bias, with and without a reference model, and at a vocab size above the 4096 vocab chunk, so both chunk loops run. - A trainable weight with a frozen bias. - A check that `_selective_logprob_backward` is asked for the weight and bias gradients only when they require grad. The frozen-head test calls `torch.compiler.reset()` first, as `test_correctness_large_seq_exercises_chunking` does. Without it, starting from a cold compile cache on CPU, two of its compiled cases hit a dynamo guard error ("sources must not be empty for symbol"). The error reproduces identically on main with the same test, so it is not related to this change. Co-authored-by: Jari Kolehmainen <[email protected]>