class_docstr: An auxiliary class that holds peer's training state, including model parameters, optimizer statistics, scheduler
  and any other variables that define the local training state (e.g. batchnorm moving averages).
  TrainingStateAveraager is intended to keep these parameters weakly synchronized across the swarm.
  
  The intended use is to call .step(optimizer_step=..., averaging_round=...) periodically, e.g. after every batch.
  If peer gets out of sync with the swarm, one should call state_averager.load_state_from_peers() to re-synchronize.
  
  Example:
  
  >>> avgr = TrainingStateAverager(optimizer=torch.optim.Adam, params=model.parameters(), ...)
  >>> # alternative interface: TrainingStateAverager(optimizer=torch.optim.Adam(model.parameters()), ...)
  >>> avgr.load_state_from_peers()
  >>> for i, batch in enumerate(training_dataloader):
  >>>     loss = compute_loss(model, batch)
  >>>     loss.backward()
  >>>     avgr.step(optimizer_step=i % 10 == 0, averaging_round=is_it_time_for_averaging(), delay_averaging=True)
  
  :note: when using delay_averaging or delay_optimizer_step, calling optimizer directly is not recommended because
    it may overlap with delayed updates from a background thread with unpredictable results. Instead, please call
    TrainingStateAverager.step(..., optimizer_step=True)
  
  :param optimizer: PyTorch Optimizer or a callable that creates a optimizer from param groups
  :param params: optional, a list/tuple of parameters or structured param groups for the optimizer
  :param scheduler: optional learning rate scheduler or callable that creates one from optimizer instance
  :note: if provided, scheduler will be updated based on averager.local_epoch, not the number of step cycles
  :param initialize_optimizer: if True, run a speculative optimizer step with zero gradients to initialize all
    state tensors. If False, user must make sure that all tensors are pre-initialized at init.
    By default, initialize optimizer unless it already has some state tensors to begin with.
  :param offload_optimizer: if True, create optimizer on top of averaged parameters which may save device memory.
  :param custom_gradients: if True, do *not* automatically load local gradients into the offloaded optimizer.
    This assumes that offloaded gradients will be populated externally, e.g. by the user or by hivemind.Optimizer.
  :param reuse_tensors: if True, reuse parameters and optimizer statistics as averaged_tensors for allreduce.
    For this to work, all parameters must be on CPU and have the appropriate dtype for use in DecentralizedAverager
    Defaults to True if offload_optimizer, False otherwise.
  :param delta_rule_averaging: if True, averaging will use delta rule to allow running local optimizer steps
    while averaging. Delta rule: `state_tensor := state_tensor + averaging_result - state_tensor_before_averaging`
  :param sync_epoch_when_averaging: if True, update local epoch to the latest epoch among averaging peers
  :param parameter_names: optionally provide parameter names in the same order as in params
  :param average_opt_statistics: names of optimizer statistics from state dict that should be averaged with peers
  :param extra_tensors: if specified, these extra tensors will also be averaged and shared in load_state_from_peers.
  :note: you can use extra_tensors to for any tensors not used by the optimizer (e.g. batchnorm statistics)
  :param kwargs: any additional parameters will be forwarded to DecentralizedAverager
docstr: Perform one or several possible actions, depending on the specified keyword args.
  The actions will be performed in the same order as specified below:
  
  :param wait_for_delayed_updates: if there are background averaging rounds, wait for them to finish
    by default, await delayed updates when scheduling the next optimizer step, otherwise do not update
  :param apply_delayed_updates: apply any averaging rounds that have finished but were not applied yet
  :param increment_epoch: increment .local_epoch and update the learning rate scheduler (if present)
  :note: if specified, it is guaranteed that epoch is incremented immediately regardless of other options
  :param optimizer_step: perform a single optimizer step and update local parameters (without changing scheduler)
  :param zero_grad: if True, reset local gradients after performing optimizer step
  :param delay_optimizer_step: if True, run optimizer step in background and apply results in a future step
  :param averaging_round: average parameters, chosen optimizer keys and extra tensors with a group of peers
  :param delay_averaging: if True, perform averaging in background and apply results in a future step
    by default, delay averaging if the optimizer step is also delayed. Set to true to delay only this phase.
  :param averaging_control: if specified, use this as a pre-scheduled averaging round. Should require_trigger.
  :param wait_for_trigger: wait for this (non-asyncio) function to finish before running optimizer step
  :note: if wait_for_trigger fails with any exception, it will abort optimizer step, zero grad and averaging
  :param grad_scaler: when using hivemind.GradScaler, one must forward it to step after calling .unscale_
  :param set_to_none: if True, zero_grad sets local gradients to None instead of zero tensors
    (default in PyTorch 2.0+)
  :param averaging_opts: a dict of keyword arguments forwarded into averaging round
