Exclude padding from RL trainer diagnostics, and fix three related metrics

docs/source/grpo_trainer.md

@@ -192,7 +192,7 @@ While training and evaluating, we record the following metrics:
 - `sampling/importance_sampling_ratio/min`: The smallest importance sampling ratio used to correct the train-inference mismatch, **after** the constraint selected by `vllm_importance_sampling_mode` has been applied (clipped to `[C_min, C_max]` for `*_truncate` modes, set to zero for `*_mask` modes). Computed over completion tokens, or over sequences for the `sequence_*` modes. Logged only when `use_vllm=True` and `vllm_importance_sampling_correction=True`.
 - `sampling/importance_sampling_ratio/mean`: The average constrained importance sampling ratio. Logged only when `use_vllm=True` and `vllm_importance_sampling_correction=True`.
 - `sampling/importance_sampling_ratio/max`: The largest constrained importance sampling ratio. Logged only when `use_vllm=True` and `vllm_importance_sampling_correction=True`.
-- `policy_loss`: The policy gradient loss value (before any entropy bonus). Logged when `entropy_coef` is nonzero or `use_adaptive_entropy=True`.
+- `policy_loss`: The policy gradient loss value, logged on every run. On the standard path it excludes the entropy bonus and the Mixture-of-Experts auxiliary loss, both of which are added afterwards. The Liger path adds neither: an entropy bonus raises `NotImplementedError` at construction, and the auxiliary loss is never computed there. Either way it does include the KL term when `beta` is nonzero, because that is folded in first. The scale differs by loss type: `grpo`, `bnpo`, `dr_grpo`, `sapo` and `luspo` report the value before the gradient accumulation rescale, while `cispo`, `dapo` and `vespo` report the per-micro-batch contribution, the value after division by `current_gradient_accumulation_steps / steps_per_generation`, which is smaller only when that ratio exceeds one.
 - `entropy`: Average entropy of token predictions across generated completions. (If `mask_truncated_completions=True`, masked sequences tokens are excluded.)
 - `entropy_coef`: The current entropy regularization coefficient. Logged when `entropy_coef` is nonzero or `use_adaptive_entropy=True`. Updated once per optimizer step when `use_adaptive_entropy=True`.
 - `aux_loss`: The load-balancing auxiliary loss of a Mixture-of-Experts model, before it is scaled by `router_aux_loss_coef` and added to the loss. Logged only when the model is a MoE model and `router_aux_loss_coef` is nonzero.

tests/experimental/test_online_dpo_trainer.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pytest
+import torch
 from datasets import Dataset, DatasetDict, features, load_dataset
 from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
 from transformers.utils import is_peft_available, is_vision_available
@@ -93,6 +94,53 @@ def test_train(self, config_name):
 
         assert "train_loss" in trainer.state.log_history[-1]
 
+    def test_completion_metrics_ignore_padding(self, monkeypatch):
+        dataset = Dataset.from_dict({"prompt": ["one", "two"]})
+        training_args = OnlineDPOConfig(
+            output_dir=self.tmp_dir,
+            per_device_train_batch_size=2,
+            max_steps=1,
+            beta=0.1,
+            report_to="none",
+        )
+        trainer = OnlineDPOTrainer(
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_funcs=self.reward_model,
+            args=training_args,
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+            reward_processing_classes=self.reward_tokenizer,
+        )
+        device = trainer.accelerator.device
+        prompt_ids = torch.ones((4, 1), dtype=torch.long, device=device)
+        prompt_mask = torch.ones_like(prompt_ids)
+        completion_ids = torch.ones((4, 3), dtype=torch.long, device=device)
+        completion_mask = torch.tensor([[1, 1, 0], [1, 0, 0], [1, 1, 1], [1, 0, 0]], dtype=torch.long, device=device)
+        logprobs = torch.tensor(
+            [[-1.0, -2.0, -100.0], [-3.0, -100.0, -100.0], [-4.0, -5.0, -6.0], [-7.0, -100.0, -100.0]],
+            device=device,
+            requires_grad=True,
+        )
+        ref_logprobs = torch.tensor(
+            [[-0.5, -1.0, -200.0], [-1.0, -200.0, -200.0], [-2.0, -2.0, -2.0], [-3.0, -200.0, -200.0]],
+            device=device,
+        )
+        rewards = torch.tensor([2.0, 1.0, 0.0, 3.0], device=device)
+        forward_calls = iter((logprobs, ref_logprobs))
+        monkeypatch.setattr(
+            trainer, "_generate", lambda *args: (prompt_ids, prompt_mask, completion_ids, completion_mask)
+        )
+        monkeypatch.setattr(trainer, "_forward", lambda *args: next(forward_calls))
+        monkeypatch.setattr(trainer, "_calculate_rewards_from_functions", lambda **kwargs: rewards)
+
+        trainer.training_step(trainer.model, {"prompt": ["one", "two"]})
+
+        assert trainer.stats["objective/kl"][-1] == pytest.approx(-4.125)
+        assert trainer.stats["objective/entropy"][-1] == pytest.approx(7.0)
+        assert trainer.stats["objective/non_score_reward"][-1] == pytest.approx(0.4125)
+        assert trainer.stats["objective/rlhf_reward"][-1] == pytest.approx(1.9125)
+
     @pytest.mark.parametrize("eval_dataset_type", ["dataset", "dataset_dict", "dict_of_dataset", "none"])
     def test_init_with_eval_dataset(self, eval_dataset_type):
         # Streaming datasets are not yet supported in OnlineDPO

