test_ppo_trainer.py

# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gc
import math
import os
from unittest.mock import patch

import pytest
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSeq2SeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    GenerationConfig,
)
from transformers.utils import is_peft_available

from trl.experimental.ppo import (
    AutoModelForCausalLMWithValueHead,
    AutoModelForSeq2SeqLMWithValueHead,
    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 (
    TrlTestCase,
    require_bitsandbytes,
    require_peft,
    require_torch_gpu_if_bnb_not_multi_backend_enabled,
)


if is_peft_available():
    from peft import LoraConfig, get_peft_model


ALL_CAUSAL_LM_MODELS = [
    "trl-internal-testing/tiny-BloomForCausalLM",
    "trl-internal-testing/tiny-CohereForCausalLM",
    # "trl-internal-testing/tiny-FalconMambaForCausalLM",  # FalconMambaForCausalLM modeling seems to be broken for now
    "trl-internal-testing/tiny-Gemma2ForCausalLM",
    "trl-internal-testing/tiny-GemmaForCausalLM",
    "trl-internal-testing/tiny-GPT2LMHeadModel",
    "trl-internal-testing/tiny-GPTNeoXForCausalLM",
    "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
    "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
    "trl-internal-testing/tiny-LlamaForCausalLM-3",
    "trl-internal-testing/tiny-MistralForCausalLM-0.1",
    "trl-internal-testing/tiny-MistralForCausalLM-0.2",
    "trl-internal-testing/tiny-OPTForCausalLM",
    "trl-internal-testing/tiny-Phi3ForCausalLM-3",
    "trl-internal-testing/tiny-Phi3ForCausalLM-3.5",
    "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
]

ALL_SEQ2SEQ_MODELS = [
    "trl-internal-testing/tiny-T5ForConditionalGeneration",
    "trl-internal-testing/tiny-BartModel",
]


class TestBatchGeneration(TrlTestCase):
    def setup_method(self):
        # Initialize the tokenizer
        self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32").to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)

        self.generation_config = GenerationConfig(
            max_new_tokens=128,
            temperature=0.5,
            do_sample=True,
            top_k=0,
            pad_token_id=self.tokenizer.pad_token_id,
        )

        # Example input
        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
        self.examples = dataset["messages"]
        self.mini_batch_size = 3

    def test_mini_batch_generation(self):
        batch = [
            self.tokenizer.apply_chat_template(example[:-1], add_generation_prompt=True, tokenize=False)
            for example in self.examples
        ]
        queries = self.tokenizer(batch, padding=True, return_tensors="pt")["input_ids"].to(self.device)
        bs, context_length = queries.shape

        query_responses, logits = batch_generation(
            self.model, queries, self.mini_batch_size, self.tokenizer.pad_token_id, self.generation_config
        )

        max_length_query = query_responses.shape[1]
        max_length_logits = max_length_query - context_length

        assert max_length_query > context_length
        assert query_responses.shape == (bs, max_length_query)
        assert logits.shape == (bs, max_length_logits, self.model.config.vocab_size)

    def test_single_batch_generation(self):
        batch = [
            self.tokenizer.apply_chat_template(example[:-1], add_generation_prompt=True, tokenize=False)
            for example in self.examples
        ]
        queries = self.tokenizer(batch, padding=True, return_tensors="pt")["input_ids"].to(self.device)
        bs, context_length = queries.shape

        query_responses, logits = batch_generation(
            self.model, queries, bs, self.tokenizer.pad_token_id, self.generation_config
        )

        max_length_query = query_responses.shape[1]
        max_length_logits = max_length_query - context_length

        assert max_length_query > context_length
        assert query_responses.shape == (bs, max_length_query)
        assert logits.shape == (bs, max_length_logits, self.model.config.vocab_size)


class BaseTester:
    class VHeadModelTester(TrlTestCase):
        all_model_names = None
        trl_model_class = None
        transformers_model_class = None

        def setup_method(self):
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        def test_value_head(self):
            r"""
            Test if the v-head is added to the model successfully
            """
            for model_name in self.all_model_names:
                model = self.trl_model_class.from_pretrained(model_name)
                assert hasattr(model, "v_head")

        def test_value_head_shape(self):
            r"""
            Test if the v-head has the correct shape
            """
            for model_name in self.all_model_names:
                model = self.trl_model_class.from_pretrained(model_name)
                assert model.v_head.summary.weight.shape[0] == 1

        def test_value_head_init_random(self):
            r"""
            Test if the v-head has been randomly initialized. We can check that by making sure the bias is different
            than zeros by default.
            """
            for model_name in self.all_model_names:
                model = self.trl_model_class.from_pretrained(model_name)
                assert not torch.allclose(model.v_head.summary.bias, torch.zeros_like(model.v_head.summary.bias))

        def test_value_head_not_str(self):
            r"""
            Test if the v-head is added to the model successfully, by passing a non `PretrainedModel` as an argument to
            `from_pretrained`.
            """
            for model_name in self.all_model_names:
                pretrained_model = self.transformers_model_class.from_pretrained(model_name)
                model = self.trl_model_class.from_pretrained(pretrained_model)
                assert hasattr(model, "v_head")

        def test_from_save_trl(self):
            """
            Test if the model can be saved and loaded from a directory and get the same weights, including the
            additional modules (e.g. v_head)
            """
            for model_name in self.all_model_names:
                model = self.trl_model_class.from_pretrained(model_name)

                model.save_pretrained(self.tmp_dir)

                model_from_save = self.trl_model_class.from_pretrained(self.tmp_dir)

                # Check if the weights are the same
                for key in model_from_save.state_dict():
                    torch.testing.assert_close(model_from_save.state_dict()[key], model.state_dict()[key])

        def test_from_save_trl_sharded(self):
            """
            Test if the model can be saved and loaded from a directory and get the same weights - sharded case
            """
            for model_name in self.all_model_names:
                model = self.trl_model_class.from_pretrained(model_name)

                model.save_pretrained(self.tmp_dir)

                model_from_save = self.trl_model_class.from_pretrained(self.tmp_dir)

                # Check if the weights are the same
                for key in model_from_save.state_dict():
                    torch.testing.assert_close(model_from_save.state_dict()[key], model.state_dict()[key])

        def test_from_save_transformers_sharded(self):
            """
            Test if the model can be saved and loaded using transformers and get the same weights - sharded case
            """
            for model_name in self.all_model_names:
                transformers_model = self.trl_model_class.transformers_parent_class.from_pretrained(model_name)

                trl_model = self.trl_model_class.from_pretrained(model_name)

                trl_model.save_pretrained(self.tmp_dir, max_shard_size="1MB")
                transformers_model_from_save = self.trl_model_class.transformers_parent_class.from_pretrained(
                    self.tmp_dir
                )

                # Check if the weights are the same
                for key in transformers_model.state_dict():
                    torch.testing.assert_close(
                        transformers_model_from_save.state_dict()[key], transformers_model.state_dict()[key]
                    )

        def test_from_save_transformers(self):
            """
            Test if the model can be saved and loaded using transformers and get the same weights. We override the test
            of the super class to check if the weights are the same.
            """
            for model_name in self.all_model_names:
                transformers_model = self.trl_model_class.transformers_parent_class.from_pretrained(model_name)

                trl_model = self.trl_model_class.from_pretrained(model_name)

                trl_model.save_pretrained(self.tmp_dir)
                transformers_model_from_save = self.trl_model_class.transformers_parent_class.from_pretrained(
                    self.tmp_dir
                )

                # Check if the weights are the same
                for key in transformers_model.state_dict():
                    torch.testing.assert_close(
                        transformers_model_from_save.state_dict()[key], transformers_model.state_dict()[key]
                    )

                # Check if the trl model has the same keys as the transformers model
                # except the v_head
                for key in trl_model.state_dict():
                    if "v_head" not in key:
                        assert key in transformers_model.state_dict()
                        # check if the weights are the same
                        torch.testing.assert_close(trl_model.state_dict()[key], transformers_model.state_dict()[key])

                # check if they have the same modules
                assert set(transformers_model_from_save.state_dict().keys()) == set(
                    transformers_model.state_dict().keys()
                )


class TestCausalLMValueHeadModel(BaseTester.VHeadModelTester, TrlTestCase):
    """
    Testing suite for v-head models.
    """

    all_model_names = ALL_CAUSAL_LM_MODELS
    trl_model_class = AutoModelForCausalLMWithValueHead
    transformers_model_class = AutoModelForCausalLM

    def teardown_method(self):
        # free memory
        gc.collect()

    def test_inference(self):
        r"""
        Test if the model can be used for inference and outputs 3 values
        - logits, loss, and value states
        """
        EXPECTED_OUTPUT_SIZE = 3

        for model_name in self.all_model_names:
            model = self.trl_model_class.from_pretrained(model_name).to(self.device)
            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
            outputs = model(input_ids)

            # Check if the outputs are of the right size - here
            # we always output 3 values - logits, loss, and value states
            assert len(outputs) == EXPECTED_OUTPUT_SIZE

    def test_dropout_config(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            pretrained_model = self.transformers_model_class.from_pretrained(model_name)
            pretrained_model.config.summary_dropout_prob = 0.5
            model = self.trl_model_class.from_pretrained(pretrained_model)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == pretrained_model.config.summary_dropout_prob

    def test_dropout_kwargs(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            v_head_kwargs = {"summary_dropout_prob": 0.5}

            model = self.trl_model_class.from_pretrained(model_name, **v_head_kwargs)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

            model = self.trl_model_class.from_pretrained(model_name, summary_dropout_prob=0.5)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

    @pytest.mark.parametrize("model_name", ALL_CAUSAL_LM_MODELS)
    def test_generate(self, model_name):
        r"""
        Test if `generate` works for every model
        """
        generation_config = GenerationConfig(max_new_tokens=9)
        model = self.trl_model_class.from_pretrained(model_name).to(self.device)
        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)

        # Just check if the generation works
        _ = model.generate(input_ids, generation_config=generation_config)

    def test_transformers_bf16_kwargs(self):
        r"""
        Test if the transformers kwargs are correctly passed. Here we check that loading a model in half precision
        works as expected, i.e. the weights of the `pretrained_model` attribute is loaded in half precision and you can
        run a dummy forward pass without any issue.
        """
        for model_name in self.all_model_names:
            trl_model = self.trl_model_class.from_pretrained(model_name, dtype=torch.bfloat16).to(self.device)

            lm_head_namings = ["lm_head", "embed_out", "output_layer"]

            assert any(hasattr(trl_model.pretrained_model, lm_head_naming) for lm_head_naming in lm_head_namings), (
                "Can't test the model because it doesn't have any of the expected lm_head namings"
            )

            for lm_head_naming in lm_head_namings:
                if hasattr(trl_model.pretrained_model, lm_head_naming):
                    assert getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16

            dummy_input = torch.LongTensor([[0, 1, 0, 1]]).to(self.device)

            # check dummy forward pass works in half precision
            _ = trl_model(dummy_input)

    @pytest.mark.skip(reason="This test needs to be run manually due to HF token issue.")
    def test_push_to_hub(self):
        for model_name in self.all_model_names:
            model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)
            if "sharded" in model_name:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True, max_shard_size="1MB")
            else:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True)

            model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(model_name + "-ppo")
            # check all keys
            assert model.state_dict().keys() == model_from_pretrained.state_dict().keys()

            for name, param in model.state_dict().items():
                (
                    torch.testing.assert_close(param, model_from_pretrained.state_dict()[name]),
                    (f"Parameter {name} is not the same after push_to_hub and from_pretrained"),
                )


