class_docstr: The PPOTrainer uses Proximal Policy Optimization to optimise language models.
  Note, this trainer is heavily inspired by the original OpenAI learning to summarize work here:
  https://github.com/openai/summarize-from-feedback
  
  Attributes:
      **config** (`PPOConfig`) -- Configuration object for PPOTrainer. Check the documentation of `PPOConfig` for more
          details.
      **model** (`PreTrainedModelWrapper`) -- Model to be optimized, Hugging Face transformer model with a value head.
          Check the documentation of `PreTrainedModelWrapper` for more details.
      **ref_model** (`PreTrainedModelWrapper`, *optional*) -- Reference model to be used for KL penalty, Hugging Face
          transformer model with a casual language modelling head. Check the documentation of `PreTrainedModelWrapper`
          for more details. If no reference model is provided, the trainer will create a reference model with the same
           architecture as the model to be optimized with shared layers.
      **tokenizer** (`PreTrainedTokenizerBase`) -- Tokenizer to be used for encoding the
          data. Check the documentation of `transformers.PreTrainedTokenizer` and
          `transformers.PreTrainedTokenizerFast` for more details.
      **dataset** (Union[`torch.utils.data.Dataset`, `datasets.Dataset`], *optional*) -- PyTorch dataset or Hugging
          Face dataset. This is used to create a PyTorch dataloader. If no dataset is provided, the dataloader must be
           created outside the trainer users needs to design their own dataloader and make sure the batch
          size that is used is the same as the one specified in the configuration object.
      **optimizer** (`torch.optim.Optimizer`, *optional*) -- Optimizer to be used for training. If no optimizer is
          provided, the trainer will create an Adam optimizer with the learning rate specified in the configuration
          object.
      **data_collator** (DataCollatorForLanguageModeling, *optional*) -- Data collator to be used for training and
          passed along the dataloader
      **num_shared_layers** (int, *optional*) -- Number of layers to be shared between the model and the reference
          model, if no reference model is passed. If no number is provided, all the layers will be shared.
      **lr_scheduler** (`torch.optim.lr_scheduler`, *optional*) -- Learning rate scheduler to be used for training.
