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:
      **cfg** (`LLM Studio Config`) -- Experiment configuration object. Check the
          documentation of `LLM Studio Config` 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.
      **tokenizer** (`Union[PreTrainedTokenizer, PreTrainedTokenizerFast]`)
          Tokenizer to be used for encoding the data. Check the documentation of
          `transformers.PreTrainedTokenizer` and
          `transformers.PreTrainedTokenizerFast` for more details.
      **optimizer** (`torch.optim.Optimizer`) -- Optimizer to be used for training.
      **lr_scheduler** (`torch.optim.lr_scheduler`) -- Learning rate scheduler to be
          used for training.
      **scaler** (`torch.cuda.amp.GradScaler`) -- Gradient scaler 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",":","*","*","cfg","*","*","(","`","LLM","Studio","Config","`",")","--","Experiment","configuration","object",".","Check","the","documentation","of","`","LLM","Studio","Config","`","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",".","*","*","tokenizer","*","*","(","`","Union","[","PreTrainedTokenizer",",","PreTrainedTokenizerFast","]","`",")","Tokenizer","to","be","used","for","encoding","the","data",".","Check","the","documentation","of","`","transformers.PreTrainedTokenizer","`","and","`","transformers.PreTrainedTokenizerFast","`","for","more","details",".","*","*","optimizer","*","*","(","`","torch.optim.Optimizer","`",")","--","Optimizer","to","be","used","for","training",".","*","*","lr_scheduler","*","*","(","`","torch.optim.lr_scheduler","`",")","--","Learning","rate","scheduler","to","be","used","for","training",".","*","*","scaler","*","*","(","`","torch.cuda.amp.GradScaler","`",")","--","Gradient","scaler","to","be","used","for","training","."]
docstr: Run a PPO optimisation step given a list of queries, model responses, and
  rewards.
  
  Args:
      queries (List[`torch.LongTensor`]):
          List of tensors containing the encoded queries of shape (`query_length`)
      responses (List[`torch.LongTensor`]):
          List of tensors containing the encoded responses of shape
          (`response_length`)
      scores (List[`torch.Tensor`]):
          List of tensors containing the scores.
  
  Returns:
      `dict[str, Any]`: A summary of the training statistics
