class_docstr: A training engine hybrid pipeline, data, and model parallel training.
  
  This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule`
  is provided.
class_docstr_tok: ["A","training","engine","hybrid","pipeline",",","data",",","and","model","parallel","training",".","This","engine","is","created","by","``","deepspeed.initialize","(",")","``","when","a",":","class",":","`","PipelineModule","`","is","provided","."]
code: def __init__(
      self, has_bool_tensors=False, *super_args, **super_kwargs
  ):
      super().__init__(*super_args, **super_kwargs)
      assert isinstance(
          self.module, PipelineModule
      ), "model must base PipelineModule"
  
      assert (
          self.zero_optimization_stage() < ZeroStageEnum.gradients
      ), "ZeRO-2 and ZeRO-3 are incompatible with pipeline parallelism"
  
      # We schedule the all-reduces, so disable it in super().backward()
      self.enable_backward_allreduce = False
      self.has_bool_tensors = has_bool_tensors
      self.eval_return_logits = False
      self.outputs = None
      # BF16 Optimizer is hardcoded for fp32 gradient accumulation
      self.using_bf16_optimizer = type(self.optimizer) == BF16_Optimizer
  
      # used to disable the pipeline all-reduce when used with 1-bit Adam/1-bit LAMB
      self.pipeline_enable_backward_allreduce = True
  
      if self.elasticity_enabled():
          if not self.is_elastic_model_parallel_supported():
              assert not self.elasticity_enabled(), (
                  "Elasticity is not currently supported"
                  " with pipeline parallelism."
              )
  
      # pipeline step for logging
      self.log_batch_step_id = -1
  
      self.micro_batch_size = self.train_micro_batch_size_per_gpu()
      self.micro_batches = self.gradient_accumulation_steps()
  
      # Set Grid and Communication Groups
      self.grid = self.module._grid
      if self.grid.get_global_rank() == 0:
          logger.info(
              f"CONFIG: micro_batches={self.micro_batches} "
              f"micro_batch_size={self.micro_batch_size}"
          )
  
      self.global_rank = self.grid.get_global_rank()
  
      assert self.dp_world_size == self.grid.data_parallel_size
      assert (
          self.train_batch_size()
          == self.micro_batch_size
          * self.micro_batches
          * self.grid.data_parallel_size
      )
  
      #  Set Stage Inf
      self.num_stages = self.grid.pipe_parallel_size
      self.stage_id = self.grid.get_stage_id()
      self.prev_stage = self.stage_id - 1
      self.next_stage = self.stage_id + 1
  
      self.data_iterator = None
      self.batch_fn = None
  
      self._force_grad_boundary = False
  
      self.batch_timer = ThroughputTimer(
          batch_size=self.train_batch_size(),
          logging_fn=self.tput_log,
          monitor_memory=False,
          steps_per_output=self.steps_per_print(),
      )
  
      # PipelineEngine needs to handle data loading specially due to only the first
      # and last stages loading inputs/labels. We construct a sampler that uses
      if self.training_data:
          self._build_data_iter(self.training_data)
  
      self.is_pipe_parallel = self.grid.pipe_parallel_size > 1
      self.is_data_parallel = self.grid.data_parallel_size > 1
      self.is_model_parallel = self.grid.model_parallel_size > 1
  
      # Partition input/output buffers
      # XXX temporarily disable while I revert some partition hacks.
      assert isinstance(self._config.pipeline["pipe_partitioned"], bool)
      assert isinstance(self._config.pipeline["grad_partitioned"], bool)
      self.is_pipe_partitioned = (
          self.is_model_parallel
          and self._config.pipeline["pipe_partitioned"]
      )
      self.is_grad_partitioned = (
          self.is_model_parallel
          and self._config.pipeline["grad_partitioned"]
      )
      logger.info(
          f"is_pipe_partitioned= {self.is_pipe_partitioned} "
          f"is_grad_partitioned= {self.is_grad_partitioned}"
      )
  
      model_parameters = filter(
          lambda p: p.requires_grad, self.module.parameters()
      )
      num_params = sum([p.numel() for p in model_parameters])
      unique_params = num_params
      # Subtract tied parameters if we don't own them
      if self.module.tied_comms:
          tied_params = 0
          for key, d in self.module.tied_comms.items():
              if self.global_rank != min(d["ranks"]):
                  tied_params += sum(
                      p.numel() for p in d["module"].parameters()
                  )
          unique_params -= tied_params
      params_tensor = torch.LongTensor(
          data=[num_params, unique_params]
      ).to(self.device)
      dist.all_reduce(
          params_tensor, group=self.grid.get_model_parallel_group()
      )
      params_tensor = params_tensor.tolist()
      total_params = params_tensor[0]
      unique_params = params_tensor[1]
      if self.grid.data_parallel_id == 0:
          logger.info(
              f"RANK={self.global_rank} "
              f"STAGE={self.stage_id} "
              f"LAYERS={self.module._local_stop - self.module._local_start} "
              f"[{self.module._local_start}, {self.module._local_stop}) "
              f"STAGE_PARAMS={num_params} ({num_params/1e6:0.3f}M) "
              f"TOTAL_PARAMS={total_params} ({total_params/1e6:0.3f}M) "
              f"UNIQUE_PARAMS={unique_params} ({unique_params/1e6:0.3f}M)"
          )
  
      # initialize peer-2-peer communication and allreduce groups
      if self.is_pipe_parallel:
          p2p.init_process_groups(self.grid)
  
      # Pipeline buffers
      self.num_pipe_buffers = 0
      self.pipe_buffers = {
          "inputs": [],  # batch input and received activations
          "labels": [],  # labels from batch input
          "outputs": [],  # activations
          "output_tensors": [],  # tensor object to preserve backward graph
      }
      self.pipe_recv_buf = None
      self.grad_layer = None
  
      self.meta_buffer = None
  
      self.first_output_send = True
      self.first_gradient_send = True
      self.pipe_partition_input_meta_cache = None
      self.pipe_partition_output_meta_cache = None
      self.pipe_partition_grad_meta_cache = None
      self.grad_partition_grad_layer_meta_cache = None
  
      # stores the loss for the current micro batch being processed
      self.loss = torch.tensor(0.0).to(self.device)
  
      # stores the loss for the entire batch
      self.total_loss = None
      self.agg_loss = torch.tensor(0.0, requires_grad=False).to(
          self.device
      )
      self.dp_group_loss = torch.tensor(0.0, requires_grad=False).to(
          self.device
      )
  
      if self._config.pipeline["activation_checkpoint_interval"] > 0:
          self.module.activation_checkpoint_interval = (
              self._config.pipeline["activation_checkpoint_interval"]
          )
          # set use_reentrant default to True.
          if self._config.pipeline.get("use_reentrant") is None:
              self._config.pipeline["use_reentrant"] = True
          if self._config.pipeline["use_reentrant"] is False:
              # set activation_checkpoint_func to non_reentrant_checkpoint func.
              self.module.activation_checkpoint_func = (
                  ds_checkpointing.non_reentrant_checkpoint
              )
              if self.grid.get_global_rank() == 0:
                  logger.info(
                      f"CONFIG: activation_checkpoint_func=non_reentrant_checkpoint"
                  )
  
      self.module.checkpoint_parallel_write_pipeline = (
          self._config.checkpoint_parallel_write_pipeline
      )
  
      if self.is_last_stage():
          self.loss_model = self.module.loss_fn
  
      self.has_attention_mask = (
          self.module.__class__.__name__ == "GPT2ModelPipe"
      )
      # Initialize pipeline communicators. Just send a 0.
      if is_even(self.stage_id):
          if not self.is_last_stage():
              p2p.send(self.loss, self.next_stage)
          if not self.is_first_stage():
              p2p.recv(self.loss, self.prev_stage)
      else:
          if not self.is_first_stage():
              p2p.recv(self.loss, self.prev_stage)
          if not self.is_last_stage():
              p2p.send(self.loss, self.next_stage)
  
      # XXX look into timer reporting timing
      # Initialize some timers because of early weirdness.
      if self.wall_clock_breakdown():
          self.timers(FORWARD_MICRO_TIMER).start()
          self.timers(FORWARD_MICRO_TIMER).stop()
          self.timers(BACKWARD_MICRO_TIMER).start()
          self.timers(BACKWARD_MICRO_TIMER).stop()
          self.timers(BACKWARD_INNER_MICRO_TIMER).start()
          self.timers(BACKWARD_INNER_MICRO_TIMER).stop()
          self.timers(BACKWARD_REDUCE_MICRO_TIMER).start()
          self.timers(BACKWARD_REDUCE_MICRO_TIMER).stop()
          self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).start()
          self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).stop()
          self.timers(STEP_MICRO_TIMER).start()
          self.timers(STEP_MICRO_TIMER).stop()
  
imports: from types import MethodType
  import torch
  from deepspeed import comm
  from deepspeed.utils import logger
  from deepspeed.utils.timer import ThroughputTimer
  from deepspeed.accelerator import get_accelerator
  from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
  from ..engine import DeepSpeedEngine, MEMORY_OPT_ALLREDUCE_SIZE
  from deepspeed.utils.timer import FORWARD_MICRO_TIMER, FORWARD_GLOBAL_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_MICRO_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, STEP_MICRO_TIMER, STEP_GLOBAL_TIMER
  from ..utils import PartitionedTensor
  from ..dataloader import RepeatingLoader
  from ..zero.config import ZeroStageEnum
  from ..activation_checkpointing import checkpointing
  from .module import PipelineModule, PipelineError
  from .None import p2p
  from .None import schedule
label_desc: Use image node_id 1 to create a new PipelineEngine object from inherited base classes: DeepSpeedEngine with example: obj = PipelineEngine(has_bool_tensors)