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","."]
docstr: Progress the pipeline to train the next batch of data. The engine will ingest
  ``self.train_batch_size()`` total samples collectively across all workers.
  
  
  An iterator that over training data should be provided as an argument
  unless ``deepspeed.initialize()`` was provided a training set. In that event,
  the training data will automatically be read.
  
  
  .. warning::
      A total of ``self.gradient_accumulation_steps()`` entries will be pulled
      from ``data_iter`` by each pipeline. There must be sufficient
      data left in ``data_iter`` or else a ``StopIteration`` will halt training.
  
      DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader`
      that wraps data loaders to automatically restart upon a ``StopIteration``.
  
  Args:
      data_iter (Iterator, optional): Iterator of training data.
  
  Returns:
      The arithmetic mean of the losses computed this batch.
code: def train_batch(self, data_iter=None):
      """Progress the pipeline to train the next batch of data. The engine will ingest
      ``self.train_batch_size()`` total samples collectively across all workers.
  
  
      An iterator that over training data should be provided as an argument
      unless ``deepspeed.initialize()`` was provided a training set. In that event,
      the training data will automatically be read.
  
  
      .. warning::
          A total of ``self.gradient_accumulation_steps()`` entries will be pulled
          from ``data_iter`` by each pipeline. There must be sufficient
          data left in ``data_iter`` or else a ``StopIteration`` will halt training.
  
          DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader`
          that wraps data loaders to automatically restart upon a ``StopIteration``.
  
      Args:
          data_iter (Iterator, optional): Iterator of training data.
  
      Returns:
          The arithmetic mean of the losses computed this batch.
      """
      if not torch._C.is_grad_enabled():
          raise RuntimeError(
              f"train_batch() requires gradients enabled. Use eval_batch() instead."
          )
  
      # Curriculum learning could change activation shape
      if self.curriculum_enabled_legacy():
          new_difficulty = (
              self.curriculum_scheduler_legacy.update_difficulty(
                  self.global_steps + 1
              )
          )
          if (
              self.global_steps == 0
              or self.curriculum_scheduler_legacy.first_step
          ):
              self.reset_activation_shape()
              self.curriculum_scheduler_legacy.first_step = False
          elif (
              new_difficulty
              != self.curriculum_scheduler_legacy.get_difficulty(
                  self.global_steps
              )
          ):
              self.reset_activation_shape()
  
      if data_iter is not None:
          self.set_dataiterator(data_iter)
  
      self.module.train()
      self.total_loss = None
      self._compute_loss = True
  
      # Do the work
      self.timers(TRAIN_BATCH_TIMER).start()
      sched = schedule.TrainSchedule(
          micro_batches=self.micro_batches,
          stages=self.num_stages,
          stage_id=self.stage_id,
      )
      self._exec_schedule(sched)
      self.agg_train_loss = self._aggregate_total_loss()
  
      self.timers(TRAIN_BATCH_TIMER).stop()
  
      if self.global_steps % self.steps_per_print() == 0:
          if self.global_rank == 0:
              elapsed = (
                  self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True)
                  / 1000.0
              )
              iter_time = elapsed / self.steps_per_print()
              tput = self.train_batch_size() / iter_time
              print(
                  f"steps: {self.global_steps} "
                  f"loss: {self.agg_train_loss:0.4f} "
                  f"iter time (s): {iter_time:0.3f} "
                  f"samples/sec: {tput:0.3f}"
              )
          else:
              self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True)
  
      # Monitoring
      if self.global_rank == 0 and self.monitor.enabled:
          self.summary_events = [
              (
                  f"Train/Samples/train_loss",
                  self.agg_train_loss.mean().item(),
                  self.global_samples,
              )
          ]
          self.monitor.write_events(self.summary_events)
  
      if (
          self.wall_clock_breakdown()
          and self.global_steps % self.steps_per_print() == 0
      ):
          self.timers.log(
              [
                  PIPE_SEND_OUTPUT_TIMER,
                  PIPE_SEND_GRAD_TIMER,
                  PIPE_RECV_INPUT_TIMER,
                  PIPE_RECV_GRAD_TIMER,
              ]
          )
  
      # TODO: should return precisely what loss returned and allow others to be queried?
      return self.agg_train_loss
  
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 9 for calling the PipelineEngine obj's underlying member method code with example usage: obj.train_batch(data_iter) and returns: self