tests/experimental/test_ppo_trainer.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import gc
+import math
 import os
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -33,6 +35,7 @@
     PPOConfig,
     PPOTrainer,
 )
+from trl.experimental.ppo import ppo_trainer as ppo_trainer_module
 from trl.experimental.ppo.ppo_trainer import batch_generation, masked_mean, masked_var, masked_whiten
 
 from ..testing_utils import (
@@ -713,6 +716,188 @@ def tokenize(example, tokenizer):
 
         self.raw_dataset = raw_dataset.map(tokenize, fn_kwargs={"tokenizer": self.tokenizer}, remove_columns="prompt")
 
+    def test_statistics_are_not_diluted_by_unwritten_slots(self):
+        """The statistics buffers used to be sized by `gradient_accumulation_steps`, while the micro-batch loop
+        writes `ceil(local_mini_batch_size / per_device_train_batch_size)` slots per minibatch and resets its index
+        each minibatch. The unwritten slots stayed at zero and the `.mean()` below averaged them in, scaling every
+        reported statistic by `ceil(gradient_accumulation_steps / num_mini_batches) / gradient_accumulation_steps`,
+        which works out to `1 / num_mini_batches` in this configuration. `ratio` pins that dilution: the micro-batch
+        body runs inside `accelerator.accumulate(model)`, which defers the parameter update to the sync micro-batch, so
+        every `ratio` recorded here is computed before any update has landed and is 1 up to floating-point noise. With
+        two minibatches it was reported as 0.5."""
+        training_args = PPOConfig(
+            output_dir=self.tmp_dir,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=2,
+            num_mini_batches=2,
+            num_ppo_epochs=1,
+            report_to="none",
+        )
+        trainer = PPOTrainer(
+            args=training_args,
+            processing_class=self.tokenizer,
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_model=self.reward_model,
+            value_model=self.value_model,
+            train_dataset=self.raw_dataset["train"],
+            eval_dataset=self.raw_dataset["test"],
+        )
+        trainer.train()
+
+        ratios = [log["val/ratio"] for log in trainer.state.log_history if "val/ratio" in log]
+        assert ratios, "no `val/ratio` was logged, so the assertion below would pass vacuously"
+        for ratio in ratios:
+            assert ratio == pytest.approx(1.0, abs=1e-4)
+
+    def test_statistics_ignore_an_all_padding_micro_batch(self):
+        """A micro-batch whose rows are all padding has no statistic to report. Its `masked_mean` is 0 by
+        construction, and writing that 0 into the buffers made the later `.mean()` count it as an observation, so
+        `val/ratio` read 0.5 when the one real micro-batch had a ratio of 1. Policy slots now stay NaN and the
+        `nanmean` reductions leave them out. Value slots use `padding_mask_p1` instead: a zero-length response retains
+        its first value timestep, so its value diagnostics must still be logged. The all-padding row is forced by
+        making generation emit the pad token first, so its sequence length comes out as -1 and every response position
+        is policy padding. The tokenizer's own pad token is used as is: padding with EOS instead makes `forward` drop
+        every EOS from the attention mask as well, and it moved `val/ratio` by up to 2e-3 even without the forced row;
+        the distinct pad token keeps the real micro-batch within 1e-7."""
+        tokenizer = AutoTokenizer.from_pretrained(self.model_id, padding_side="left")
+        pad_token_id = tokenizer.pad_token_id
+        real_batch_generation = ppo_trainer_module.batch_generation
+        real_forward = ppo_trainer_module.forward
+        context_length = None
+        expected_entropies = []
+
+        def batch_generation_with_one_empty_row(model, queries, local_rollout_forward_batch_size, pad_id, config):
+            nonlocal context_length
+            context_length = queries.shape[1]
+            query_responses, logitss = real_batch_generation(
+                model, queries, local_rollout_forward_batch_size, pad_id, config
+            )
+            query_responses[0, queries.shape[1]] = pad_token_id
+            return query_responses, logitss
+
+        def forward_with_distinct_empty_value(model, query_responses, pad_id):
+            output = real_forward(model, query_responses, pad_id)
+            if isinstance(model, ppo_trainer_module.PolicyAndValueWrapper):
+                policy_output, values = output
+                responses = query_responses[:, context_length:]
+                sequence_lengths = ppo_trainer_module.first_true_indices(responses == pad_token_id) - 1
+                response_idxs = torch.arange(responses.shape[1], device=responses.device).expand_as(responses)
+                policy_mask = response_idxs <= sequence_lengths.unsqueeze(1)
+                policy_logits = policy_output.logits.clone()
+                response_logits = policy_logits[:, context_length - 1 : -1]
+                response_logits[~policy_mask] = -1000
+                response_logits[..., 0] = torch.where(
+                    policy_mask, response_logits[..., 0], torch.zeros_like(response_logits[..., 0])
+                )
+                policy_output.logits = policy_logits
+                logits = policy_output.logits[:, context_length - 1 : -1] / (training_args.temperature + 1e-7)
+                prob_dist = torch.nn.functional.softmax(logits, dim=-1)
+                entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1)
+                if policy_mask.any():
+                    expected_entropies.append(masked_mean(entropy, policy_mask).item())
+
+                empty_rows = sequence_lengths == -1
+                values = values.clone()
+                values[empty_rows, context_length - 1 : -1] = values[empty_rows, context_length - 1 : -1] * 0 + 1000
+                output = policy_output, values
+            return output
+
+        training_args = PPOConfig(
+            output_dir=self.tmp_dir,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=2,
+            num_mini_batches=1,
+            num_ppo_epochs=1,
+            total_episodes=2,
+            report_to="none",
+        )
+        trainer = PPOTrainer(
+            args=training_args,
+            processing_class=tokenizer,
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_model=self.reward_model,
+            value_model=self.value_model,
+            train_dataset=self.raw_dataset["train"],
+            eval_dataset=self.raw_dataset["test"],
+        )
+        with (
+            patch.object(ppo_trainer_module, "batch_generation", batch_generation_with_one_empty_row),
+            patch.object(ppo_trainer_module, "forward", forward_with_distinct_empty_value),
+        ):
+            trainer.train()
+
+        logs = [log for log in trainer.state.log_history if "val/ratio" in log]
+        assert logs, "no PPO statistics were logged, so the assertions below would pass vacuously"
+        assert len(expected_entropies) == len(logs)
+        # The one real micro-batch reports a ratio of 1 (measured within 1.2e-7); counting the empty slot as a 0 halves
+        # it to 0.5.
+        for log, expected_entropy in zip(logs, expected_entropies, strict=True):
+            assert log["val/ratio"] == pytest.approx(1.0, abs=1e-4)
+            assert log["policy/entropy_avg"] == pytest.approx(expected_entropy)
+            assert log["loss/value_avg"] > 100_000
+
+    def test_statistics_do_not_carry_over_from_the_previous_update(self):
+        """The statistic buffers are NaN-initialised and a slot is written only for a micro-batch with a valid token.
+        Allocated once for the whole run, a slot skipped in one update kept the value the previous update wrote there,
+        and `nanmean` folded that stale number into the current update's averages. The buffers are now fresh for every
+        update. The oracle: the first update has two real micro-batches, the second has one all-padding and one real.
+        With fresh buffers the second update has a single written ratio slot, so `val/ratio_var`, the unbiased variance
+        over the written slots, is NaN; a stale first-update slot makes it a finite number. An update with no valid
+        token at all is not a reachable state, `masked_whiten` refuses it, so the padding is confined to one row."""
+        tokenizer = AutoTokenizer.from_pretrained(self.model_id, padding_side="left")
+        pad_token_id = tokenizer.pad_token_id
+        real_batch_generation = ppo_trainer_module.batch_generation
+        calls = []
+
+        def batch_generation_with_an_empty_row_in_the_second_update(
+            model, queries, local_rollout_forward_batch_size, pad_id, config
+        ):
+            query_responses, logitss = real_batch_generation(
+                model, queries, local_rollout_forward_batch_size, pad_id, config
+            )
+            calls.append(len(calls))
+            if (
+                len(calls) == 2
+            ):  # row 0 of the second update starts with the pad token: sequence length -1, all padding
+                query_responses[0, queries.shape[1]] = pad_token_id
+            return query_responses, logitss
+
+        training_args = PPOConfig(
+            output_dir=self.tmp_dir,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=2,
+            num_mini_batches=1,
+            num_ppo_epochs=1,
+            total_episodes=6,  # three updates of two episodes each
+            report_to="none",
+        )
+        trainer = PPOTrainer(
+            args=training_args,
+            processing_class=tokenizer,
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_model=self.reward_model,
+            value_model=self.value_model,
+            train_dataset=self.raw_dataset["train"],
+            eval_dataset=self.raw_dataset["test"],
+        )
+        with patch.object(
+            ppo_trainer_module, "batch_generation", batch_generation_with_an_empty_row_in_the_second_update
+        ):
+            trainer.train()
+
+        logs = [log for log in trainer.state.log_history if "val/ratio_var" in log]
+        assert len(logs) == 3, f"expected one PPO log per update, got {len(logs)}"
+        for log in logs:
+            assert log["val/ratio"] == pytest.approx(1.0, abs=1e-4)
+        assert not math.isnan(logs[0]["val/ratio_var"]), "two real micro-batches give a finite variance"
+        assert math.isnan(logs[1]["val/ratio_var"]), (
+            f"one written slot must give a NaN variance, got {logs[1]['val/ratio_var']}: a stale slot was counted"
+        )
+        assert not math.isnan(logs[2]["val/ratio_var"])
+
     def test_basic_training(self):
         """Test basic PPO training configuration and verify model updates."""
         # Capture initial weights