class TestSeq2SeqValueHeadModel(BaseTester.VHeadModelTester, TrlTestCase):
    """
    Testing suite for v-head models.
    """

    all_model_names = ALL_SEQ2SEQ_MODELS
    trl_model_class = AutoModelForSeq2SeqLMWithValueHead
    transformers_model_class = AutoModelForSeq2SeqLM

    def teardown_method(self):
        # free memory
        gc.collect()

    def test_inference(self):
        r"""
        Test if the model can be used for inference and outputs 3 values
        - logits, loss, and value states
        """
        EXPECTED_OUTPUT_SIZE = 3

        for model_name in self.all_model_names:
            model = self.trl_model_class.from_pretrained(model_name).to(self.device)
            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
            decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
            outputs = model(input_ids, decoder_input_ids=decoder_input_ids)

            # Check if the outputs are of the right size - here
            # we always output 3 values - logits, loss, and value states
            assert len(outputs) == EXPECTED_OUTPUT_SIZE

    def test_dropout_config(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            pretrained_model = self.transformers_model_class.from_pretrained(model_name)
            pretrained_model.config.summary_dropout_prob = 0.5
            model = self.trl_model_class.from_pretrained(pretrained_model)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == pretrained_model.config.summary_dropout_prob

    def test_dropout_kwargs(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            v_head_kwargs = {"summary_dropout_prob": 0.5}

            model = self.trl_model_class.from_pretrained(model_name, **v_head_kwargs)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

            model = self.trl_model_class.from_pretrained(model_name, summary_dropout_prob=0.5)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

    @pytest.mark.parametrize("model_name", ALL_SEQ2SEQ_MODELS)
    def test_generate(self, model_name):
        r"""
        Test if `generate` works for every model
        """
        generation_config = GenerationConfig(max_new_tokens=9)
        model = self.trl_model_class.from_pretrained(model_name).to(self.device)
        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
        decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)

        # Just check if the generation works
        _ = model.generate(input_ids, decoder_input_ids=decoder_input_ids, generation_config=generation_config)

    @pytest.mark.skip(reason="This test needs to be run manually due to HF token issue.")
    def test_push_to_hub(self):
        for model_name in self.all_model_names:
            model = self.trl_model_class.from_pretrained(model_name)
            if "sharded" in model_name:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True, max_shard_size="1MB")
            else:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True)

            model_from_pretrained = self.trl_model_class.from_pretrained(model_name + "-ppo")
            # check all keys
            assert model.state_dict().keys() == model_from_pretrained.state_dict().keys()

            for name, param in model.state_dict().items():
                (
                    torch.testing.assert_close(param, model_from_pretrained.state_dict()[name]),
                    (f"Parameter {name} is not the same after push_to_hub and from_pretrained"),
                )

    def test_transformers_bf16_kwargs(self):
        r"""
        Test if the transformers kwargs are correctly passed. Here we check that loading a model in half precision
        works as expected, i.e. the weights of the `pretrained_model` attribute is loaded in half precision and you can
        run a dummy forward pass without any issue.
        """
        for model_name in self.all_model_names:
            trl_model = self.trl_model_class.from_pretrained(model_name, dtype=torch.bfloat16).to(self.device)

            lm_head_namings = self.trl_model_class.lm_head_namings

            assert any(hasattr(trl_model.pretrained_model, lm_head_naming) for lm_head_naming in lm_head_namings)

            for lm_head_naming in lm_head_namings:
                if hasattr(trl_model.pretrained_model, lm_head_naming):
                    assert getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16

            dummy_input = torch.LongTensor([[0, 1, 0, 1]]).to(self.device)

            # check dummy forward pass works in half precision
            _ = trl_model(input_ids=dummy_input, decoder_input_ids=dummy_input)


