linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(ppo): skip dW/dB calculation when lm_head parameters are frozen (#1488)

Browse Source 
## 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]>
This commit is contained in:
Daniel Lenton
2026-09-24 21:36:46 +00:00
committed by GitHub
co-authored by Jari Kolehmainen
parent fbdea0ff3b
commit ec10acbbe6
2 changed files with 126 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/chunked_loss/fused_linear_ppo.py
+44 -21
View File
@@ -68,10 +68,22 @@ def _selective_logprob_forward(
return logprobs, log_z
def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logprobs, temperature, vocab_chunk_size):
def _selective_logprob_backward(
hidden,
weight,
targets,
bias,
log_z,
grad_logprobs,
temperature,
vocab_chunk_size,
compute_grad_weight=True,
compute_grad_bias=True,
):
"""Dual-chunked (sequence × vocab) backward for selective logprob.
Recomputes logits per chunk for memory efficiency.
Recomputes logits per chunk for memory efficiency. The weight and bias gradients (a buffer the size of
the weight and one GEMM per chunk) are only computed when requested, e.g. not for a frozen lm_head.
"""
inv_t = 1.0 / temperature
n_rows, _ = hidden.shape
@@ -80,8 +92,9 @@ def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logpr
seq_chunk_size = _SELECTIVE_LOGPROB_SEQ_CHUNK_SIZE
grad_hidden = torch.zeros(hidden.shape, device=hidden.device, dtype=torch.float32)
grad_weight = torch.zeros(weight.shape, device=weight.device, dtype=torch.float32)
grad_bias = torch.zeros((vocab_size,), device=weight.device, dtype=torch.float32) if has_bias else None
compute_grad_bias = compute_grad_bias and has_bias
grad_weight = torch.zeros(weight.shape, device=weight.device, dtype=torch.float32) if compute_grad_weight else None
grad_bias = torch.zeros((vocab_size,), device=weight.device, dtype=torch.float32) if compute_grad_bias else None
grad_logprobs = grad_logprobs.to(torch.float32)
@@ -110,8 +123,9 @@ def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logpr
grad_logits.mul_(inv_t)
grad_hidden[seq_start:seq_end].add_(grad_logits @ weight_chunk.float())
grad_weight[vocab_start:vocab_end].add_(grad_logits.t() @ hidden_chunk.float())
if has_bias:
if compute_grad_weight:
grad_weight[vocab_start:vocab_end].add_(grad_logits.t() @ hidden_chunk.float())
if compute_grad_bias:
grad_bias[vocab_start:vocab_end].add_(grad_logits.sum(dim=0))
return grad_hidden, grad_weight, grad_bias
@@ -147,12 +161,14 @@ class _ChunkedSelectiveLogProbFunction(torch.autograd.Function):
grad_logprobs=grad_logprobs,
temperature=ctx.temperature,
vocab_chunk_size=ctx.vocab_chunk_size,
compute_grad_weight=ctx.needs_input_grad[1],
compute_grad_bias=ctx.has_bias and ctx.needs_input_grad[3],
)
return (
grad_hidden.to(hidden.dtype),
grad_weight.to(weight.dtype),
grad_weight.to(weight.dtype) if grad_weight is not None else None,
None,
grad_bias.to(bias.dtype) if ctx.has_bias else None,
grad_bias.to(bias.dtype) if grad_bias is not None else None,
None,
None,
)
@@ -231,9 +247,12 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
vllm_is_ratio = vllm_is_ratio.unsqueeze(-1) # (B,) -> (B, 1) for broadcasting
# Initialize accumulators
loss_acc = torch.zeros((), device=_input.device, dtype=torch.float32)
grad_weight = torch.zeros_like(weight) # [V, H]
# Determine which parameters require gradients (a frozen lm_head under LoRA/PEFT does not)
weight_requires_grad = weight.requires_grad
bias_requires_grad = bias is not None and bias.requires_grad
grad_weight = torch.zeros_like(weight) if weight_requires_grad else None # [V, H]
grad_inputs = []
grad_bias = torch.zeros_like(bias) if bias is not None else None # [V]
grad_bias = torch.zeros_like(bias) if bias_requires_grad else None # [V]
aggregated_metrics = []
# Only compile the loss math, NOT chunk_forward (which uses custom autograd.Function)
@@ -272,8 +291,8 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
"""Fused forward and backward for a chunk."""
with torch.enable_grad():
input_chunk = input_chunk.detach().requires_grad_(True)
weight_local = weight.detach().requires_grad_(True)
bias_local = bias.detach().requires_grad_(True) if bias is not None else None
weight_local = weight.detach().requires_grad_(weight_requires_grad)
bias_local = bias.detach().requires_grad_(bias_requires_grad) if bias is not None else None
# Step 1: compute logprobs OUTSIDE compile (custom autograd, memory-efficient)
per_token_logps = LigerFusedLinearPPOBase.chunk_forward(
@@ -305,8 +324,10 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
vllm_is_ratio_chunk=vllm_is_ratio_chunk,
)
grad_targets = [input_chunk, weight_local]
if bias_local is not None:
grad_targets = [input_chunk]
if weight_requires_grad:
grad_targets.append(weight_local)
if bias_requires_grad:
grad_targets.append(bias_local)
grads = torch.autograd.grad(chunk_loss, grad_targets)
return grads, (chunk_loss.detach(), tuple(metric.detach() for metric in chunk_metrics))
@@ -321,7 +342,7 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
ref_input_chunk=None,
vllm_is_ratio_chunk=None,
):
(chunk_grad_input, chunk_grad_weight, *chunk_grad_bias), (chunk_loss, chunk_metrics) = fused_fwd_bwd(
(chunk_grad_input, *chunk_grad_params), (chunk_loss, chunk_metrics) = fused_fwd_bwd(
input_chunk,
selected_token_ids_chunk,
attention_mask_chunk,
@@ -331,11 +352,12 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
ref_input_chunk,
vllm_is_ratio_chunk,
)
if bias is not None:
grad_bias.add_(chunk_grad_bias[0])
# Accumulate gradients and loss
grad_weight.add_(chunk_grad_weight)
# Accumulate gradients and loss (weight and bias gradients follow the input's, in that order,
# for whichever of them require grad)
if weight_requires_grad:
grad_weight.add_(chunk_grad_params.pop(0))
if bias_requires_grad:
grad_bias.add_(chunk_grad_params.pop(0))
grad_inputs.append(chunk_grad_input)
loss_acc.add_(chunk_loss)
# Initialize storage for metrics on first chunk
@@ -551,7 +573,8 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
if grad_output != 1.0:
grad_input = grad_input * grad_output
grad_weight = grad_weight * grad_output
if grad_weight is not None:
grad_weight = grad_weight * grad_output
if grad_bias is not None:
grad_bias = grad_bias * grad_output
test/chunked_loss/test_grpo_loss.py
+82
View File
@@ -1240,3 +1240,85 @@ def _reference_per_token_loss(
def _masked_mean(values, mask):
mask = mask.to(values.dtype)
return (values * mask).sum() / mask.sum().clamp(min=1.0)
@pytest.mark.parametrize("compiled", [False, True])
@pytest.mark.parametrize("bias", [False, True])
@pytest.mark.parametrize("use_ref_model", [False, True])
@pytest.mark.parametrize("V", [64, 4096 + 37]) # the second exceeds the vocab chunk of the selective logprob
def test_grpo_frozen_lm_head(compiled, bias, use_ref_model, V):
"""A frozen lm_head (LoRA/PEFT): no weight or bias gradient is computed or returned, and the loss and
input gradient are unchanged."""
torch.compiler.reset() # as in test_correctness_large_seq_exercises_chunking: start from a clean compile state
B, T, H = 2, 16, 32
torch.manual_seed(42)
_input = torch.randn(B, T, H, device=device, dtype=torch.float32)
weight = torch.randn(V, H, device=device, dtype=torch.float32)
lin_bias = torch.randn(V, device=device, dtype=torch.float32) if bias else None
selected_token_ids = torch.randint(0, V, (B, T), device=device)
attention_mask = torch.ones(B, T, device=device)
attention_mask[:, -3:] = 0
advantages = torch.randn(B, device=device, dtype=torch.float32)
ref_kwargs = {}
if use_ref_model:
ref_kwargs = dict(
ref_input=_input + 0.01 * torch.randn_like(_input),
ref_weight=weight.clone(),
ref_bias=lin_bias.clone() if bias else None,
)
def run(weight_requires_grad, bias_requires_grad):
x = _input.clone().requires_grad_(True)
w = weight.clone().requires_grad_(weight_requires_grad)
b = lin_bias.clone().requires_grad_(bias_requires_grad) if bias else None
loss_fn = LigerFusedLinearGRPOLoss(
beta=0.04 if use_ref_model else 0.0, use_ref_model=use_ref_model, compiled=compiled, chunk_size=1
)
loss, _ = loss_fn(x, w, selected_token_ids, attention_mask, advantages, bias=b, **ref_kwargs)
loss.backward()
return loss, x, w, b
loss_full, x_full, w_full, b_full = run(True, True)
loss_lora, x_lora, w_lora, b_lora = run(False, False)
assert w_full.grad is not None
assert w_lora.grad is None
if bias:
assert b_full.grad is not None
assert b_lora.grad is None
assert_verbose_allclose(loss_full, loss_lora, atol=1e-5, rtol=1e-5)
assert_verbose_allclose(x_full.grad, x_lora.grad, atol=1e-5, rtol=1e-5)
if bias:
# trainable weight, frozen bias: the weight gradient is unchanged and the bias gets none
_, x_mixed, w_mixed, b_mixed = run(True, False)
assert b_mixed.grad is None
assert_verbose_allclose(w_full.grad, w_mixed.grad, atol=1e-5, rtol=1e-5)
assert_verbose_allclose(x_full.grad, x_mixed.grad, atol=1e-5, rtol=1e-5)
@pytest.mark.parametrize("bias", [False, True])
def test_selective_logprob_backward_skips_frozen_params(monkeypatch, bias):
"""The chunked selective-logprob backward is asked for weight/bias gradients only when they require grad,
so a frozen lm_head costs neither the gradient buffer nor the per-chunk GEMM."""
import liger_kernel.chunked_loss.fused_linear_ppo as ppo
calls = []
real_backward = ppo._selective_logprob_backward
def spy(*args, **kwargs):
calls.append((kwargs["compute_grad_weight"], kwargs["compute_grad_bias"]))
return real_backward(*args, **kwargs)
monkeypatch.setattr(ppo, "_selective_logprob_backward", spy)
torch.manual_seed(0)
N, H, V = 10, 16, 50
targets = torch.randint(0, V, (N,), device=device)
hidden = torch.randn(N, H, device=device, requires_grad=True)
for weight_requires_grad, bias_requires_grad in ((False, False), (True, False), (True, True)):
weight = torch.randn(V, H, device=device, requires_grad=weight_requires_grad)
b = torch.randn(V, device=device, requires_grad=bias_requires_grad) if bias else None
out = ppo._ChunkedSelectiveLogProbFunction.apply(hidden, weight, targets, b, 1.0, 16)
out.sum().backward()
assert calls[-1] == (weight_requires_grad, bias and bias_requires_grad)
assert (weight.grad is not None) == weight_requires_grad