class_docstr_tok: ["The","PPOTrainer","uses","Proximal","Policy","Optimization","to","optimise","language","models",".","Note",",","this","trainer","is","heavily","inspired","by","the","original","OpenAI","learning","to","summarize","work","here",":","https",":","\/\/github.com\/openai\/summarize-from-feedback","Attributes",":","*","*","config","*","*","(","`","PPOConfig","`",")","--","Configuration","object","for","PPOTrainer",".","Check","the","documentation","of","`","PPOConfig","`","for","more","details",".","*","*","model","*","*","(","`","PreTrainedModelWrapper","`",")","--","Model","to","be","optimized",",","Hugging","Face","transformer","model","with","a","value","head",".","Check","the","documentation","of","`","PreTrainedModelWrapper","`","for","more","details",".","*","*","ref_model","*","*","(","`","PreTrainedModelWrapper","`",",","*","optional","*",")","--","Reference","model","to","be","used","for","KL","penalty",",","Hugging","Face","transformer","model","with","a","casual","language","modelling","head",".","Check","the","documentation","of","`","PreTrainedModelWrapper","`","for","more","details",".","If","no","reference","model","is","provided",",","the","trainer","will","create","a","reference","model","with","the","same","architecture","as","the","model","to","be","optimized","with","shared","layers",".","*","*","tokenizer","*","*","(","`","PreTrainedTokenizerBase","`",")","--","Tokenizer","to","be","used","for","encoding","the","data",".","Check","the","documentation","of","`","transformers.PreTrainedTokenizer","`","and","`","transformers.PreTrainedTokenizerFast","`","for","more","details",".","*","*","dataset","*","*","(","Union","[","`","torch.utils.data.Dataset","`",",","`","datasets.Dataset","`","]",",","*","optional","*",")","--","PyTorch","dataset","or","Hugging","Face","dataset",".","This","is","used","to","create","a","PyTorch","dataloader",".","If","no","dataset","is","provided",",","the","dataloader","must","be","created","outside","the","trainer","users","needs","to","design","their","own","dataloader","and","make","sure","the","batch","size","that","is","used","is","the","same","as","the","one","specified","in","the","configuration","object",".","*","*","optimizer","*","*","(","`","torch.optim.Optimizer","`",",","*","optional","*",")","--","Optimizer","to","be","used","for","training",".","If","no","optimizer","is","provided",",","the","trainer","will","create","an","Adam","optimizer","with","the","learning","rate","specified","in","the","configuration","object",".","*","*","data_collator","*","*","(","DataCollatorForLanguageModeling",",","*","optional","*",")","--","Data","collator","to","be","used","for","training","and","passed","along","the","dataloader","*","*","num_shared_layers","*","*","(","int",",","*","optional","*",")","--","Number","of","layers","to","be","shared","between","the","model","and","the","reference","model",",","if","no","reference","model","is","passed",".","If","no","number","is","provided",",","all","the","layers","will","be","shared",".","*","*","lr_scheduler","*","*","(","`","torch.optim.lr_scheduler","`",",","*","optional","*",")","--","Learning","rate","scheduler","to","be","used","for","training","."]
docstr: Initialize PPOTrainer.
  
  Args:
      config (`PPOConfig`):
          Configuration object for PPOTrainer. Check the documentation of `PPOConfig` for more details.
      model (`PreTrainedModelWrapper`):
          Hugging Face transformer model with a value head.
      ref_model (`PreTrainedModelWrapper`):
          Hugging Face transformer model with a casual language modelling head. Used for KL penalty
      tokenizer (`transformers.PreTrainedTokenizerBase`):
          Hugging Face tokenizer
      dataset (Optional[Union[`torch.utils.data.Dataset`, `datasets.Dataset`]]):
          PyTorch dataset or Hugging Face dataset. If a Hugging Face dataset is passed, the dataset
          will be preprocessed by removing the columns that are not used by the model. If none is passed,
          a warning will be raised in a multi-GPU setting.
      optimizer (Optional[`torch.optim.Optimizer`]):
          Optimizer used for training. If `None`, the `Adam` is used as default.
      data_collator (Optional[function]):
          Data collator function.
      num_shared_layers (Optional[int]):
          Number of shared layers between the model and the reference model. If `None`, all layers are shared.
          used only if `ref_model` is `None`.
      lr_scheduler (Optional[`torch.optim.lr_scheduler`]):
          Learning rate scheduler used for training.
