class_docstr: An auxiliary averaging class that is responsible for accumulating gradients and aggregating them with peers.
  GradientAverager is meant to be used within hivemind.Optimizer, but it can be used standalone (see example below).
  
  GradientAverager manages three sets of buffers:
  (1) model gradients - the gradients associated with local model parameters by PyTorch (param.grad).
      These tensors are typically stored on device and updated by torch autograd
  (2) gradient accumulators - an [optional] set of buffers where local gradients are accumulated.
    - note: if reuse_grad_buffers is True, the averager will use gradients from parameters as local accumulators,
      which reduces RAM usage but requires the user to avoid calling zero_grad / clip_grad manually
  (3) averaged gradients - gradient buffers that are aggregated in-place with peers, always in host memory
  
  :param parameters: pytorch parameters for which to aggregate gradients
  :param dht: a DHT instance connected to the rest of the swarm. See hivemind.DHT docs
  :param prefix: a unique DHT key used for matchmaking. E.g. this can be your experiment name with optional suffixes
  :param reuse_grad_buffers: if True, use model's .grad buffers for accumulating gradients over multiple steps.
    This is more memory efficient, but it requires that the user does *not* call zero_grad or clip_by_whatever at all
  :param accumulate_grads_on: if specified, accumulate gradients on this device. By default, this will use the same
    device as model parameters. One can specify a different device (e.g. 'cpu' vs 'cuda') to save device memory at
    the cost of extra time per step. If reuse_grad_buffers is True, this parameter has no effect.
  :param client_mode: if False, this averager will accept incoming requests from other peers.
    if True, the averager will only join existing groups where at least one peer has client_mode=False.
    By default, this flag is copied from DHTNode inside the ``dht`` instance.
  :param warn: if True, warn when the averager did not reset accumulators after use or did not use averaging results
  :param averaged_grads: if provided, it will be used as a set of averagable gradients
  :param kwargs: see DecentralizedAverager keyword arguments for additional parameters
  
  
  Example:
  
  >>> model = SuchModelMuchLayers()
  >>> opt = torch.optim.Adam(model.parameters())
  >>> grad_averager = GradientAverager(model.parameters(), dht=hivemind.DHT(...))
  >>> next_step_time = hivemind.get_dht_time() + 60   # runs global steps every 60 seconds
  >>> next_step_control = None
  >>> while True:
  >>>    # accumulate as many gradients as you can before next_step_time
  >>>    loss = compute_loss(model, batch_size=32)
  >>>    loss.backward()
  >>>    grad_averager.accumulate_grads_(batch_size=32)
  >>>    # [optional] next step in 5 seconds, start looking for peers in advance
  >>>    if next_step_time - hivemind.get_dht_time() <= 5
  >>>        next_step_control = grad_averager.schedule_step(scheduled_time=next_step_time)
  >>>    # aggregate gradients and perform optimizer step
  >>>    if hivemind.get_dht_time() >= next_step_time:
  >>>        grad_averager.step(control=next_step_control)
  >>>        with grad_averager.use_averaged_gradients():  # this will fill param.grads with aggregated gradients
  >>>            opt.step()  # update model parameters using averaged gradients
  >>>        grad_averager.reset_accumulated_grads_()  # prepare for next step
  >>>        next_step_time = hivemind.get_dht_time() + 60
  >>>        next_step_control = None
