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","."]
code: def _prepare_deepspeed(self, model: PreTrainedModelWrapper):
      # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473
      deepspeed_plugin = self.accelerator.state.deepspeed_plugin
      config_kwargs = deepspeed_plugin.deepspeed_config
      if model is not None:
          if hasattr(model, "config"):
              hidden_size = (
                  max(model.config.hidden_sizes)
                  if getattr(model.config, "hidden_sizes", None)
                  else getattr(model.config, "hidden_size", None)
              )
              if (
                  hidden_size is not None
                  and config_kwargs["zero_optimization"]["stage"] == 3
              ):
                  # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`
                  # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081
                  config_kwargs.update(
                      {
                          "zero_optimization.reduce_bucket_size": hidden_size
                          * hidden_size,
                          "zero_optimization.stage3_param_persistence_threshold": 10
                          * hidden_size,
                          "zero_optimization.stage3_prefetch_bucket_size": 0.9
                          * hidden_size
                          * hidden_size,
                      }
                  )
  
      # If ZeRO-3 is used, we shard both the active and reference model.
      # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)
      if config_kwargs["zero_optimization"]["stage"] != 3:
          config_kwargs["zero_optimization"]["stage"] = 0
      model, *_ = deepspeed.initialize(
          model=model, config=config_kwargs
      )
      model.eval()
      return model
  
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 24 for calling the PPOTrainer obj's underlying member method code with example usage: obj._prepare_deepspeed(model) and returns: model