docstr: DeepSpeed allows specifying train_batch_size, train_micro_batch_size_per_gpu, and
  gradient_accumulation_steps. The three are related as follows:
  train_batch_size = train_micro_batch_size * gradient_accumulation_steps.
  Hence, if two are specified, the third can be inferred.
  
  For pipeline parallel training, DeepSpeed will automatically interleave
  gradient_accumulation_steps worth of micro batches in one train_batch/eval_batch call.
  
  With the default DeepSpeed model engine (no pipeline parallel training), the backward
  and optimizer step calls track micro batches and will automatically update model weights
  and lr scheduler if micro batches % gradient_accumulation_steps == 0.
  
  Comparing training with and without pipeline parallel is a common goal.  Since DeepSpeed's
  PipelineEngine trains on a number of micro batches equal to gradient accumulation steps,
  we automatically perform gradient accumulation by default when pipeline parallelism is not
  enabled.  This makes it fair to compare training with and without pipeline parallelism
  at a given batch idx. This can be turned off by setting
  context.disable_auto_grad_accumulation.
code: def _train_for_step(
      self, step_id: int, num_batches: int, total_batches_processed: int
  ) -> workload.Metrics:
      """
      DeepSpeed allows specifying train_batch_size, train_micro_batch_size_per_gpu, and
      gradient_accumulation_steps. The three are related as follows:
      train_batch_size = train_micro_batch_size * gradient_accumulation_steps.
      Hence, if two are specified, the third can be inferred.
  
      For pipeline parallel training, DeepSpeed will automatically interleave
      gradient_accumulation_steps worth of micro batches in one train_batch/eval_batch call.
  
      With the default DeepSpeed model engine (no pipeline parallel training), the backward
      and optimizer step calls track micro batches and will automatically update model weights
      and lr scheduler if micro batches % gradient_accumulation_steps == 0.
  
      Comparing training with and without pipeline parallel is a common goal.  Since DeepSpeed's
      PipelineEngine trains on a number of micro batches equal to gradient accumulation steps,
      we automatically perform gradient accumulation by default when pipeline parallelism is not
      enabled.  This makes it fair to compare training with and without pipeline parallelism
      at a given batch idx. This can be turned off by setting
      context.disable_auto_grad_accumulation.
      """
      self.prof.set_training(True)
      assert step_id > 0, "step_id should be greater than 0"
      step_start_time = time.time()
      self.context.reset_reducers()
  
      # Set the behavior of certain layers (e.g., dropout) that are different
      # between training and inference.
      for model in self.context.models:
          model.train()
  
      start = total_batches_processed
      end = start + num_batches
  
      per_batch_metrics = []  # type: List[Dict]
      num_inputs = 0
  
      for batch_idx in range(start, end):
          self.steps_completed += 1
          self.prof.update_batch_idx(batch_idx)
          batch_start_time = time.time()
          self.context._current_batch_idx = batch_idx
          if self.context.is_epoch_start():
              for callback in self.callbacks.values():
                  with self.prof.record_timing(
                      f"callbacks.{callback.__class__.__name__}.on_training_epoch_start"
                  ):
                      callback.on_training_epoch_start(
                          self.get_epoch_idx(batch_idx)
                      )
          # This can be inaccurate if the user's data loader does not return batches with
          # the micro batch size.  It is also slightly inaccurate if the data loader can return
          # partial batches.  The same sort of assumptions is made in the DeepSpeed
          # model engine's accounting and profiling computations.
          batch_inputs = (
              self.context.train_micro_batch_size_per_gpu
              * self.context.num_micro_batches_per_slot
          )
          num_inputs += batch_inputs
          num_train_batch_calls = (
              self.context.num_micro_batches_per_slot
          )
          if (
              self.context.use_pipeline_parallel
              or self.context._manual_grad_accumulation
          ):
              num_train_batch_calls = 1
          self.context._loss_ids = {}
          for _ in range(num_train_batch_calls):
              with contextlib.ExitStack() as exit_stack:
                  exit_stack.enter_context(
                      self.prof.record_timing(
                          "train_batch",
                          requires_sync=False,
                          accumulate=True,
                      )
                  )
                  if self.context.profiler:
                      exit_stack.enter_context(self.context.profiler)
  
                  tr_metrics = self.trial.train_batch(
                      self.training_iterator,
                      self.get_epoch_idx(batch_idx),
                      batch_idx,
                  )
  
                  if self.context.profiler:
                      self.context.profiler.step()
  
              if self.context._mpu.should_report_metrics:
                  if isinstance(tr_metrics, torch.Tensor):
                      tr_metrics = {"loss": tr_metrics}
                  if not isinstance(tr_metrics, dict):
                      raise det.errors.InvalidExperimentException(
                          "train_batch must return a dictionary "
                          f"mapping string names to Tensor metrics, got {type(tr_metrics)}",
                      )
  
                  for name, metric in tr_metrics.items():
                      # Convert PyTorch metric values to NumPy, so that
                      # `det.util.encode_json` handles them properly without
                      # needing a dependency on PyTorch.
                      if isinstance(metric, torch.Tensor):
                          metric = metric.cpu().detach().numpy()
                      tr_metrics[name] = metric
                  per_batch_metrics.append(tr_metrics)
          # We do a check here to make sure that we do indeed process `num_micro_batches_per_slot`
          # micro batches when training a batch for models that do not use pipeline parallelism.
          model0 = self.context.models[0]
          if not isinstance(model0, deepspeed.PipelineEngine):
              assert (
                  model0.micro_steps
                  % self.context.num_micro_batches_per_slot
                  == 0
              ), "did not train for gradient accumulation steps"
  
          batch_dur = time.time() - batch_start_time
          samples_per_second = batch_inputs / batch_dur
          samples_per_second *= (
              self.context._mpu.data_parallel_world_size
          )
          self.prof.record_metric(
              "samples_per_second", samples_per_second
          )
  
          if self.context.is_epoch_end():
              for callback in self.callbacks.values():
                  with self.prof.record_timing(
                      f"callbacks.{callback.__class__.__name__}.on_training_epoch_end"
                  ):
                      callback.on_training_epoch_end(
                          self.get_epoch_idx(batch_idx)
                      )
  
      # Aggregate and reduce training metrics from all the training processes.
      if (
          self.context.distributed.size > 1
          and self.context._average_training_metrics
      ):
          with self.prof.record_timing("average_training_metrics"):
              per_batch_metrics = (
                  pytorch._combine_and_average_training_metrics(
                      self.context.distributed, per_batch_metrics
                  )
              )
      num_inputs *= self.context._mpu.data_parallel_world_size
      metrics = det.util.make_metrics(num_inputs, per_batch_metrics)
  
      # Ignore batch_metrics entirely for custom reducers; there's no guarantee that per-batch
      # metrics are even logical for a custom reducer.
      with self.prof.record_timing("reduce_metrics"):
          metrics["avg_metrics"].update(
              pytorch._convert_metrics_to_numpy(
                  self.context.reduce_metrics(for_training=True)
              )
          )
  
      if self.is_chief:
          step_duration = time.time() - step_start_time
          logger.info(
              det.util.make_timing_log(
                  "trained", step_duration, num_inputs, num_batches
              )
          )
          self.prof.set_training(False)
  
          if self.context.get_enable_tensorboard_logging():
              det.pytorch._log_tb_metrics(
                  self.context.get_tensorboard_writer(),
                  "train",
                  self.steps_completed,
                  metrics["avg_metrics"],
                  metrics["batch_metrics"],
              )
  
      if not self.is_chief:
          return {}
  
      return metrics
  
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 8 for calling the DeepSpeedTrialController obj's underlying member method code with example usage: obj._train_for_step(step_id, num_batches, total_batches_processed) and returns: metrics, dict