class_docstr: An optimizer that wraps a _DistributedOptimizer, intercepting push-pull operations.
  This class enables overlapping gradient push-pull with both backward and forward propagation while maintaining
  correct dependencies. It can achieve even higher training performance than the default BytePS with proper system
  parameters. To understand the principles behind barrier crossing, check the paper
  https://dl.acm.org/citation.cfm?id=3359642
class_docstr_tok: ["An","optimizer","that","wraps","a","_DistributedOptimizer",",","intercepting","push-pull","operations",".","This","class","enables","overlapping","gradient","push-pull","with","both","backward","and","forward","propagation","while","maintaining","correct","dependencies",".","It","can","achieve","even","higher","training","performance","than","the","default","BytePS","with","proper","system","parameters",".","To","understand","the","principles","behind","barrier","crossing",",","check","the","paper","https",":","\/\/dl.acm.org\/citation.cfm","?","id=3359642"]
docstr: Add hook before forward propagation of each layer to block forward computation until the push-pull and
  parameter update is finished. The blocking is implemented using a lock.
docstr_tok: ["Add","hook","before","forward","propagation","of","each","layer","to","block","forward","computation","until","the","push-pull","and","parameter","update","is","finished",".","The","blocking","is","implemented","using","a","lock","."]
code: def _register_forward_hooks(self):
      """Add hook before forward propagation of each layer to block forward computation until the push-pull and
      parameter update is finished. The blocking is implemented using a lock.
      """
      # Recursively find all submodules
      submodules = []
      q = queue.LifoQueue()
      for mod in self._model.children():
          q.put(mod)
      while not q.empty():
          mod = q.get()
          if len(list(mod.children())) == 0:
              submodules.append(mod)
          else:
              for m in mod.children():
                  q.put(m)
  
      def pre_forward_hook(mod, input):
          for p in mod.parameters():
              if p in self._handles:
                  del self._handles[p]
              if p not in self._locks:
                  continue
              with self._locks[p]:
                  self._logger.debug(
                      "{} {} is ready.".format(
                          self._desc, self._get_parameter_name(p)
                      )
                  )
  
          self._logger.debug(
              "{} starts forward {}.".format(self._desc, mod)
          )
  
      def after_forward_hook(mod, input, result):
          self._logger.debug(
              "{} finished forward {}.".format(self._desc, mod)
          )
  
      # Register pre-hook and hook for each module
      for mod in reversed(submodules):
          self._logger.debug(
              "{} registers forward hook on module {}".format(
                  self._desc, mod
              )
          )
          mod.register_forward_pre_hook(pre_forward_hook)
          mod.register_forward_hook(after_forward_hook)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from byteps.torch.compression import Compression
  from byteps.torch.ops import push_pull_async_inplace
  from byteps.torch.ops import poll, synchronize
  from byteps.torch.ops import init, shutdown
  from byteps.torch.ops import size, local_size, rank, local_rank
  import threading
  import logging
  import time
  import math
  import torch
  import byteps.torch
label_desc: Use image node_id 10 for calling the _CrossBarrier obj's underlying member method code with example usage: obj._register_forward_hooks() without return types