class_docstr: Wraps an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>`
  optimizer and shards its state as described by ZeRO_.
  ::
  
      opt = OSS(params, optim=torch.optim.Adam, lr=0.01)
  
  .. _ZeRO: https://arxiv.org/abs/1910.02054
  
  We use a greedy algorithm to pack a number of parameters
  at each rank. Each parameter belongs to a single rank and
  is not divided among rank.
  
  After each rank completed their parameter update, they broadcast
  the new version of the parameters to all other ranks to synchronize
  the parameters for next round forward/backward computation.
  
  Args:
      params (list of tensors):
          parameters to be optimized
  Keyword Args:
      optim (torch.nn.Optimizer):
          optimizer to shard (default: SGD)
      group (group):
          torch.distributed group (default: group.WORLD)
      broadcast_buffer_size (int):
          (deprecated) used to cap the size of the broadcast buffers, not being used anymore.
      broadcast_fp16 (bool):
          Compress the model shards in fp16 before sharing them in between ranks.
          This is safe to use when PyTorch AMP is activated. Without torch AMP this will lead to a slight
          degradation in terms of accuracy.
      force_broadcast_object (bool):
          If True, '_broadcast_object' will be used for rebuilding the sharded optimizer.
          If False, whether to use '_broadcast_object' or 'dist.broadcast_object_list' will be determined by GPU capabilities.
          This feature is needed since some newer GPUs still get some memory issues when applying dist.broadcast_object_list.
  
  .. warning: the communication patterns that OSS use depend on the "trainability" graph,
      meaning that all the parameters which `require_grad` are handled differently. This is
      not reevaluated at every step, please use `refresh_trainable()` if your model changed
      (freeze or unfreeze for instance).
      If used with :class:<fairscale.nn.ShardedDDP> then an automatic change detection is possible,
      via the `auto_refresh_trainable` parameter.
class_docstr_tok: ["Wraps","an","arbitrary",":","class",":","`","optim.Optimizer","<","torch.optim.Optimizer",">","`","optimizer","and","shards","its","state","as","described","by","ZeRO_",".",":",":","opt","=","OSS","(","params",",","optim=torch.optim.Adam",",","lr=0.01",")","..","_ZeRO",":","https",":","\/\/arxiv.org\/abs\/1910.02054","We","use","a","greedy","algorithm","to","pack","a","number","of","parameters","at","each","rank",".","Each","parameter","belongs","to","a","single","rank","and","is","not","divided","among","rank",".","After","each","rank","completed","their","parameter","update",",","they","broadcast","the","new","version","of","the","parameters","to","all","other","ranks","to","synchronize","the","parameters","for","next","round","forward\/backward","computation",".","Args",":","params","(","list","of","tensors",")",":","parameters","to","be","optimized","Keyword","Args",":","optim","(","torch.nn.Optimizer",")",":","optimizer","to","shard","(","default",":","SGD",")","group","(","group",")",":","torch.distributed","group","(","default",":","group.WORLD",")","broadcast_buffer_size","(","int",")",":","(","deprecated",")","used","to","cap","the","size","of","the","broadcast","buffers",",","not","being","used","anymore",".","broadcast_fp16","(","bool",")",":","Compress","the","model","shards","in","fp16","before","sharing","them","in","between","ranks",".","This","is","safe","to","use","when","PyTorch","AMP","is","activated",".","Without","torch","AMP","this","will","lead","to","a","slight","degradation","in","terms","of","accuracy",".","force_broadcast_object","(","bool",")",":","If","True",",","'_broadcast_object","'","will","be","used","for","rebuilding","the","sharded","optimizer",".","If","False",",","whether","to","use","'_broadcast_object","'","or","'dist.broadcast_object_list","'","will","be","determined","by","GPU","capabilities",".","This","feature","is","needed","since","some","newer","GPUs","still","get","some","memory","issues","when","applying","dist.broadcast_object_list",".","..","warning",":","the","communication","patterns","that","OSS","use","depend","on","the","``","trainability","''","graph",",","meaning","that","all","the","parameters","which","`","require_grad","`","are","handled","differently",".","This","is","not","reevaluated","at","every","step",",","please","use","`","refresh_trainable","(",")","`","if","your","model","changed","(","freeze","or","unfreeze","for","instance",")",".","If","used","with",":","class",":","<","fairscale.nn.ShardedDDP",">","then","an","automatic","change","detection","is","possible",",","via","the","`","auto_refresh_trainable","`","parameter","."]
code: def __init__(
      self,
      params: _params_t,
      optim: Type[Optimizer] = SGD,
      group: Optional[Any] = None,
      broadcast_buffer_size: int = -1,
      broadcast_fp16: bool = False,
      force_broadcast_object: bool = False,
      **default: Any,
  ):
      # Hold all the model params in the root .param_groups
      self.in_super_constructor = True
      super().__init__(params, default)
      self.in_super_constructor = False
  
      # Partition information. lazy evaluation, computed when requested
      self.__per_device_params: Dict[
          torch.device, List[List[Parameter]]
      ] = OrderedDict()  # device, rank, params
      self.__param_rank: Dict[torch.Tensor, int] = {}
      self._partition_parameters: List[List[dict]] = []
      self.__param_to_index: Dict[int, int] = {}
      self.__local_params: Optional[List[torch.Tensor]] = None
  
      # Default empty values + immutables
      self._optim_defaults = default
      self._optim_constructor = optim
  
      self.group = group if group is not None else dist.group.WORLD
      self.world_size = dist.get_world_size(self.group)
      self.backend = dist.get_backend(self.group)
      self.rank = dist.get_rank(self.group)
      self.global_rank = get_global_rank(self.group, self.rank)
      self._local_to_global_rank = [
          get_global_rank(self.group, i) for i in range(self.world_size)
      ]
  
      self.broadcast_fp16 = broadcast_fp16
      self.force_broadcast_object = force_broadcast_object
      self.buckets: Dict[torch.device, Dict[int, ParamBucket]] = {}
      self._all_states: List[
          Dict[str, Any]
      ] = []  # Optional consolidated optimizer state
      self._default_device = torch.device("cpu")
  
      # Setup everything which is related to the parameters to be trained
      # (partition and optimizer for the shard)
      self.refresh_trainable()
  
imports: from collections import OrderedDict
  import copy
  import io
  from itertools import chain
  import logging
  from math import inf
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
  import torch
  from torch.autograd import profiler
  import torch.distributed
  from torch.nn import Parameter
  from torch.optim import SGD, Optimizer
  from fairscale.internal.params import calc_grad_norm, get_global_rank, recursive_copy_to_device
  from fairscale.nn.misc import ParamBucket
label_desc: Use image node_id 1 to create a new OSS object from inherited base classes: Optimizer with example: obj = OSS(params, optim, group, broadcast_buffer_size, broadcast_fp16, force_broadcast_object)