code: def step(
      self,
      queries: List[torch.LongTensor],
      responses: List[torch.LongTensor],
      scores: List[torch.Tensor],
  ):
      """
      Run a PPO optimisation step given a list of queries, model responses, and
      rewards.
  
      Args:
          queries (List[`torch.LongTensor`]):
              List of tensors containing the encoded queries of shape (`query_length`)
          responses (List[`torch.LongTensor`]):
              List of tensors containing the encoded responses of shape
              (`response_length`)
          scores (List[`torch.Tensor`]):
              List of tensors containing the scores.
  
      Returns:
          `dict[str, Any]`: A summary of the training statistics
      """
  
      queries, responses, scores = self._step_safety_checker(
          queries, responses, scores
      )
  
      timing = dict()
      t0 = time.time()
  
      t = time.time()
  
      model_inputs = self.prepare_model_inputs(queries, responses)
  
      model_inputs_names = list(model_inputs.keys())
  
      with torch.no_grad():
          (
              all_logprobs,
              logits_or_none,
              values,
              masks,
          ) = self.batched_forward_pass(
              self.model,
              queries,
              responses,
              model_inputs,
              return_logits=self.cfg.training.full_kl_penalty,
          )
          with unwrap_model(self.model).backbone.disable_adapter():
              (
                  ref_logprobs,
                  ref_logits_or_none,
                  _,
                  _,
              ) = self.batched_forward_pass(
                  self.model,
                  queries,
                  responses,
                  model_inputs,
                  return_logits=self.cfg.training.full_kl_penalty,
                  return_values=False,
              )
  
      timing["time/ppo/forward_pass"] = time.time() - t
  
      with torch.no_grad():
          t = time.time()
          if self.cfg.training.full_kl_penalty:
              active_full_logprobs = logprobs_from_logits(
                  logits_or_none, None, gather=False
              )
              ref_full_logprobs = logprobs_from_logits(
                  ref_logits_or_none, None, gather=False
              )
  
              rewards, non_score_reward = self.compute_rewards(
                  scores, active_full_logprobs, ref_full_logprobs, masks
              )
          else:
              rewards, non_score_reward = self.compute_rewards(
                  scores, all_logprobs, ref_logprobs, masks
              )
          timing["time/ppo/compute_rewards"] = time.time() - t
  
          t = time.time()
          values, advantages, returns = self.compute_advantages(
              values, rewards, masks
          )
          timing["time/ppo/compute_advantages"] = time.time() - t
  
      # upcast to float32 to avoid dataset issues
      mini_batch_dict = {
          "queries": queries,
          "responses": responses,
          "logprobs": all_logprobs.to(torch.float32),
          "values": values.to(torch.float32),
          "masks": masks,
          "advantages": advantages,
          "returns": returns,
      }
  
      def collator(data: List[Dict[str, torch.Tensor]]):
          return_dict: Dict[str, Any] = dict()
          keys = data[0].keys()
          for key in keys:
              if key in ["queries", "responses"]:
                  return_dict[key] = [d[key] for d in data]
              else:
                  return_dict[key] = torch.stack(
                      [d[key] for d in data]
                  ).to(self.current_device)
          return return_dict
  
      mini_batch_dict.update(model_inputs)
      mini_batch_data = Dataset.from_dict(mini_batch_dict)
      mini_batch_data.set_format("torch")
      mini_batch_dataloader = DataLoader(
          mini_batch_data,
          batch_size=self.cfg.training.ppo_batch_size,
          shuffle=True,
          collate_fn=collator,
      )
  
      t = time.time()
      all_stats = []
  
      if (
          self.cfg.training.ppo_batch_size
          > self.cfg.training.batch_size
          * self.cfg.training.rollout_steps
      ):
          raise ValueError(
              f"ppo_batch_size ({self.cfg.training.ppo_batch_size}) "
              "must not be larger than the rollout "
              f"({self.cfg.training.batch_size * self.cfg.training.rollout_steps})"
          )
  
      for _ in range(self.cfg.training.ppo_epochs):
          for batch in mini_batch_dataloader:
              self.num_updates += 1
  
              model_inputs = {k: batch[k] for k in model_inputs_names}
              logprobs, logits, vpreds, _ = self.batched_forward_pass(
                  model=self.model,
                  queries=batch["queries"],
                  responses=batch["responses"],
                  model_inputs=model_inputs,
                  return_logits=True,
              )
  
              loss_p, loss_v, train_stats = self.loss(
                  old_logprobs=batch["logprobs"],
                  values=batch["values"],
                  logits=logits,
                  vpreds=vpreds,
                  logprobs=logprobs,
                  mask=batch["masks"],
                  advantages=batch["advantages"],
                  returns=batch["returns"],
              )
              loss = loss_p + loss_v
  
              # loss is a mean loss per batch/sample
              # as grad_accumulations sums up the gradients, this loss must be scaled
              # by the number of grad_accumulations, to have similar behavior for
              # BS * grad_accumulations = const.
              if self.cfg.training.grad_accumulation != 1:
                  loss = loss / self.cfg.training.grad_accumulation
  
              # Backward pass
              if self.cfg.environment.mixed_precision:
                  self.scaler.scale(loss).backward()
                  if (
                      self.num_updates
                      % self.cfg.training.grad_accumulation
                      == 0
                  ):
                      if self.cfg.training.gradient_clip > 0:
                          self.scaler.unscale_(self.optimizer)
                          torch.nn.utils.clip_grad_norm_(
                              self.model.parameters(),
                              self.cfg.training.gradient_clip,
                          )
                      self.scaler.step(self.optimizer)
                      self.scaler.update()
                      self.optimizer.zero_grad(set_to_none=True)
              else:
                  loss.backward()
                  if (
                      self.num_updates
                      % self.cfg.training.grad_accumulation
                      == 0
                  ):
                      if self.cfg.training.gradient_clip > 0:
                          torch.nn.utils.clip_grad_norm_(
                              self.model.parameters(),
                              self.cfg.training.gradient_clip,
                          )
                      self.optimizer.step()
                      self.optimizer.zero_grad(set_to_none=True)
  
              if self.cfg.environment._distributed:
                  torch.cuda.synchronize(device=self.current_device)
  
              del logprobs, logits, vpreds
  
              all_stats.append(train_stats)
  
      timing["time/ppo/ppo_steps"] = time.time() - t
  
      t = time.time()
      train_stats = stack_dicts(all_stats)
  
      # reshape advantages/ratios such that they are not averaged.
      train_stats["policy/advantages"] = torch.flatten(
          train_stats["policy/advantages"]
      ).unsqueeze(0)
      train_stats["policy/advantages"] = torch.nan_to_num(
          train_stats["policy/advantages"], -1
      )
      train_stats["policy/ratio"] = torch.flatten(
          train_stats["policy/ratio"]
      ).unsqueeze(0)
  
      stats = self.record_step_stats(
          scores=scores,
          logprobs=all_logprobs,
          ref_logprobs=ref_logprobs,
          non_score_reward=non_score_reward,
          train_stats=train_stats,
          kl_coef=self.kl_ctl.value,
          masks=masks,
          queries=queries,
          responses=responses,
      )
      stats = stats_to_np(stats)
      timing["time/ppo/calc_stats"] = time.time() - t
      stats["ppo/learning_rate"] = self.optimizer.param_groups[0]["lr"]
  
      # Update the KL control - multiply the batch_size by the number of processes
      self.kl_ctl.update(
          stats["objective/kl"], self.cfg.training.batch_size
      )
  
      # Log the total ppo time
      timing["time/ppo/total"] = time.time() - t0
      stats.update(timing)
  
      if self.lr_scheduler is not None:
          self.lr_scheduler.step()
  
      return stats
  
imports: import math
  import time
  import warnings
  from collections.abc import Mapping
  from typing import Any, Dict, List, Optional, Union
  import numpy
  import torch
  import torch.nn.functional
  from datasets import Dataset
  from huggingface_hub import PyTorchModelHubMixin
  from torch.cuda.amp import autocast
  from torch.nn.utils.rnn import pad_sequence
  from torch.utils.data import DataLoader
  from transformers import DataCollatorForLanguageModeling, PreTrainedTokenizer, PreTrainedTokenizerFast
  from llm_studio.src.utils.modeling_utils import unwrap_model
label_desc: Use image node_id 3 for calling the PPOTrainer obj's underlying member method code with example usage: obj.step(queries, responses, scores) and returns: stats, return_dict