class_docstr_tok: ["An","auxiliary","averaging","class","that","is","responsible","for","accumulating","gradients","and","aggregating","them","with","peers",".","GradientAverager","is","meant","to","be","used","within","hivemind.Optimizer",",","but","it","can","be","used","standalone","(","see","example","below",")",".","GradientAverager","manages","three","sets","of","buffers",":","(","1",")","model","gradients","-","the","gradients","associated","with","local","model","parameters","by","PyTorch","(","param.grad",")",".","These","tensors","are","typically","stored","on","device","and","updated","by","torch","autograd","(","2",")","gradient","accumulators","-","an","[","optional","]","set","of","buffers","where","local","gradients","are","accumulated",".","-","note",":","if","reuse_grad_buffers","is","True",",","the","averager","will","use","gradients","from","parameters","as","local","accumulators",",","which","reduces","RAM","usage","but","requires","the","user","to","avoid","calling","zero_grad","\/","clip_grad","manually","(","3",")","averaged","gradients","-","gradient","buffers","that","are","aggregated","in-place","with","peers",",","always","in","host","memory",":","param","parameters",":","pytorch","parameters","for","which","to","aggregate","gradients",":","param","dht",":","a","DHT","instance","connected","to","the","rest","of","the","swarm",".","See","hivemind.DHT","docs",":","param","prefix",":","a","unique","DHT","key","used","for","matchmaking",".","E.g",".","this","can","be","your","experiment","name","with","optional","suffixes",":","param","reuse_grad_buffers",":","if","True",",","use","model","'s",".grad","buffers","for","accumulating","gradients","over","multiple","steps",".","This","is","more","memory","efficient",",","but","it","requires","that","the","user","does","*","not","*","call","zero_grad","or","clip_by_whatever","at","all",":","param","accumulate_grads_on",":","if","specified",",","accumulate","gradients","on","this","device",".","By","default",",","this","will","use","the","same","device","as","model","parameters",".","One","can","specify","a","different","device","(","e.g",".","'cpu","'","vs","'cuda","'",")","to","save","device","memory","at","the","cost","of","extra","time","per","step",".","If","reuse_grad_buffers","is","True",",","this","parameter","has","no","effect",".",":","param","client_mode",":","if","False",",","this","averager","will","accept","incoming","requests","from","other","peers",".","if","True",",","the","averager","will","only","join","existing","groups","where","at","least","one","peer","has","client_mode=False",".","By","default",",","this","flag","is","copied","from","DHTNode","inside","the","``","dht","``","instance",".",":","param","warn",":","if","True",",","warn","when","the","averager","did","not","reset","accumulators","after","use","or","did","not","use","averaging","results",":","param","averaged_grads",":","if","provided",",","it","will","be","used","as","a","set","of","averagable","gradients",":","param","kwargs",":","see","DecentralizedAverager","keyword","arguments","for","additional","parameters","Example",":",">",">",">","model","=","SuchModelMuchLayers","(",")",">",">",">","opt","=","torch.optim.Adam","(","model.parameters","(",")",")",">",">",">","grad_averager","=","GradientAverager","(","model.parameters","(",")",",","dht=hivemind.DHT","(","...",")",")",">",">",">","next_step_time","=","hivemind.get_dht_time","(",")","+","60","#","runs","global","steps","every","60","seconds",">",">",">","next_step_control","=","None",">",">",">","while","True",":",">",">",">","#","accumulate","as","many","gradients","as","you","can","before","next_step_time",">",">",">","loss","=","compute_loss","(","model",",","batch_size=32",")",">",">",">","loss.backward","(",")",">",">",">","grad_averager.accumulate_grads_","(","batch_size=32",")",">",">",">","#","[","optional","]","next","step","in","5","seconds",",","start","looking","for","peers","in","advance",">",">",">","if","next_step_time","-","hivemind.get_dht_time","(",")","<","=","5",">",">",">","next_step_control","=","grad_averager.schedule_step","(","scheduled_time=next_step_time",")",">",">",">","#","aggregate","gradients","and","perform","optimizer","step",">",">",">","if","hivemind.get_dht_time","(",")",">","=","next_step_time",":",">",">",">","grad_averager.step","(","control=next_step_control",")",">",">",">","with","grad_averager.use_averaged_gradients","(",")",":","#","this","will","fill","param.grads","with","aggregated","gradients",">",">",">","opt.step","(",")","#","update","model","parameters","using","averaged","gradients",">",">",">","grad_averager.reset_accumulated_grads_","(",")","#","prepare","for","next","step",">",">",">","next_step_time","=","hivemind.get_dht_time","(",")","+","60",">",">",">","next_step_control","=","None"]
docstr: Begin matchmaking: look for a group of peers and prepare for averaging gradients at a specified time.
  
  :param scheduled_time: expected time when to perform all-reduce. Can be changed using control.scheduled_time
  :param kwargs: any additional keyword args from DecentralizedAverager.step, such as gather, allow_retries, etc
  :note: setting weight at this stage is not supported, please leave this parameter as None
  :returns: step_control - a handle that can be passed into GradientAverager.step to use the pre-scheduled group
  :note: in the current implementation, each step_control can only be used in one step.
docstr_tok: ["Begin","matchmaking",":","look","for","a","group","of","peers","and","prepare","for","averaging","gradients","at","a","specified","time",".",":","param","scheduled_time",":","expected","time","when","to","perform","all-reduce",".","Can","be","changed","using","control.scheduled_time",":","param","kwargs",":","any","additional","keyword","args","from","DecentralizedAverager.step",",","such","as","gather",",","allow_retries",",","etc",":","note",":","setting","weight","at","this","stage","is","not","supported",",","please","leave","this","parameter","as","None",":","returns",":","step_control","-","a","handle","that","can","be","passed","into","GradientAverager.step","to","use","the","pre-scheduled","group",":","note",":","in","the","current","implementation",",","each","step_control","can","only","be","used","in","one","step","."]
code: def schedule_step(
      self, scheduled_time: Optional[DHTExpiration] = None, **kwargs
  ) -> StepControl:
      """
      Begin matchmaking: look for a group of peers and prepare for averaging gradients at a specified time.
  
      :param scheduled_time: expected time when to perform all-reduce. Can be changed using control.scheduled_time
      :param kwargs: any additional keyword args from DecentralizedAverager.step, such as gather, allow_retries, etc
      :note: setting weight at this stage is not supported, please leave this parameter as None
      :returns: step_control - a handle that can be passed into GradientAverager.step to use the pre-scheduled group
      :note: in the current implementation, each step_control can only be used in one step.
      """
      assert (
          kwargs.get("weight") is None
      ), "setting weight in schedule_step is not supported"
      return super().step(
          scheduled_time=scheduled_time,
          wait=False,
          require_trigger=True,
          **kwargs,
      )
  
file_path: hivemind/hivemind/optim/grad_averager.py
imports: import contextlib
  from typing import Callable, Iterable, Iterator, Optional, Sequence, TypeVar
  import torch
  from hivemind.averaging import DecentralizedAverager
  from hivemind.averaging.control import StepControl
  from hivemind.dht import DHT
  from hivemind.utils import DHTExpiration, get_logger
label_desc: Use image node_id 5 for calling the GradientAverager obj's underlying member method code with example usage: obj.schedule_step(scheduled_time) and returns: super