class_docstr: A high-level interface to DecentralizedAverager that averages trainable params or gradients for an optimizer.
  
  This averager implements a number of typical use cases that arise in collaborative optimization
  - averaging parameters or gradients or both (in future, this will support averaging learning rates as well)
  - this peer's weight (e.g. based on its batch size) can be specified via averager.step(weight=...)
  - when out of sync, the averager will load the entire optimizer state from an up-to-date peer
  
  :param opt: a pytorch optimizer to be averaged between peers (complete with model parameters)
  :param average_parameters: whether or not to average model parameters in self.step(...)
  :param average_gradients: whether or not to average model gradients in self.step(...)
  :param average_opt_statistics: if specified, average optimizer statistics with corresponding names in statedict
  :param initialize_optimizer: if True, this will run a speculative optimizer step with
    zero gradients to initialize all tensors. If False, please initialize the optimizer state manually.
  :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 for averaging tensors that are updated outside of opt.step (e.g. batchnorm stats)
  :param kwargs: any additional parameters will be forwarded to DecentralizedAverager
class_docstr_tok: ["A","high-level","interface","to","DecentralizedAverager","that","averages","trainable","params","or","gradients","for","an","optimizer",".","This","averager","implements","a","number","of","typical","use","cases","that","arise","in","collaborative","optimization","-","averaging","parameters","or","gradients","or","both","(","in","future",",","this","will","support","averaging","learning","rates","as","well",")","-","this","peer","'s","weight","(","e.g",".","based","on","its","batch","size",")","can","be","specified","via","averager.step","(","weight=","...",")","-","when","out","of","sync",",","the","averager","will","load","the","entire","optimizer","state","from","an","up-to-date","peer",":","param","opt",":","a","pytorch","optimizer","to","be","averaged","between","peers","(","complete","with","model","parameters",")",":","param","average_parameters",":","whether","or","not","to","average","model","parameters","in","self.step","(","...",")",":","param","average_gradients",":","whether","or","not","to","average","model","gradients","in","self.step","(","...",")",":","param","average_opt_statistics",":","if","specified",",","average","optimizer","statistics","with","corresponding","names","in","statedict",":","param","initialize_optimizer",":","if","True",",","this","will","run","a","speculative","optimizer","step","with","zero","gradients","to","initialize","all","tensors",".","If","False",",","please","initialize","the","optimizer","state","manually",".",":","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","for","averaging","tensors","that","are","updated","outside","of","opt.step","(","e.g",".","batchnorm","stats",")",":","param","kwargs",":","any","additional","parameters","will","be","forwarded","to","DecentralizedAverager"]
docstr: Average optimizer weights and gradients with peers.
  
  :param data_lock: averager locks it when model parameters are modified. Otherwise it's assumed that no model
  modifications occur during averaging step
docstr_tok: ["Average","optimizer","weights","and","gradients","with","peers",".",":","param","data_lock",":","averager","locks","it","when","model","parameters","are","modified",".","Otherwise","it","'s","assumed","that","no","model","modifications","occur","during","averaging","step"]
code: def step(
      self,
      data_lock: Optional[Lock] = None,
      wait: bool = True,
      **kwargs,
  ):
      """
      Average optimizer weights and gradients with peers.
  
      :param data_lock: averager locks it when model parameters are modified. Otherwise it's assumed that no model
      modifications occur during averaging step
      """
      if not wait:
          return self.step_executor.submit(
              self.step, data_lock, wait=True, **kwargs
          )
  
      # if data_lock is supplied, tensors might change during averaging, so we need to copy them
      use_old_local_tensors = data_lock is not None
      if data_lock is None:
          data_lock = nullcontext()
  
      local_tensors = list(self.local_tensors())
      with self.lock_averager_step, torch.no_grad():
          # fill averager's tensors with current local tensors
          self.pending_updates_done.clear()
          with data_lock, self.get_tensors() as averaged_tensors:
              if use_old_local_tensors:
                  old_local_tensors = tuple(
                      x.cpu().float().clone() for x in local_tensors
                  )
              assert len(local_tensors) == len(
                  averaged_tensors
              ), "The number of optimized parameters should not change."
              for averaged_tensor, local_tensor in zip(
                  averaged_tensors, local_tensors
              ):
                  averaged_tensor[...] = local_tensor.cpu().float()
          self.pending_updates_done.set()
  
          # find a group and hopefully average tensors with peers, use batch sizes as weights
          gathered = super().step(**kwargs)
          if gathered is not None:
              # load averaged tensors back into model
              self.pending_updates_done.clear()
              with data_lock, self.get_tensors() as averaged_tensors:
                  if len(averaged_tensors) != len(local_tensors):
                      raise RuntimeError(
                          "The number of optimized parameters should not change"
                      )
  
                  if use_old_local_tensors:
                      # since tensors might have changed, we subtract old_local_tensor and add averaged. This prevents
                      # losing local updates that might have occurred during averaging
                      for (
                          averaged_tensor,
                          local_tensor,
                          old_local_tensor,
                      ) in zip(
                          averaged_tensors,
                          local_tensors,
                          old_local_tensors,
                      ):
                          averaged_tensor = averaged_tensor.to(
                              dtype=local_tensor.dtype,
                              device=local_tensor.device,
                              non_blocking=True,
                          )
                          old_local_tensor = old_local_tensor.to(
                              dtype=local_tensor.dtype,
                              device=local_tensor.device,
                              non_blocking=True,
                          )
  
                          local_tensor.add_(
                              averaged_tensor - old_local_tensor
                          )
                  else:
                      for averaged_tensor, local_tensor in zip(
                          averaged_tensors, local_tensors
                      ):
                          local_tensor.copy_(
                              averaged_tensor, non_blocking=True
                          )
              self.pending_updates_done.set()
  
          self.local_step += 1
          return gathered
  
imports: from concurrent.futures import ThreadPoolExecutor
  from contextlib import nullcontext
  from itertools import chain
  from threading import Event, Lock
  from typing import Dict, Iterator, Optional, Sequence
  import torch
  from hivemind.averaging import DecentralizedAverager
  from hivemind.compression import CompressionInfo, TensorRole
  from hivemind.utils import get_logger, nested_flatten, nested_pack
label_desc: Use image node_id 2 for calling the TrainingAverager obj's underlying member method code with example usage: obj.step(data_lock, wait) and returns: self, gathered