docstr_tok: ["Initialize","PPOTrainer",".","Args",":","config","(","`","PPOConfig","`",")",":","Configuration","object","for","PPOTrainer",".","Check","the","documentation","of","`","PPOConfig","`","for","more","details",".","model","(","`","PreTrainedModelWrapper","`",")",":","Hugging","Face","transformer","model","with","a","value","head",".","ref_model","(","`","PreTrainedModelWrapper","`",")",":","Hugging","Face","transformer","model","with","a","casual","language","modelling","head",".","Used","for","KL","penalty","tokenizer","(","`","transformers.PreTrainedTokenizerBase","`",")",":","Hugging","Face","tokenizer","dataset","(","Optional","[","Union","[","`","torch.utils.data.Dataset","`",",","`","datasets.Dataset","`","]","]",")",":","PyTorch","dataset","or","Hugging","Face","dataset",".","If","a","Hugging","Face","dataset","is","passed",",","the","dataset","will","be","preprocessed","by","removing","the","columns","that","are","not","used","by","the","model",".","If","none","is","passed",",","a","warning","will","be","raised","in","a","multi-GPU","setting",".","optimizer","(","Optional","[","`","torch.optim.Optimizer","`","]",")",":","Optimizer","used","for","training",".","If","`","None","`",",","the","`","Adam","`","is","used","as","default",".","data_collator","(","Optional","[","function","]",")",":","Data","collator","function",".","num_shared_layers","(","Optional","[","int","]",")",":","Number","of","shared","layers","between","the","model","and","the","reference","model",".","If","`","None","`",",","all","layers","are","shared",".","used","only","if","`","ref_model","`","is","`","None","`",".","lr_scheduler","(","Optional","[","`","torch.optim.lr_scheduler","`","]",")",":","Learning","rate","scheduler","used","for","training","."]
code: def __init__(
      self,
      config: PPOConfig = None,
      model: PreTrainedModelWrapper = None,
      ref_model: Optional[PreTrainedModelWrapper] = None,
      tokenizer: PreTrainedTokenizerBase = None,
      dataset: Optional[
          Union[torch.utils.data.Dataset, Dataset]
      ] = None,
      optimizer: Optional[torch.optim.Optimizer] = None,
      data_collator: Optional[typing.Callable] = None,
      num_shared_layers: Optional[int] = None,
      lr_scheduler: Optional[
          torch.optim.lr_scheduler._LRScheduler
      ] = None,
  ):
      """
      Initialize PPOTrainer.
  
      Args:
          config (`PPOConfig`):
              Configuration object for PPOTrainer. Check the documentation of `PPOConfig` for more details.
          model (`PreTrainedModelWrapper`):
              Hugging Face transformer model with a value head.
          ref_model (`PreTrainedModelWrapper`):
              Hugging Face transformer model with a casual language modelling head. Used for KL penalty
          tokenizer (`transformers.PreTrainedTokenizerBase`):
              Hugging Face tokenizer
          dataset (Optional[Union[`torch.utils.data.Dataset`, `datasets.Dataset`]]):
              PyTorch dataset or Hugging Face dataset. If a Hugging Face dataset is passed, the dataset
              will be preprocessed by removing the columns that are not used by the model. If none is passed,
              a warning will be raised in a multi-GPU setting.
          optimizer (Optional[`torch.optim.Optimizer`]):
              Optimizer used for training. If `None`, the `Adam` is used as default.
          data_collator (Optional[function]):
              Data collator function.
          num_shared_layers (Optional[int]):
              Number of shared layers between the model and the reference model. If `None`, all layers are shared.
              used only if `ref_model` is `None`.
          lr_scheduler (Optional[`torch.optim.lr_scheduler`]):
              Learning rate scheduler used for training.
      """
      super().__init__(config)
  
      # initial seed for reproducible experiments
      set_seed(config.seed)
  
      # Step 0: check positional arguments validity
      if not isinstance(config, PPOConfig):
          raise ValueError(
              f"config must be a PPOConfig, got {type(config)}"
          )
      if not isinstance(tokenizer, (PreTrainedTokenizerBase)):
          raise ValueError(
              f"tokenizer must be a PreTrainedTokenizerBase like a PreTrainedTokenizer or a PreTrainedTokenizerFast, got {type(tokenizer)}"
          )
      if not isinstance(model, (SUPPORTED_ARCHITECTURES)):
          raise ValueError(
              f"model must be a PreTrainedModelWrapper, got {type(model)} - supported architectures are: {SUPPORTED_ARCHITECTURES}"
          )
      # Step 1: Initialize Accelerator
      self.accelerator = Accelerator(
          log_with=config.log_with,
          gradient_accumulation_steps=config.gradient_accumulation_steps,
          project_config=ProjectConfiguration(**config.project_kwargs),
          **config.accelerator_kwargs,
      )
  
      # Step 1.1 Runtime variables filled by the accelerator
      config.world_size = self.accelerator.num_processes
      config.global_backward_batch_size = (
          config.backward_batch_size * config.world_size
      )
      config.global_batch_size = config.batch_size * config.world_size
  
      self.model = model
      self.model_params = filter(
          lambda p: p.requires_grad, self.model.parameters()
      )
      self.is_encoder_decoder = hasattr(
          self.model, "is_encoder_decoder"
      )
      self.is_peft_model = getattr(self.model, "is_peft_model", False)
      config.is_encoder_decoder = self.is_encoder_decoder
      config.is_peft_model = self.is_peft_model
  
      is_using_tensorboard = (
          config.log_with is not None
          and config.log_with == "tensorboard"
      )
      self.accelerator.init_trackers(
          config.tracker_project_name,
          config=dict(trl_ppo_trainer_config=config.to_dict())
          if not is_using_tensorboard
          else config.to_dict(),
          init_kwargs=config.tracker_kwargs,
      )
      self.is_using_text_environment = getattr(
          config, "use_text_environment", False
      )
  
      if isinstance(ref_model, SUPPORTED_ARCHITECTURES):
          self.ref_model = ref_model
          if num_shared_layers is not None:
              warnings.warn(
                  "num_shared_layers is ignored when ref_model is provided. Two different models are used for the "
                  "model and the reference model and no layers are shared.",
                  UserWarning,
              )
      elif ref_model is None and not self.is_peft_model:
          self.ref_model = create_reference_model(
              self.model, num_shared_layers=num_shared_layers
          )
      elif self.is_peft_model:
          self.ref_model = None
      else:
          raise ValueError(
              f"ref_model must be a PreTrainedModelWrapper or `None`, got {type(ref_model)} - supported "
              f"architectures are: {SUPPORTED_ARCHITECTURES} "
          )
      self.optional_peft_ctx = (
          self.accelerator.unwrap_model(
              self.model
          ).pretrained_model.disable_adapter
          if self.is_peft_model
          else nullcontext
      )
  
      if not (
          isinstance(tokenizer, PreTrainedTokenizer)
          or isinstance(tokenizer, PreTrainedTokenizerFast)
      ):
          raise ValueError(
              "tokenizer must be a transformers.PreTrainedTokenizer or transformers.PreTrainedTokenizerFast"
          )
      self.tokenizer = tokenizer
  
      if dataset is not None and not (
          isinstance(dataset, torch.utils.data.Dataset)
          or isinstance(dataset, Dataset)
      ):
          raise ValueError(
              "dataset must be a torch.utils.data.Dataset or datasets.Dataset"
          )
      elif dataset is None:
          warnings.warn(
              "No dataset is provided. Make sure to set config.batch_size to the correct value before training.",
              UserWarning,
          )
      self.dataset = dataset
      self._signature_columns = None
      if self.dataset is not None:
          self.dataloader = self.prepare_dataloader(
              self.dataset, data_collator
          )
      elif self.dataset is None and self.accelerator.num_processes > 1:
          warnings.warn(
              "No dataset is provided. In a multi-GPU setting, this will lead to an error. You should"
              " prepare your dataloader yourself with `dataloader = ppo_trainer.accelerator.prepare(dataloader)`"
              " and using `torch.utils.data.DataLoader`, or pass a dataset to the `PPOTrainer`. Please "
              " refer to the documentation for more details.",
              UserWarning,
          )
          self.dataloader = None
      else:
          self.dataloader = None
  
      # Step 3: Initialize optimizer and data collator
      self.data_collator = DataCollatorForLanguageModeling(
          self.tokenizer, mlm=False
      )
      if optimizer is None:
          self.optimizer = Adam(
              filter(
                  lambda p: p.requires_grad, self.model.parameters()
              ),
              lr=self.config.learning_rate,
          )
      else:
          self.optimizer = optimizer
  
      self.lr_scheduler = lr_scheduler
      if self.lr_scheduler is not None:
          lr_scheduler_class = (
              torch.optim.lr_scheduler._LRScheduler
              if not is_torch_greater_2_0()
              else torch.optim.lr_scheduler.LRScheduler
          )
  
          if not isinstance(self.lr_scheduler, lr_scheduler_class):
              raise ValueError(
                  "lr_scheduler must be a torch.optim.lr_scheduler._LRScheduler or torch.optim.lr_scheduler.LRScheduler (for torch >= 2.0)"
              )
  
      if self.config.adap_kl_ctrl:
          self.kl_ctl = AdaptiveKLController(
              self.config.init_kl_coef,
              self.config.target,
              self.config.horizon,
          )
      else:
          self.kl_ctl = FixedKLController(self.config.init_kl_coef)
  
      # Safety checkers for DS integration
      is_deepspeed_used = (
          self.accelerator.distributed_type == "DEEPSPEED"
          and hasattr(self.accelerator.state, "deepspeed_plugin")
      )
  
      (
          self.model,
          self.optimizer,
          self.data_collator,
          self.dataloader,
          self.lr_scheduler,
      ) = self.accelerator.prepare(
          self.model,
          self.optimizer,
          self.data_collator,
          self.dataloader,
          self.lr_scheduler,
      )
      if is_deepspeed_used:
          # Quantized models are already set on the correct device
          if not self.is_peft_model and not (
              getattr(
                  self.ref_model.pretrained_model,
                  "is_loaded_in_8bit",
                  False,
              )
              or getattr(
                  self.ref_model.pretrained_model,
                  "is_loaded_in_4bit",
                  False,
              )
          ):
              self.ref_model = self._prepare_deepspeed(self.ref_model)
      else:
          self.ref_model = self.accelerator.prepare(self.ref_model)
  
      # In a distributed setup, only logging needs to be performed on the main process
      # check: https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
      # or: https://discuss.pytorch.org/t/use-distributed-data-parallel-correctly/82500/11
      self.is_distributed = self.accelerator.num_processes > 1
  
      # init the current step
      self.current_step = 0
  
      # init variables for pushing model to hub
      if config.push_to_hub_if_best_kwargs:
          if "repo_id" not in config.push_to_hub_if_best_kwargs:
              raise ValueError(
                  "You have to specify repo_id in order to push the model to the hub!"
              )
          self.push_to_hub_kwargs = config.push_to_hub_if_best_kwargs
          self.compare_step = 0
          self.highest_reward = torch.tensor(-float("inf"))
  
      # post process for PP
      if not getattr(self.model, "is_sequential_parallel", False):
          self.current_device = self.accelerator.device
      else:
          if is_xpu_available():
              self.current_device = torch.device("xpu:0")
          elif is_npu_available():
              self.current_device = torch.device("npu:0")
          else:
              self.current_device = torch.device("cuda:0")
  
      PPODecorators.optimize_device_cache = (
          self.config.optimize_device_cache
      )
  
      self.running = RunningMoments(self.accelerator)
  
