class_docstr: ModuleBackend is a wrapper around torch module that allows it to run tasks asynchronously with Runtime
  By default, ModuleBackend handles three types of requests:
  
   - forward - receive inputs and compute outputs. Concurrent requests will be batched for better GPU utilization.
   - backward - receive gradients w.r.t. outputs, compute gradients w.r.t. inputs and **update expert**. Also batched.
   - get_info - return expert metadata. Not batched.
  
  :param module: nn.Module to be wrapped into a backend. Arbitrary pytorch module with a few limitations:
  
   - Experts must always receive the same set of args and kwargs and produce output tensors of same type
   - All args, kwargs and outputs must be **tensors** where 0-th dimension represents to batch size
   - We recommend using experts that are ~invariant to the order in which they process batches
   - Using randomness (e.g. Dropout) leads to different samples at forward and backward. If you want consistency,
      you should explicitly register these random variables as model inputs or outputs.
      See hivemind.utils.custom_layers.DeterministicDropout for an example
  
  :param optimizer: torch optimizer to be applied on every backward call
  :param scheduler: a function to create the learning rate scheduler for the expert
  :param args_schema: description of positional arguments to expert.forward, list of BatchTensorProto
  :param kwargs_schema: description of keyword arguments to expert.forward, dict of BatchTensorProto
  :param outputs_schema: description of outputs from expert.forward, nested structure of BatchTensorProto
  :param kwargs: extra parameters to be forwarded into TaskPool.__init__
class_docstr_tok: ["ModuleBackend","is","a","wrapper","around","torch","module","that","allows","it","to","run","tasks","asynchronously","with","Runtime","By","default",",","ModuleBackend","handles","three","types","of","requests",":","-","forward","-","receive","inputs","and","compute","outputs",".","Concurrent","requests","will","be","batched","for","better","GPU","utilization",".","-","backward","-","receive","gradients","w.r.t",".","outputs",",","compute","gradients","w.r.t",".","inputs","and","*","*","update","expert","*","*",".","Also","batched",".","-","get_info","-","return","expert","metadata",".","Not","batched",".",":","param","module",":","nn.Module","to","be","wrapped","into","a","backend",".","Arbitrary","pytorch","module","with","a","few","limitations",":","-","Experts","must","always","receive","the","same","set","of","args","and","kwargs","and","produce","output","tensors","of","same","type","-","All","args",",","kwargs","and","outputs","must","be","*","*","tensors","*","*","where","0-th","dimension","represents","to","batch","size","-","We","recommend","using","experts","that","are","~invariant","to","the","order","in","which","they","process","batches","-","Using","randomness","(","e.g",".","Dropout",")","leads","to","different","samples","at","forward","and","backward",".","If","you","want","consistency",",","you","should","explicitly","register","these","random","variables","as","model","inputs","or","outputs",".","See","hivemind.utils.custom_layers.DeterministicDropout","for","an","example",":","param","optimizer",":","torch","optimizer","to","be","applied","on","every","backward","call",":","param","scheduler",":","a","function","to","create","the","learning","rate","scheduler","for","the","expert",":","param","args_schema",":","description","of","positional","arguments","to","expert.forward",",","list","of","BatchTensorProto",":","param","kwargs_schema",":","description","of","keyword","arguments","to","expert.forward",",","dict","of","BatchTensorProto",":","param","outputs_schema",":","description","of","outputs","from","expert.forward",",","nested","structure","of","BatchTensorProto",":","param","kwargs",":","extra","parameters","to","be","forwarded","into","TaskPool.__init__"]
docstr: Apply backward pass to an aggregated batch of requests. Used by Runtime, do not call this manually
  To submit a request for asynchronous processing, please use ``ModuleBackend.backward_pool.submit_task``.
  
  Subclassing:
     This method receives a sequence of torch tensors following ``nested_flatten(self.backward_schema)``;
  
     It should return gradients w.r.t. inputs that follow ``nested_flatten(self.forward_schema)``;
  
     Runtime doesn't guarantee that backward will be performed in the same order and for the same data
     as forward, so we recommend stateless backward pass that re-runs expert forward pass inside backward.
  
     Please make sure to call ``ModuleBackend.on_backward`` after each call to backward