@require_peft
class TestPeftModel(TrlTestCase):
    def setup_method(self):
        self.causal_lm_model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
        self.lora_config = LoraConfig(
            r=16,
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
        )

    def test_create_peft_model(self):
        r"""
        Simply creates a peft model and checks that it can be loaded.
        """
        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)

        _ = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)

    def test_peft_requires_grad(self):
        r"""
        Check that the value head of the returned model has requires_grad=True.
        """
        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)

        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)

        # Check that the value head has requires_grad=True
        assert model.v_head.summary.weight.requires_grad

    def test_check_peft_model_nb_trainable_params(self):
        r"""
        Check that the number of trainable parameters is correct.
        """
        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)

        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)

        # Check that the number of trainable parameters is correct
        nb_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        assert nb_trainable_params == 905

        # Check that the number of trainable param for the non-peft model is correct
        non_peft_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.causal_lm_model_id)
        nb_trainable_params = sum(p.numel() for p in non_peft_model.parameters() if p.requires_grad)
        assert nb_trainable_params == 2428641

    def test_create_peft_model_from_config(self):
        r"""
        Simply creates a peft model and checks that it can be loaded.
        """
        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(
            self.causal_lm_model_id, peft_config=self.lora_config
        )
        # Check that the number of trainable parameters is correct
        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)
        assert nb_trainable_params == 905

        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(causal_lm_model, peft_config=self.lora_config)
        # Check that the number of trainable parameters is correct
        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)
        assert nb_trainable_params == 905

    @require_bitsandbytes
    @require_torch_gpu_if_bnb_not_multi_backend_enabled
    def test_create_bnb_peft_model_from_config(self):
        r"""
        Simply creates a peft model and checks that it can be loaded.
        """
        from bitsandbytes.nn import Linear8bitLt
        from transformers import BitsAndBytesConfig

        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(
            self.causal_lm_model_id,
            peft_config=self.lora_config,
            quantization_config=BitsAndBytesConfig(load_in_8bit=True),
        )
        # Check that the number of trainable parameters is correct
        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)
        assert nb_trainable_params == 905
        assert isinstance(trl_model.pretrained_model.model.model.layers[0].mlp.gate_proj, Linear8bitLt)

        causal_lm_model = AutoModelForCausalLM.from_pretrained(
            self.causal_lm_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True), device_map="auto"
        )
        trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(causal_lm_model, peft_config=self.lora_config)
        # Check that the number of trainable parameters is correct
        nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)
        assert nb_trainable_params == 905
        assert isinstance(trl_model.pretrained_model.model.model.layers[0].mlp.gate_proj, Linear8bitLt)

    def test_save_pretrained_peft(self):
        r"""
        Check that the model can be saved and loaded properly.
        """
        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)

        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)

        model.save_pretrained(self.tmp_dir)

        # check that the files `adapter_model.safetensors` and `adapter_config.json` are in the directory
        assert os.path.isfile(f"{self.tmp_dir}/adapter_model.safetensors"), (
            f"{self.tmp_dir}/adapter_model.safetensors does not exist"
        )
        assert os.path.exists(f"{self.tmp_dir}/adapter_config.json"), (
            f"{self.tmp_dir}/adapter_config.json does not exist"
        )

        # check also for `pytorch_model.bin` and make sure it only contains `v_head` weights
        assert os.path.exists(f"{self.tmp_dir}/pytorch_model.bin"), f"{self.tmp_dir}/pytorch_model.bin does not exist"

        # check that only keys that starts with `v_head` are in the dict
        maybe_v_head = torch.load(f"{self.tmp_dir}/pytorch_model.bin", weights_only=True)
        assert all(k.startswith("v_head") for k in maybe_v_head.keys()), (
            f"keys in {self.tmp_dir}/pytorch_model.bin do not start with `v_head`"
        )

        model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(self.tmp_dir)

        # check all the weights are the same
        for p1, p2 in zip(model.named_parameters(), model_from_pretrained.named_parameters(), strict=True):
            torch.testing.assert_close(p1[1], p2[1], msg=f"{p1[0]} != {p2[0]}")

    def test_load_pretrained_peft(self):
        r"""
        Check that the model saved with peft class interface can be loaded properly.
        """
        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)

        model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)

        pretrained_model.save_pretrained(self.tmp_dir)
        model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(self.tmp_dir)

        # check that the files `adapter_model.safetensors` and `adapter_config.json` are in the directory
        assert os.path.isfile(f"{self.tmp_dir}/adapter_model.safetensors"), (
            f"{self.tmp_dir}/adapter_model.safetensors does not exist"
        )
        assert os.path.exists(f"{self.tmp_dir}/adapter_config.json"), (
            f"{self.tmp_dir}/adapter_config.json does not exist"
        )

        # check all the weights are the same
        for p1, p2 in zip(model.named_parameters(), model_from_pretrained.named_parameters(), strict=True):
            if p1[0] not in ["v_head.summary.weight", "v_head.summary.bias"]:
                torch.testing.assert_close(p1[1], p2[1], msg=f"{p1[0]} != {p2[0]}")

    def test_continue_training_peft_model(self):
        r"""
        Load peft and checks that it can continue training.
        """
        causal_lm_model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
        pretrained_model = get_peft_model(causal_lm_model, self.lora_config)

        pretrained_model.save_pretrained(self.tmp_dir)
        # set is_trainable to True
        model = AutoModelForCausalLMWithValueHead.from_pretrained(self.tmp_dir, is_trainable=True)
        # Check that the number of trainable parameters is correct
        nb_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        assert nb_trainable_params == 905


