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: Copy parameters from offloaded optimizer to the main model
code: def _apply_optimizer_parameters_(self):
      """Copy parameters from offloaded optimizer to the main model"""
      assert (
          self.offload_optimizer
      ), "Applying offloaded optimizer updates requires offloaded optimizer"
      offloaded_parameters = [
          param
          for group in self.optimizer.param_groups
          for param in group["params"]
      ]
      assert len(offloaded_parameters) == len(
          self.main_parameters
      ), "Optimizer parameters changed during training"
      for main_param, offloaded_param in zip(
          self.main_parameters, offloaded_parameters
      ):
          main_param.copy_(offloaded_param, non_blocking=True)
  
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 13 for calling the TrainingStateAverager obj's underlying member method code with example usage: obj._apply_optimizer_parameters_() without return types