docstr_tok: ["Apply","backward","pass","to","an","aggregated","batch","of","requests",".","Used","by","Runtime",",","do","not","call","this","manually","To","submit","a","request","for","asynchronous","processing",",","please","use","``","ModuleBackend.backward_pool.submit_task","``",".","Subclassing",":","This","method","receives","a","sequence","of","torch","tensors","following","``","nested_flatten","(","self.backward_schema",")","``",";","It","should","return","gradients","w.r.t",".","inputs","that","follow","``","nested_flatten","(","self.forward_schema",")","``",";","Runtime","does","n't","guarantee","that","backward","will","be","performed","in","the","same","order","and","for","the","same","data","as","forward",",","so","we","recommend","stateless","backward","pass","that","re-runs","expert","forward","pass","inside","backward",".","Please","make","sure","to","call","``","ModuleBackend.on_backward","``","after","each","call","to","backward"]
code: def backward(self, *inputs: torch.Tensor) -> Tuple[torch.Tensor, ...]:
      """
      Apply backward pass to an aggregated batch of requests. Used by Runtime, do not call this manually
      To submit a request for asynchronous processing, please use ``ModuleBackend.backward_pool.submit_task``.
  
      Subclassing:
         This method receives a sequence of torch tensors following ``nested_flatten(self.backward_schema)``;
  
         It should return gradients w.r.t. inputs that follow ``nested_flatten(self.forward_schema)``;
  
         Runtime doesn't guarantee that backward will be performed in the same order and for the same data
         as forward, so we recommend stateless backward pass that re-runs expert forward pass inside backward.
  
         Please make sure to call ``ModuleBackend.on_backward`` after each call to backward
      """
      (args, kwargs), grad_outputs = nested_pack(
          inputs, structure=self.backward_schema
      )
  
      with torch.enable_grad():
          args = [
              tensor.detach().requires_grad_(True)
              if tensor.is_floating_point()
              else tensor.detach()
              for tensor in args
          ]
          kwargs = {
              input_key: (
                  tensor.detach().requires_grad_(True)
                  if tensor.is_floating_point()
                  else tensor.detach()
              )
              for input_key, tensor in kwargs.items()
          }
  
          batch_size = args[0].size(0)
  
          outputs = self.module(*args, **kwargs)
          assert nested_compare(
              outputs, grad_outputs
          ), "outputs and grad_outputs must have the same structure"
  
          outputs_flat = tuple(nested_flatten(outputs))
  
          grad_outputs_flat = tuple(
              map(
                  lambda grad, out: grad.to(
                      device=out.device,
                      dtype=out.dtype,
                      non_blocking=True,
                  ),
                  nested_flatten(grad_outputs),
                  outputs_flat,
              )
          )
          torch.autograd.backward(
              outputs_flat,
              grad_tensors=grad_outputs_flat,
              create_graph=False,
              retain_graph=False,
          )
          self.on_backward(batch_size)
  
      return tuple(
          x.grad
          if isinstance(x.grad, torch.Tensor)
          else torch.zeros_like(x)
          for x in nested_flatten((args, kwargs))
      )
  
file_path: hivemind/hivemind/moe/server/module_backend.py
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union
  import torch
  from torch import nn
  from hivemind.moe.server.task_pool import TaskPool
  from hivemind.utils.logging import get_logger
  from hivemind.utils.nested import nested_compare, nested_flatten, nested_map, nested_pack
  from hivemind.utils.tensor_descr import DUMMY_BATCH_SIZE, BatchTensorDescriptor
label_desc: Use image node_id 3 for calling the ModuleBackend obj's underlying member method code with example usage: obj.backward() and returns: tuple