class TestCore(TrlTestCase):
    """
    A wrapper class for testing core utils functions
    """

    def setup_method(self):
        self.test_input = torch.Tensor([1, 2, 3, 4])
        self.test_mask = torch.Tensor([0, 1, 1, 0])
        self.test_input_unmasked = self.test_input[1:3]

    def test_masked_mean(self):
        assert torch.mean(self.test_input_unmasked) == masked_mean(self.test_input, self.test_mask)

    def test_masked_var(self):
        assert torch.var(self.test_input_unmasked) == masked_var(self.test_input, self.test_mask)

    def test_masked_whiten(self):
        def whiten(values: torch.Tensor) -> torch.Tensor:
            mean, var = torch.mean(values), torch.var(values)
            return (values - mean) * torch.rsqrt(var + 1e-8)

        whiten_unmasked = whiten(self.test_input_unmasked)
        whiten_masked = masked_whiten(self.test_input, self.test_mask)[1:3]
        diffs = (whiten_unmasked - whiten_masked).sum()
        assert abs(diffs.item()) < 0.00001


class TestPPOTrainer(TrlTestCase):
    def setup_method(self):
        # Set up the models and tokenizer using the test model
        self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
        self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
        self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, padding_side="left")
        self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})

        # Add reward and value models as in ppo.py
        reward_model_id = "trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5"
        self.value_model = AutoModelForSequenceClassification.from_pretrained(reward_model_id, num_labels=1)
        self.reward_model = AutoModelForSequenceClassification.from_pretrained(reward_model_id, num_labels=1)

        # Load dataset
        raw_dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only")

        def tokenize(example, tokenizer):
            tokenized = tokenizer(text=example["prompt"])
            if tokenizer.eos_token_id is not None and tokenized["input_ids"][-1] != tokenizer.eos_token_id:
                tokenized["input_ids"] = tokenized["input_ids"] + [tokenizer.eos_token_id]
                tokenized["attention_mask"] = tokenized["attention_mask"] + [1]
            return tokenized

        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
        initial_critic_weights = {}
        initial_policy_weights = {}
        for name, param in self.value_model.named_parameters():
            initial_critic_weights[name] = param.clone().detach()
        for name, param in self.model.named_parameters():
            initial_policy_weights[name] = param.clone().detach()

        # Configure training args similar to example script
        training_args = PPOConfig(
            output_dir=self.tmp_dir,
            per_device_train_batch_size=4,
            per_device_eval_batch_size=2,
            num_ppo_epochs=2,  # Decrease number of PPO epochs to speed up test
            report_to="none",
        )

        # Create trainer
        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"],
        )

        # Train
        trainer.train()

        # Check if critic weights have been updated
        critic_weights_updated = False
        for name, param in trainer.model.value_model.named_parameters():
            if not torch.equal(initial_critic_weights[name], param.to("cpu")):
                critic_weights_updated = True
                break

        # Check if policy weights have been updated
        policy_weights_updated = False
        for name, param in trainer.model.policy.named_parameters():
            if not torch.equal(initial_policy_weights[name], param.to("cpu")):
                policy_weights_updated = True
                break

        assert critic_weights_updated, "Critic weights were not updated during training"
        assert policy_weights_updated, "Policy weights were not updated during training"

    @require_peft
    def test_peft_training(self):
        """Test PPO training with PEFT configuration and verify model updates."""
        # Capture initial weights
        initial_critic_weights = {}
        initial_policy_weights = {}
        for name, param in self.value_model.named_parameters():
            initial_critic_weights[name] = param.clone().detach()
        for name, param in self.model.named_parameters():
            initial_policy_weights[name] = param.clone().detach()

        # Configure training args
        training_args = PPOConfig(
            output_dir=self.tmp_dir,
            per_device_train_batch_size=4,
            per_device_eval_batch_size=2,
            num_ppo_epochs=2,  # Decrease number of PPO epochs to speed up test
            report_to="none",
        )

        # Configure PEFT
        peft_config = LoraConfig(
            r=32,
            lora_alpha=16,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
        )

        # Create trainer with PEFT
        trainer = PPOTrainer(
            args=training_args,
            processing_class=self.tokenizer,
            model=self.model,
            ref_model=None,
            reward_model=self.reward_model,
            value_model=self.value_model,
            train_dataset=self.raw_dataset["train"],
            eval_dataset=self.raw_dataset["test"],
            peft_config=peft_config,
        )

        # Train
        trainer.train()

        # Check if critic weights have been updated
        critic_weights_updated = False
        for name, param in trainer.model.value_model.named_parameters():
            if name in initial_critic_weights and not torch.equal(initial_critic_weights[name], param.to("cpu")):
                critic_weights_updated = True
                break

        # Check if policy weights have been updated - for PEFT we check the LoRA weights
        policy_weights_updated = False
        for name, param in trainer.model.policy.named_parameters():
            if "lora" in name.lower() and param.requires_grad:  # Only check LoRA weights
                # New weights should be non-zero if they've been updated
                if not torch.allclose(param, torch.zeros_like(param)):
                    policy_weights_updated = True
                    break

        assert critic_weights_updated, "Critic weights were not updated during training"
        assert policy_weights_updated, "Policy LoRA weights were not updated during training"