class_docstr: Wrap the model, and reduce the gradients to the right rank during the backward pass.
  
  - the partition is given by the sharded optimizer
  - wrap the base model with a model which knows where to reduce each gradient
  - add an autograd function which calls the model grad dispatch on the way back
  
   Args:
      module (nn.Module):
          model to be wrapped
      sharded_optimizer (OSS, or list of OSS):
          the sharded optimizer(s) which will decide the gradient partitioning
  
  Keyword Args:
      process_group (group):
          torch.distributed group (default: group.WORLD)
      broadcast_buffers (bool):
          Whether to additionally broadcast model buffers in between ranks at the beginning of each forward pass.
          Same setting as in Pytorch DDP, this is in addition to the broadcast and reduction of the model parameters.
      sync_models_at_startup (bool):
          Synchronize the models in between the ranks when starting up. Not needed if each rank has the same seed,
          or the training restarts from a saved state
      reduce_buffer_size (int):
          The max size of the buffer used to batch the small parameter tensors, in number of elements (default 0 - unused).
          this will impact the long term memory consumption, because these buckets correspond to parameters which will not be sharded.
          Set to 0 to remove all bucketing, 1M to 8M is usually reasonable.
      auto_refresh_trainable (bool):
          (default: True) Check whether the parameters trainability (`requires_grad`) has changed and update both ShardedDDP
          and OSS automatically if this is the case. If set to False, `refresh_trainable()` needs to be called anytime
          a parameter is frozen or unfrozen.
      reduce_fp16 (bool):
          cast the grads to fp16 before reducing. Not needed if the model is already fp16, but will probably improve performance
          for multi node jobs using PyTorch AMP. The effect is similar to DDP's fp16_compress_hook_ and will also save some memory.
      warn_on_trainable_params_changed (bool):
          When set to False no warning will be logged whenever a parameter trainability change has been detected.
          Default is True.
  
  .. _fp16_compress_hook: https://pytorch.org/docs/1.8.0/ddp_comm_hooks.html?highlight=fp16#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook
  
  .. warning:
      ShardedDDP implements gradient sharding, meaning that each rank only owns a unique shard of the model gradients
      after the backward pass, in order to save memory and some communication bandwidth.
  
  .. warning:
      As a consequence of sharding:
          * in case of gradient clipping, one has to use the `clip_grad_norm` exposed by
              the `optimizer state sharding wrapper <fairscale.optim.OSS>`
  
          * after loss.backward() (or equivalent) each rank will have `None` in place of some param.grad
  
          * Pytorch and Apex AMP implementations will hang when used in conjunction with `ShardedDDP`.
              One needs a `shard-aware grad scaler<ShardedGradScaler>`, which is proposed in `fairscale.optim.grad_scaler`,
              compatible with PytorchAMP.
  
  .. warning:
      If `auto_refresh_trainable` is set to `True` (this is the default) then any trainability change in the model graph will be handled
      automatically.
      If `auto_refresh_trainable` is set to `False`, ShardedDDP will not refresh its assumptions with respect to trainable parameters
      for every forward pass, in the hope of saving some time. If some parameters are frozen or unfrozen over time, please refresh
      ShardedDDP assumptions by calling `refresh_trainable()` just after said change (before the next forward pass).
