class_docstr: DeepSpeed trials are created by subclassing this abstract class.
  
  We can do the following things in this trial class:
  
  * **Define the DeepSpeed model engine which includes the model, optimizer, and lr_scheduler**.
  
     In the :meth:`__init__` method, initialize models and, optionally, optimizers and
     LR schedulers and pass them to ``deepspeed.initialize`` to build the model engine.  Then
     pass the created model engine to ``wrap_model_engine`` provided by
     :class:`~determined.pytorch.deepspeed.DeepSpeedTrialContext`.
     We support multiple DeepSpeed model engines if they only use data parallelism or if
     they use the same model parallel unit.
  
  * **Run forward and backward passes**.
  
     In :meth:`train_batch`, use the methods provided by the DeepSpeed model engine to perform
     the backward pass and optimizer step.  These methods will differ depending on whether
     you are using pipeline parallelism or not.
class_docstr_tok: ["DeepSpeed","trials","are","created","by","subclassing","this","abstract","class",".","We","can","do","the","following","things","in","this","trial","class",":","*","*","*","Define","the","DeepSpeed","model","engine","which","includes","the","model",",","optimizer",",","and","lr_scheduler","*","*",".","In","the",":","meth",":","`","__init__","`","method",",","initialize","models","and",",","optionally",",","optimizers","and","LR","schedulers","and","pass","them","to","``","deepspeed.initialize","``","to","build","the","model","engine",".","Then","pass","the","created","model","engine","to","``","wrap_model_engine","``","provided","by",":","class",":","`","~determined.pytorch.deepspeed.DeepSpeedTrialContext","`",".","We","support","multiple","DeepSpeed","model","engines","if","they","only","use","data","parallelism","or","if","they","use","the","same","model","parallel","unit",".","*","*","*","Run","forward","and","backward","passes","*","*",".","In",":","meth",":","`","train_batch","`",",","use","the","methods","provided","by","the","DeepSpeed","model","engine","to","perform","the","backward","pass","and","optimizer","step",".","These","methods","will","differ","depending","on","whether","you","are","using","pipeline","parallelism","or","not","."]
docstr: Train one full batch (i.e. train on ``train_batch_size`` samples, perhaps consisting
  of multiple micro-batches).
  
  If training without pipeline parallelism, users should implement this function by doing
  the following things:
  
  #. Get a batch from the ``dataloader_iter`` and pass it to the GPU.
  #. Compute the loss in the forward pass.
  #. Perform the backward pass.
  #. Perform an optimizer step.
  #. Return training metrics in a dictionary.
  
  Here is a code example.
  
  .. code-block:: python
  
      # Assume one model_engine wrapped in ``__init__``.
  
      batch = self.context.to_device(next(dataloader_iter))
      loss = self.model_engine(batch)
      self.model_engine.backward(loss)
      self.model_engine.step()
      return {"loss": loss}
  
  If using gradient accumulation over multiple micro-batches, Determined will automatically
  call ``train_batch`` multiple times according to ``gradient_accumulation_steps`` in the
  DeepSpeed config.
  
  With pipeline parallelism there is no need to manually get a batch from the
  ``dataloader_iter`` and the forward, backward, optimizer steps are combined in the
  model engine's ``train_batch`` method.
  
  .. code-block:: python
  
      # Assume one model_engine wrapped in ``__init__``.
  
      loss = self.model_engine.train_batch(dataloader_iter)
      return {"loss": loss}
  
  Arguments:
      dataloader_iter (Iterator[torch.utils.data.DataLoader], optional): iterator over
          the train DataLoader.
      epoch_idx (integer): index of the current epoch among all the batches processed
          per device (slot) since the start of training.
      batch_idx (integer): index of the current batch among all the epochs processed
          per device (slot) since the start of training.
  Returns:
      torch.Tensor or Dict[str, Any]:
          training metrics to return.
code: def train_batch(
      self,
      dataloader_iter: Optional[Iterator[pytorch.TorchData]],
      epoch_idx: int,
      batch_idx: int,
  ) -> Union[torch.Tensor, Dict[str, Any]]:
      """
      Train one full batch (i.e. train on ``train_batch_size`` samples, perhaps consisting
      of multiple micro-batches).
  
      If training without pipeline parallelism, users should implement this function by doing
      the following things:
  
      #. Get a batch from the ``dataloader_iter`` and pass it to the GPU.
      #. Compute the loss in the forward pass.
      #. Perform the backward pass.
      #. Perform an optimizer step.
      #. Return training metrics in a dictionary.
  
      Here is a code example.
  
      .. code-block:: python
  
          # Assume one model_engine wrapped in ``__init__``.
  
          batch = self.context.to_device(next(dataloader_iter))
          loss = self.model_engine(batch)
          self.model_engine.backward(loss)
          self.model_engine.step()
          return {"loss": loss}
  
      If using gradient accumulation over multiple micro-batches, Determined will automatically
      call ``train_batch`` multiple times according to ``gradient_accumulation_steps`` in the
      DeepSpeed config.
  
      With pipeline parallelism there is no need to manually get a batch from the
      ``dataloader_iter`` and the forward, backward, optimizer steps are combined in the
      model engine's ``train_batch`` method.
  
      .. code-block:: python
  
          # Assume one model_engine wrapped in ``__init__``.
  
          loss = self.model_engine.train_batch(dataloader_iter)
          return {"loss": loss}
  
      Arguments:
          dataloader_iter (Iterator[torch.utils.data.DataLoader], optional): iterator over
              the train DataLoader.
          epoch_idx (integer): index of the current epoch among all the batches processed
              per device (slot) since the start of training.
          batch_idx (integer): index of the current batch among all the epochs processed
              per device (slot) since the start of training.
      Returns:
          torch.Tensor or Dict[str, Any]:
              training metrics to return.
      """
      pass
  
imports: import abc
  import contextlib
  import logging
  import os
  import pathlib
  import pickle
  import random
  import time
  import warnings
  from typing import Any, Callable, Dict, Iterator, List, Optional, Type, Union, cast
  import deepspeed
  import numpy
  import torch
  from deepspeed.runtime import dataloader
  import determined
  from determined import layers, pytorch, util, workload
  from determined.pytorch import deepspeed
  from determined.pytorch import dsat
label_desc: Use image node_id 2 for calling the DeepSpeedTrial obj's underlying member method code with example usage: obj.train_batch(dataloader_iter, epoch_idx, batch_idx) without return types