policy_loss is logged at two different scales depending on loss_type

[behroozazarkhalili]

Summary

policy_loss is logged at two different scales depending on loss_type. The grpo, bnpo, dr_grpo, sapo and luspo branches capture it before the gradient accumulation rescale, so the logged number is the value for the whole optimizer window. The cispo, dapo and vespo branch folds the normalizer into the loss first and captures afterwards, so the logged number is one micro-batch's contribution. The two differ by a factor of current_gradient_accumulation_steps / steps_per_generation.

Comparing policy_loss against loss, or across two loss types, therefore gives a wrong answer, and nothing in the code or the docs says so.

Where

trl/trainer/grpo_trainer.py, in _compute_loss:

branch	branch line	capture	rescale	scale reported
grpo, sapo	3242	3245	3246	window
bnpo	3247	3250	3251	window
dr_grpo	3252	3255	3256	window
cispo, dapo, vespo	3257	3263	3262, before the capture	per micro-batch
luspo	3264	3270	3271	window

Four branches capture before dividing by the normalizer, one divides first. The single append at 3302 then logs whichever quantity the branch produced.

Reproduction

Two micro-batches through _compute_loss with current_gradient_accumulation_steps = 2, comparing the appended metric against the sum of the returned losses:

import torch
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer

for loss_type in ("grpo", "bnpo", "dr_grpo", "sapo", "luspo", "dapo", "cispo", "vespo"):
    args = GRPOConfig(
        output_dir=f"/tmp/scale-{loss_type}", per_device_train_batch_size=3, num_generations=3,
        max_completion_length=2, gradient_accumulation_steps=2, loss_type=loss_type,
        use_liger_kernel=False, use_cpu=True, bf16=False, report_to="none",
    )
    trainer = GRPOTrainer(
        model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
        reward_funcs=lambda completions, **kwargs: [0.0] * len(completions),
        args=args, train_dataset=Dataset.from_dict({"prompt": ["p"] * 6}),
    )
    trainer.current_gradient_accumulation_steps = 2
    trainer.model.train()
    inputs = {
        "prompt_ids": torch.tensor([[1]] * 3), "prompt_mask": torch.ones(3, 1, dtype=torch.long),
        "completion_ids": torch.tensor([[2, 3]] * 3),
        "completion_mask": torch.tensor([[1, 0], [1, 1], [1, 1]]),
        "advantages": torch.tensor([-1.0, 0.0, 1.0]),
        "num_items_in_batch": torch.tensor(10.0),
    }
    returned = [trainer._compute_loss(trainer.model, inputs).item() for _ in range(2)]
    appended = list(trainer._metrics["train"]["policy_loss"])[0]
    window = sum(returned)
    print(loss_type, "window_sum", round(window, 6), "appended", round(appended, 6),
          "ratio", "n/a" if window == 0 else round(appended / window, 4))

Measured output:

loss_type	window sum	appended	ratio	scale
bnpo	-0.200000	-0.200000	1.0000	window
dr_grpo	-0.166667	-0.166667	1.0000	window
sapo	-0.031746	-0.031746	1.0000	window
luspo	-0.333333	-0.333333	1.0000	window
dapo	-0.200000	-0.100000	0.5000	per micro-batch
cispo	0.477227	0.238614	0.5000	per micro-batch
vespo	2.386140	1.193070	0.5000	per micro-batch
grpo	0.000000	0.000000	n/a	shares the sapo branch

Every evaluable ratio is exactly 1 or exactly 1/gradient_accumulation_steps, so the split is structural rather than numerical noise. grpo cannot be measured this way because its loss here is exactly zero: with old_per_token_logps unset the ratio is exactly 1, and GRPO advantages are centred within each group, so a micro-batch holding one whole group sums to zero. It takes the same branch as sapo.

Scope

The same block is duplicated in trl/experimental/gmpo/gmpo_trainer.py and trl/experimental/gspo_token/grpo_trainer.py, and trl/experimental/grpo_with_replay_buffer inherits it. Whatever is decided should apply to all four together.

Why this is filed rather than fixed

AGENTS.md asks that a defect found in duplicated logic be reported for a dedicated sweep rather than corrected inside a PR that happens to touch a copy. This came up while adding the Liger policy_loss metric in #6861: that metric mirrors the standard path exactly, including this split, which is what the consistency rule asks for. Picking one scale is a behaviour change for seven loss types across four trainers and belongs in its own change.

Possible directions

Report one scale everywhere, most naturally the window value, so policy_loss is comparable with loss and across loss types. Failing that, document the split in docs/source/grpo_trainer.md so the difference is at least visible.