class_docstr_tok: ["Wrap","the","model",",","and","reduce","the","gradients","to","the","right","rank","during","the","backward","pass",".","-","the","partition","is","given","by","the","sharded","optimizer","-","wrap","the","base","model","with","a","model","which","knows","where","to","reduce","each","gradient","-","add","an","autograd","function","which","calls","the","model","grad","dispatch","on","the","way","back","Args",":","module","(","nn.Module",")",":","model","to","be","wrapped","sharded_optimizer","(","OSS",",","or","list","of","OSS",")",":","the","sharded","optimizer","(","s",")","which","will","decide","the","gradient","partitioning","Keyword","Args",":","process_group","(","group",")",":","torch.distributed","group","(","default",":","group.WORLD",")","broadcast_buffers","(","bool",")",":","Whether","to","additionally","broadcast","model","buffers","in","between","ranks","at","the","beginning","of","each","forward","pass",".","Same","setting","as","in","Pytorch","DDP",",","this","is","in","addition","to","the","broadcast","and","reduction","of","the","model","parameters",".","sync_models_at_startup","(","bool",")",":","Synchronize","the","models","in","between","the","ranks","when","starting","up",".","Not","needed","if","each","rank","has","the","same","seed",",","or","the","training","restarts","from","a","saved","state","reduce_buffer_size","(","int",")",":","The","max","size","of","the","buffer","used","to","batch","the","small","parameter","tensors",",","in","number","of","elements","(","default","0","-","unused",")",".","this","will","impact","the","long","term","memory","consumption",",","because","these","buckets","correspond","to","parameters","which","will","not","be","sharded",".","Set","to","0","to","remove","all","bucketing",",","1M","to","8M","is","usually","reasonable",".","auto_refresh_trainable","(","bool",")",":","(","default",":","True",")","Check","whether","the","parameters","trainability","(","`","requires_grad","`",")","has","changed","and","update","both","ShardedDDP","and","OSS","automatically","if","this","is","the","case",".","If","set","to","False",",","`","refresh_trainable","(",")","`","needs","to","be","called","anytime","a","parameter","is","frozen","or","unfrozen",".","reduce_fp16","(","bool",")",":","cast","the","grads","to","fp16","before","reducing",".","Not","needed","if","the","model","is","already","fp16",",","but","will","probably","improve","performance","for","multi","node","jobs","using","PyTorch","AMP",".","The","effect","is","similar","to","DDP","'s","fp16_compress_hook_","and","will","also","save","some","memory",".","warn_on_trainable_params_changed","(","bool",")",":","When","set","to","False","no","warning","will","be","logged","whenever","a","parameter","trainability","change","has","been","detected",".","Default","is","True",".","..","_fp16_compress_hook",":","https",":","\/\/pytorch.org\/docs\/1.8.0\/ddp_comm_hooks.html","?","highlight=fp16","#","torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook","..","warning",":","ShardedDDP","implements","gradient","sharding",",","meaning","that","each","rank","only","owns","a","unique","shard","of","the","model","gradients","after","the","backward","pass",",","in","order","to","save","memory","and","some","communication","bandwidth",".","..","warning",":","As","a","consequence","of","sharding",":","*","in","case","of","gradient","clipping",",","one","has","to","use","the","`","clip_grad_norm","`","exposed","by","the","`","optimizer","state","sharding","wrapper","<","fairscale.optim.OSS",">","`","*","after","loss.backward","(",")","(","or","equivalent",")","each","rank","will","have","`","None","`","in","place","of","some","param.grad","*","Pytorch","and","Apex","AMP","implementations","will","hang","when","used","in","conjunction","with","`","ShardedDDP","`",".","One","needs","a","`","shard-aware","grad","scaler","<","ShardedGradScaler",">","`",",","which","is","proposed","in","`","fairscale.optim.grad_scaler","`",",","compatible","with","PytorchAMP",".","..","warning",":","If","`","auto_refresh_trainable","`","is","set","to","`","True","`","(","this","is","the","default",")","then","any","trainability","change","in","the","model","graph","will","be","handled","automatically",".","If","`","auto_refresh_trainable","`","is","set","to","`","False","`",",","ShardedDDP","will","not","refresh","its","assumptions","with","respect","to","trainable","parameters","for","every","forward","pass",",","in","the","hope","of","saving","some","time",".","If","some","parameters","are","frozen","or","unfrozen","over","time",",","please","refresh","ShardedDDP","assumptions","by","calling","`","refresh_trainable","(",")","`","just","after","said","change","(","before","the","next","forward","pass",")","."]
docstr: Moves and/or casts the parameters and buffers.
  
  Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
  floating point desired :attr:`dtype` s. In addition, this method will
  only cast the floating point parameters and buffers to :attr:`dtype`
  (if given). The integral parameters and buffers will be moved
  :attr:`device`, if that is given, but with dtypes unchanged. When
  :attr:`non_blocking` is set, it tries to convert/move asynchronously
  with respect to the host if possible, e.g., moving CPU Tensors with
  pinned memory to CUDA devices.
  
  .. note::
      This method modifies the module in-place.
  
  .. warning:
      Device changes are not supported, and this will raise an exception. The issue in that case is not
      really ShardedDDP, but OSS which will not be aware of the device change, and whose buffers will be
      in a broken state.
  
  Arguments:
      device (:class:`torch.device`): the desired device of the parameters and buffers in this module.
      dtype (:class:`torch.dtype`): the desired floating point type of the floating point parameters and buffers.
      non_blocking (bool): make it an asynchronous call.
  
  Returns:
      Module: self.