imports: import inspect
  import math
  import os
  import time
  import typing
  import warnings
  from contextlib import nullcontext
  from typing import Callable, List, Optional, Union
  import datasets
  import numpy
  import torch
  import torch.nn.functional
  from accelerate import Accelerator
  from accelerate.utils import ProjectConfiguration, gather_object, is_deepspeed_available
  from datasets import Dataset
  from huggingface_hub import whoami
  from packaging import version
  from torch.optim import Adam
  from transformers import DataCollatorForLanguageModeling, PreTrainedTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizerFast
  from ..core import WANDB_PADDING, PPODecorators, clip_by_value, convert_to_scalar, entropy_from_logits, flatten_dict, logprobs_from_logits, masked_mean, masked_var, masked_whiten, set_seed, stack_dicts, stats_to_np
  from ..import_utils import is_npu_available, is_torch_greater_2_0, is_xpu_available
  from ..models import SUPPORTED_ARCHITECTURES, PreTrainedModelWrapper, create_reference_model
  from .None import AdaptiveKLController, BaseTrainer, FixedKLController, PPOConfig, RunningMoments
label_desc: Use image node_id 1 to create a new PPOTrainer object from inherited base classes: BaseTrainer with example: obj = PPOTrainer(config, model, ref_model, tokenizer, dataset, optimizer, data_collator, num_shared_layers, lr_scheduler)