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: Initializes a trial using the provided ``context``. The general steps are:
  
  #. Initialize the model(s) and, optionally, the optimizer and lr_scheduler.  The latter
     two can also be configured using the DeepSpeed config.
  #. Build the DeepSpeed model engine by calling ``deepspeed.initialize`` with the model
     (optionally optimizer and lr scheduler) and a DeepSpeed config.  Wrap it with
     ``context.wrap_model_engine``.
  #. If you want, use a custom model parallel unit by calling ``context.set_mpu``.
  #. If you want, disable automatic gradient accumulation by calling
     ``context.disable_auto_grad_accumulation``.
  #. If you want, use a custom data loader by calling
     ``context.disable_dataset_reproducibility_checks``.
  
  Here is a code example.
  
  .. code-block:: python
  
      self.context = context
      self.args = AttrDict(self.context.get_hparams())
  
      # Build deepspeed model engine.
      model = ... # build model
      model_engine, optimizer, lr_scheduler, _ = deepspeed.initialize(
          args=self.args,
          model=model,
      )
  
      self.model_engine = self.context.wrap_model_engine(model_engine)
code: def __init__(self, context: det_ds.DeepSpeedTrialContext) -> None:
      """
      Initializes a trial using the provided ``context``. The general steps are:
  
      #. Initialize the model(s) and, optionally, the optimizer and lr_scheduler.  The latter
         two can also be configured using the DeepSpeed config.
      #. Build the DeepSpeed model engine by calling ``deepspeed.initialize`` with the model
         (optionally optimizer and lr scheduler) and a DeepSpeed config.  Wrap it with
         ``context.wrap_model_engine``.
      #. If you want, use a custom model parallel unit by calling ``context.set_mpu``.
      #. If you want, disable automatic gradient accumulation by calling
         ``context.disable_auto_grad_accumulation``.
      #. If you want, use a custom data loader by calling
         ``context.disable_dataset_reproducibility_checks``.
  
      Here is a code example.
  
      .. code-block:: python
  
          self.context = context
          self.args = AttrDict(self.context.get_hparams())
  
          # Build deepspeed model engine.
          model = ... # build model
          model_engine, optimizer, lr_scheduler, _ = deepspeed.initialize(
              args=self.args,
              model=model,
          )
  
          self.model_engine = self.context.wrap_model_engine(model_engine)
      """
      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 1 to create a new DeepSpeedTrial object from inherited base classes: det with example: obj = DeepSpeedTrial(context)