docstr_tok: ["Moves","and\/or","casts","the","parameters","and","buffers",".","Its","signature","is","similar","to",":","meth",":","`","torch.Tensor.to","`",",","but","only","accepts","floating","point","desired",":","attr",":","`","dtype","`","s.","In","addition",",","this","method","will","only","cast","the","floating","point","parameters","and","buffers","to",":","attr",":","`","dtype","`","(","if","given",")",".","The","integral","parameters","and","buffers","will","be","moved",":","attr",":","`","device","`",",","if","that","is","given",",","but","with","dtypes","unchanged",".","When",":","attr",":","`","non_blocking","`","is","set",",","it","tries","to","convert\/move","asynchronously","with","respect","to","the","host","if","possible",",","e.g.",",","moving","CPU","Tensors","with","pinned","memory","to","CUDA","devices",".","..","note",":",":","This","method","modifies","the","module","in-place",".","..","warning",":","Device","changes","are","not","supported",",","and","this","will","raise","an","exception",".","The","issue","in","that","case","is","not","really","ShardedDDP",",","but","OSS","which","will","not","be","aware","of","the","device","change",",","and","whose","buffers","will","be","in","a","broken","state",".","Arguments",":","device","(",":","class",":","`","torch.device","`",")",":","the","desired","device","of","the","parameters","and","buffers","in","this","module",".","dtype","(",":","class",":","`","torch.dtype","`",")",":","the","desired","floating","point","type","of","the","floating","point","parameters","and","buffers",".","non_blocking","(","bool",")",":","make","it","an","asynchronous","call",".","Returns",":","Module",":","self","."]
code: def to(  # type: ignore
      self,
      device: Optional[torch.device],
      dtype: Optional[torch.dtype] = None,
      non_blocking: bool = False,
  ) -> "ShardedDataParallel":
      """
      Moves and/or casts the parameters and buffers.
  
      Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
      floating point desired :attr:`dtype` s. In addition, this method will
      only cast the floating point parameters and buffers to :attr:`dtype`
      (if given). The integral parameters and buffers will be moved
      :attr:`device`, if that is given, but with dtypes unchanged. When
      :attr:`non_blocking` is set, it tries to convert/move asynchronously
      with respect to the host if possible, e.g., moving CPU Tensors with
      pinned memory to CUDA devices.
  
      .. note::
          This method modifies the module in-place.
  
      .. warning:
          Device changes are not supported, and this will raise an exception. The issue in that case is not
          really ShardedDDP, but OSS which will not be aware of the device change, and whose buffers will be
          in a broken state.
  
      Arguments:
          device (:class:`torch.device`): the desired device of the parameters and buffers in this module.
          dtype (:class:`torch.dtype`): the desired floating point type of the floating point parameters and buffers.
          non_blocking (bool): make it an asynchronous call.
  
      Returns:
          Module: self.
      """
      if isinstance(device, str):
          device = torch.device(device)
  
      assert (
          device is None
          or len(self._buckets.keys()) == 0
          or device.type in map(lambda x: x.type, self._buckets.keys())
      ), "Changing devices is not supported, because this would break OSSs state"
  
      assert (
          len(self._buckets.keys()) < 2
      ), "Several devices specified to begin with, incompatible with setting a single device here"
  
      self.module.to(
          device=device, dtype=dtype, non_blocking=non_blocking
      )
  
      # Re-build the buckets, hooks, etc..
      self.refresh_trainable()
  
imports: from collections import deque
  import contextlib
  import functools
  from itertools import chain
  import logging
  from typing import Any, Callable, Deque, Dict, Generator, List, Optional, Union
  import torch
  from torch import nn
  from torch.autograd import Variable
  import torch.autograd.profiler
  import torch.distributed
  from fairscale.internal.params import Workhandle, get_global_rank
  from fairscale.nn.misc import GradBucket
  from fairscale.optim import OSS
label_desc: Use image node_id 3 for calling the ShardedDataParallel obj's underlying member method code with example usage: obj.to(device, dtype, non_blocking) without return types