code: def step(
      self,
      wait_for_delayed_updates: bool = None,
      apply_delayed_updates: bool = True,
      increment_epoch: bool = False,
      optimizer_step: bool = False,
      zero_grad: bool = False,
      delay_optimizer_step: bool = False,
      averaging_round: bool = False,
      delay_averaging: Optional[bool] = None,
      averaging_control: Optional[StepControl] = None,
      wait_for_trigger: Optional[Callable[[], Any]] = None,
      grad_scaler: Optional[GradScaler] = None,
      set_to_none: bool = ZERO_GRAD_SET_TO_NONE_DEFAULT,
      averaging_opts: Optional[Dict[str, Any]] = None,
  ):
      """
      Perform one or several possible actions, depending on the specified keyword args.
      The actions will be performed in the same order as specified below:
  
      :param wait_for_delayed_updates: if there are background averaging rounds, wait for them to finish
        by default, await delayed updates when scheduling the next optimizer step, otherwise do not update
      :param apply_delayed_updates: apply any averaging rounds that have finished but were not applied yet
      :param increment_epoch: increment .local_epoch and update the learning rate scheduler (if present)
      :note: if specified, it is guaranteed that epoch is incremented immediately regardless of other options
      :param optimizer_step: perform a single optimizer step and update local parameters (without changing scheduler)
      :param zero_grad: if True, reset local gradients after performing optimizer step
      :param delay_optimizer_step: if True, run optimizer step in background and apply results in a future step
      :param averaging_round: average parameters, chosen optimizer keys and extra tensors with a group of peers
      :param delay_averaging: if True, perform averaging in background and apply results in a future step
        by default, delay averaging if the optimizer step is also delayed. Set to true to delay only this phase.
      :param averaging_control: if specified, use this as a pre-scheduled averaging round. Should require_trigger.
      :param wait_for_trigger: wait for this (non-asyncio) function to finish before running optimizer step
      :note: if wait_for_trigger fails with any exception, it will abort optimizer step, zero grad and averaging
      :param grad_scaler: when using hivemind.GradScaler, one must forward it to step after calling .unscale_
      :param set_to_none: if True, zero_grad sets local gradients to None instead of zero tensors
        (default in PyTorch 2.0+)
      :param averaging_opts: a dict of keyword arguments forwarded into averaging round
      """
      if delay_averaging is None:
          delay_averaging = delay_optimizer_step
      should_wait = (
          averaging_round or optimizer_step or zero_grad
          if self.delta_rule_averaging
          else averaging_round
      )
      if wait_for_delayed_updates is None:
          wait_for_delayed_updates = should_wait
      if should_wait and not (
          wait_for_delayed_updates and apply_delayed_updates
      ):
          raise ValueError(
              "Should wait for background operation to finish before scheduling new one"
          )
      assert (
          not delay_optimizer_step or delay_averaging
      ), "Delayed optimizer step requires delayed averaging"
      if delay_optimizer_step:
          assert (
              self.offload_optimizer
          ), "Delayed optimizer step is only available with offload_optimizer"
          assert (
              not averaging_round or delay_averaging
          ), "Averaging after delayed optimizer should also be delayed"
      if averaging_opts and not averaging_round:
          logger.warning(
              f"Averaging parameters not used because averaging_round=False: {averaging_opts}"
          )
      if averaging_control is not None:
          assert (
              averaging_round
          ), "averaging_control is unused if averaging_round is not performed"
      if wait_for_trigger is not None:
          assert (
              optimizer_step or zero_grad or averaging_round
          ), "trigger is only used for updating parameters"
          if not (self.reuse_tensors or self.custom_gradients):
              # averager was asked to wait_for_trigger in background, but it is not clear which version of gradients
              # should be used for optimizer step (e.g. the gradients that were present during the call to .step or
              # the possibly different gradients when wait_for_trigger has finished).
              raise ValueError(
                  "wait_for_trigger is a low-level option that requires manual gradient manipulation. "
                  "If you know what you're doing, please refer to the comments in the source code for details"
              )
      output = None
  
      if wait_for_delayed_updates:
          for pending_update in self.pending_updates:
              try:
                  timeout = (averaging_opts or {}).get(
                      "averaging_timeout", self._allreduce_timeout
                  )
                  logger.log(
                      self.status_loglevel,
                      "Waiting for delayed updates to finish...",
                  )
                  output = pending_update.result(timeout)
              except BaseException:
                  # exception will be reported below
                  if not pending_update.done():
                      pending_update.cancel()
  
      # remove finished updates, log any exceptions
      finished_updates = {
          pending_update
          for pending_update in self.pending_updates
          if pending_update.done()
      }
      self.pending_updates = {
          pending_update
          for pending_update in self.pending_updates
          if not pending_update.done()
      }
      for finished_update in finished_updates:
          if finished_update.cancelled() or finished_update.exception():
              logger.log(
                  self.status_loglevel,
                  f"Background update failed: {finished_update}",
              )
  
      if apply_delayed_updates:
          if self.finished_averaging_round.is_set():
              if not self.reuse_tensors:
                  self._apply_averaging_results_()
              if (
                  self.offload_optimizer
                  and not self.finished_optimizer_step.is_set()
              ):
                  self._apply_optimizer_parameters_()
              logger.log(
                  self.status_loglevel,
                  "Received parameters from background averaging round",
              )
              self.finished_averaging_round.clear()
  
          if self.finished_optimizer_step.is_set():
              if self.offload_optimizer:
                  self._apply_optimizer_parameters_()
              logger.debug(
                  "Received parameters from background optimizer step"
              )
              self.finished_optimizer_step.clear()
  
      if increment_epoch:
          self.local_epoch += 1
  
      if optimizer_step or zero_grad or averaging_round:
          if self.offload_optimizer and not self.custom_gradients:
              self._load_local_grads_into_optimizer_()
  
          pending_update = self.step_executor.submit(
              self._do,
              wait_for_trigger,
              optimizer_step,
              zero_grad,
              averaging_round,
              averaging_control,
              grad_scaler,
              set_to_none,
              **averaging_opts or {},
          )
          self.pending_updates.add(pending_update)
  
          should_await_optimizer = (
              optimizer_step or zero_grad
          ) and not delay_optimizer_step
          should_await_averaging = (
              averaging_round and not delay_averaging
          )
  
          if should_await_optimizer:
              self.finished_optimizer_step.wait()
              self.finished_optimizer_step.clear()
              if self.offload_optimizer and not should_await_averaging:
                  self._apply_optimizer_parameters_()
              logger.debug("Finished optimizer step")
  
          if should_await_averaging:
              self.finished_averaging_round.wait()
              self.finished_averaging_round.clear()
              if not self.reuse_tensors:
                  self._apply_averaging_results_()
              if self.offload_optimizer:
                  self._apply_optimizer_parameters_()
              logger.log(
                  self.status_loglevel, "Finished averaging round"
              )
  
          async_averaging = averaging_round and delay_averaging
          async_optimizer = (
              optimizer_step or zero_grad
          ) and delay_optimizer_step
  
          if not (async_averaging or async_optimizer):
              try:
                  output = pending_update.result()
              finally:
                  self.pending_updates.remove(pending_update)
  
      return output
  
imports: import logging
  import threading
  import time
  from concurrent.futures import ThreadPoolExecutor
  from contextlib import nullcontext
  from itertools import chain
  from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Sequence, Tuple, Union
  import torch
  from packaging.version import Version
  import hivemind
  from hivemind.averaging import DecentralizedAverager
  from hivemind.averaging.control import StepControl
  from hivemind.compression import CompressionInfo, TensorRole
  from hivemind.optim.grad_scaler import GradScaler
  from hivemind.utils import DHTExpiration, PerformanceEMA, get_logger, nested_flatten, nested_pack
label_desc: Use image node_id 10 for calling the TrainingStateAverager obj's underlying member method code with example usage: obj.step(wait_for_delayed_updates, apply_delayed_updates, increment_epoch, optimizer_step, zero_grad, delay_optimizer_step, averaging_round, delay_averaging, averaging_control, wait_for_trigger, grad_scaler, set_to_none, averaging_opts) and returns: output