tests/experimental/test_xpo_trainer.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import defaultdict
+from types import SimpleNamespace
+
 import pytest
+import torch
 from datasets import DatasetDict, load_dataset
 from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
 from transformers.utils import is_peft_available
@@ -62,6 +66,40 @@ def test_xpo_trainer_training(self, config_name):
 
         assert "train_loss" in trainer.state.log_history[-1]
 
+    def test_logps_metrics_contain_only_policy_logps(self):
+        trainer = SimpleNamespace(
+            accelerator=SimpleNamespace(gather_for_metrics=lambda tensor: tensor),
+            stats=defaultdict(list),
+            beta=0.1,
+            alpha=0.2,
+            processing_class=SimpleNamespace(eos_token_id=99),
+        )
+        model_data = {"input_ids": torch.tensor([[0, 1, 99], [0, 1, 2]])}
+        ref_data = {"input_ids": torch.tensor([[0, 3, 4], [0, 5, 99]])}
+        model_logprobs_model_data = torch.tensor([[1.0, 2.0], [10.0, 20.0]])
+        model_logprobs_ref_data = torch.tensor([[4.0, 5.0], [40.0, 50.0]])
+        ref_logprobs_ref_data = torch.tensor([[0.4, 0.5], [4.0, 5.0]])
+        ref_logprobs_model_data = torch.tensor([[0.1, 0.2], [1.0, 2.0]])
+
+        XPOTrainer._log_statistics(
+            trainer,
+            model_data,
+            ref_data,
+            model_logprobs_model_data,
+            model_logprobs_ref_data,
+            ref_logprobs_ref_data,
+            ref_logprobs_model_data,
+            torch.tensor([True, False]),
+            torch.tensor([0.1, 0.2]),
+            torch.tensor([0.3, 0.4]),
+            1,
+            torch.tensor([1.0, 2.0]),
+            torch.tensor([0.5, 1.5]),
+        )
+
+        assert trainer.stats["logps/chosen"] == [pytest.approx(46.5)]
+        assert trainer.stats["logps/rejected"] == [pytest.approx(19.5)]
+
     @pytest.mark.parametrize("eval_dataset_type", ["dataset", "dataset_dict", "dict_of_dataset", "none"])
     def test_init_with_eval_dataset(self, eval_dataset_type):
         # Streaming datasets are not yet supported in XPO