Repository navigation
[BUG] GRPOLoss kl_to_ref penalty is averaged over the padded tensor, so it scales with the fraction of valid tokens #4490
Description
Describe the bug
GRPOLoss._kl_to_ref (torchrl/objectives/llm/grpo.py) zeroes the masked positions of ref_log_prob and cur_log_prob and then takes (diff.expm1() - diff).mean() over the whole padded [B, T] tensor. The k3 estimate is therefore sum over the valid tokens / (B * T) instead of a mean over the valid tokens, so loss_kl_to_ref (and loss_kl_to_inference, which goes through the same method) is scaled by the fraction of valid tokens in the batch: the effective coefficient changes from batch to batch with the amount of padding, and the same responses padded to a longer T get a smaller penalty. The aggregation setting ("token_mean" / "prompt_mean") is not applied to the KL terms either, since they are added to td_out after the aggregation loop in forward.
The other loss terms go through _aggregate_loss_value, which divides by mask.sum() (token_mean) or by the per-sample token count (prompt_mean), and SFTLoss._kl_to_ref concatenates the unpadded log-probs, so it is a mean over the valid tokens as well. The logged kl_approx is a mean over the padded tensor too (logging only).
To Reproduce
Two responses of 2 and 3 valid tokens, left padded to T = 3, then the same tokens padded to T = 7. _FixedLogProbPolicy is the fixed-log-prob policy of test/llm/test_llm_objectives.py.
import torch
from tensordict import TensorDict
from torchrl.objectives.llm.grpo import GRPOLoss
from test_llm_objectives import _FixedLogProbPolicy # test/llm
cur = torch.tensor([[-1.0, -2.0, -0.5], [-0.2, -1.5, -3.0]])
ref = torch.tensor([[-1.5, -1.0, -0.5], [-0.4, -1.0, -2.0]])
mask = torch.tensor([[False, True, True], [True, True, True]])
def data(cur, ref, mask):
batch, seq = cur.shape
return TensorDict(
{
"current_log_prob": cur,
"mask": mask,
("tokens", "full"): torch.zeros(batch, seq, dtype=torch.long),
("log_probs", "full"): cur.clone(),
"advantage": torch.ones(batch, seq, 1),
("next", "ref_log_probs", "full"): ref,
},
batch_size=[batch],
)
loss_fn = GRPOLoss(_FixedLogProbPolicy(), clip_epsilon=0.2, entropy_bonus=False, kl_to_ref_coeff=0.5)
out = loss_fn(data(cur, ref, mask))
pad, pad_mask = torch.zeros(2, 4), torch.zeros(2, 4, dtype=torch.bool)
out_padded = loss_fn(
data(torch.cat([pad, cur], -1), torch.cat([pad, ref], -1), torch.cat([pad_mask, mask], -1))
)
diff = ref - cur
kl_token = diff.expm1() - diff
print("kl_to_ref, T=3:", out.kl_to_ref.item())
print("kl_to_ref, T=7:", out_padded.kl_to_ref.item())
print("mean over the 5 valid tokens:", kl_token[mask].mean().item())
print("sum over the valid tokens / 6:", (kl_token[mask].sum() / 6).item(), " / 14:", (kl_token[mask].sum() / 14).item())kl_to_ref, T=3: 0.2673359513282776
kl_to_ref, T=7: 0.11457256227731705
mean over the 5 valid tokens: 0.320803165435791
sum over the valid tokens / 6: 0.2673359513282776 / 14: 0.11457256227731705
Expected behavior
kl_to_ref is 0.3208 for both paddings (token_mean), or the mean of the per-sample means (prompt_mean), and loss_kl_to_ref is kl_to_ref_coeff times that, as for loss_objective and loss_entropy.
Screenshots
N/A
System info
- torchrl main at b8b6f0c (the same code is in 0.14.0), tensordict 0.14.2
- torch 2.14.0+cpu, numpy 2.2.6, Python 3.13, Linux
Additional context
kl_to_inference_coeff is affected in the same way. clip_fraction is also a mean over the padded tensor (padded positions have log_weight = 0, so they are never clipped); it is logging only and it is computed in _compute_policy_objective, which CISPOLoss overrides, so it is left alone here.
Reason and Possible fixes
Compute the per-token k3 term and aggregate it with _aggregate_loss_value when a mask is given, so the KL terms follow the same aggregation and reduction as the other loss terms (a mean over the valid tokens by default, per-token with reduction="none"), and take the logged kl_approx over the valid tokens. Patch and test below (test_kl_to_ref_is_aggregated_over_valid_tokens, parametrized over token_mean and prompt_mean; it fails on main with Expected 0.3208 but got 0.2673 and passes with the patch; test/llm/test_llm_objectives.py 75 passed, 5 skipped for vllm and CUDA; ufmt and flake8 clean). I can open the PR if the fix looks right to you.
git diff against main (2 files, +73 -2)
diff --git a/test/llm/test_llm_objectives.py b/test/llm/test_llm_objectives.py
index cd20858..f8ce3f9 100644
--- a/test/llm/test_llm_objectives.py
+++ b/test/llm/test_llm_objectives.py
@@ -1567,6 +1567,67 @@ class TestGRPOLossRefactorBehavior:
torch.testing.assert_close(out.kl_to_ref, expected_kl)
torch.testing.assert_close(out.loss_kl_to_ref, kl_coeff * expected_kl)
+ @pytest.mark.parametrize("aggregation", ["token_mean", "prompt_mean"])
+ def test_kl_to_ref_is_aggregated_over_valid_tokens(self, aggregation):
+ """The KL penalty follows the loss aggregation over the valid tokens.
+
+ Before the fix the k3 estimate was averaged over the padded [B, T] tensor
+ with the masked positions set to 0, so the penalty (and the logged
+ kl_approx) was scaled by the fraction of valid tokens in the batch and
+ changed with the amount of padding.
+ """
+ cur_lp = torch.tensor([[-1.0, -2.0, -0.5], [-0.2, -1.5, -3.0]])
+ ref_lp = torch.tensor([[-1.5, -1.0, -0.5], [-0.4, -1.0, -2.0]])
+ mask = torch.tensor([[False, True, True], [True, True, True]])
+ diff = ref_lp - cur_lp
+ kl_token = diff.expm1() - diff
+ if aggregation == "token_mean":
+ expected_kl = kl_token[mask].mean()
+ else:
+ expected_kl = (
+ (kl_token * mask).sum(-1, keepdim=True) / mask.sum(-1, keepdim=True)
+ ).mean(0)
+
+ def make_data(cur_lp, ref_lp, mask):
+ batch, seq = cur_lp.shape
+ return TensorDict(
+ {
+ "current_log_prob": cur_lp,
+ "mask": mask,
+ ("tokens", "full"): torch.zeros(batch, seq, dtype=torch.long),
+ ("log_probs", "full"): cur_lp + 0.1,
+ "advantage": torch.ones(batch, seq, 1),
+ ("next", "ref_log_probs", "full"): ref_lp,
+ },
+ batch_size=[batch],
+ )
+
+ kl_coeff = 0.5
+ loss_fn = GRPOLoss(
+ _FixedLogProbPolicy(),
+ clip_epsilon=0.2,
+ entropy_bonus=False,
+ kl_to_ref_coeff=kl_coeff,
+ aggregation=aggregation,
+ )
+ out = loss_fn(make_data(cur_lp, ref_lp, mask))
+ torch.testing.assert_close(out.kl_to_ref, expected_kl)
+ torch.testing.assert_close(out.loss_kl_to_ref, kl_coeff * expected_kl)
+ torch.testing.assert_close(out.kl_approx, torch.tensor(0.1))
+
+ # The same tokens with more left padding give the same penalty.
+ pad = torch.zeros(2, 4)
+ pad_mask = torch.zeros(2, 4, dtype=torch.bool)
+ out_padded = loss_fn(
+ make_data(
+ torch.cat([pad, cur_lp], -1),
+ torch.cat([pad, ref_lp], -1),
+ torch.cat([pad_mask, mask], -1),
+ )
+ )
+ torch.testing.assert_close(out_padded.kl_to_ref, expected_kl)
+ torch.testing.assert_close(out_padded.kl_approx, torch.tensor(0.1))
+
def test_dapo_can_be_instantiated_with_actor_network(self):
"""Before the fix, DAPO.__init__ took tensordict as its first positional
argument (a copy-paste of _kl_to_ref body). Calling
diff --git a/torchrl/objectives/llm/grpo.py b/torchrl/objectives/llm/grpo.py
index 2f5decc..e4cda7f 100644
--- a/torchrl/objectives/llm/grpo.py
+++ b/torchrl/objectives/llm/grpo.py
@@ -723,7 +723,9 @@ class GRPOLoss(LossModule):
)
td_out = TensorDict({"loss_objective": loss_objective})
td_out.set("clip_fraction", clip_fraction)
- td_out.set("kl_approx", kl_approx.detach().mean()) # for logging
+ td_out.set(
+ "kl_approx", kl_approx.detach()[expand_as_right(mask, kl_approx)].mean()
+ ) # for logging
if self.entropy_bonus:
entropy = self._get_entropy(dist, adv_shape=advantage.shape[:-1])
@@ -912,7 +914,15 @@ class GRPOLoss(LossModule):
expand_as_right(mask, cur_log_prob), cur_log_prob, 0.0
)
diff = ref_log_prob - cur_log_prob
- kl_penalty = (diff.expm1() - diff).mean()
+ kl_token = diff.expm1() - diff
+ if mask is None:
+ kl_penalty = kl_token.mean()
+ else:
+ # Aggregate over the valid tokens like the other loss terms. A mean over the padded
+ # tensor would scale the penalty by the fraction of valid tokens in the batch.
+ kl_penalty = self._aggregate_loss_value(
+ kl_token.unsqueeze(-1), mask, tensordict=tensordict
+ )
return coeff * kl_penalty, kl_penalty
def _log_weight(Checklist
- I have checked that there is no similar issue in the repo (required)
- I have read the documentation (required)
- I have provided a minimal working example to reproduce the bug (required)
Activity