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","."]
docstr: Return the last known global optimizer state. The returned state is compatible with Pytorch, in that the
  sharded properties are not exposed.
  
  
  Arguments:
      all_ranks (bool): materialize the state on all ranks. In that case, `.state_dict()` needs to be called on
      all ranks
  
  Returns:
      a dict with two entries
          * state - a dict holding current optimization state. Its content
              differs between optimizer classes.
  
          * param_groups - a dict containing all parameter groups
  
  .. warning:
      Returning the global state is limited to the replica which was responsible for the consolidation,
      if `all_ranks` was not set to `True`. In that case, the state may also not be up to date,
      depending on when `consolidate_state_dict` was last called.
docstr_tok: ["Return","the","last","known","global","optimizer","state",".","The","returned","state","is","compatible","with","Pytorch",",","in","that","the","sharded","properties","are","not","exposed",".","Arguments",":","all_ranks","(","bool",")",":","materialize","the","state","on","all","ranks",".","In","that","case",",","`",".state_dict","(",")","`","needs","to","be","called","on","all","ranks","Returns",":","a","dict","with","two","entries","*","state","-","a","dict","holding","current","optimization","state",".","Its","content","differs","between","optimizer","classes",".","*","param_groups","-","a","dict","containing","all","parameter","groups","..","warning",":","Returning","the","global","state","is","limited","to","the","replica","which","was","responsible","for","the","consolidation",",","if","`","all_ranks","`","was","not","set","to","`","True","`",".","In","that","case",",","the","state","may","also","not","be","up","to","date",",","depending","on","when","`","consolidate_state_dict","`","was","last","called","."]
code: def state_dict(self, all_ranks: bool = False) -> Dict[str, Any]:
      """Return the last known global optimizer state. The returned state is compatible with Pytorch, in that the
      sharded properties are not exposed.
  
  
      Arguments:
          all_ranks (bool): materialize the state on all ranks. In that case, `.state_dict()` needs to be called on
          all ranks
  
      Returns:
          a dict with two entries
              * state - a dict holding current optimization state. Its content
                  differs between optimizer classes.
  
              * param_groups - a dict containing all parameter groups
  
      .. warning:
          Returning the global state is limited to the replica which was responsible for the consolidation,
          if `all_ranks` was not set to `True`. In that case, the state may also not be up to date,
          depending on when `consolidate_state_dict` was last called.
      """
  
      if not all_ranks and len(self._all_states) == 0:
          raise RuntimeError(
              "Optimizer state has not been consolidated on this rank. \
                  Please call `consolidate_state_dict()` on all ranks beforehand if you meant to save the global state"
          )
  
      if all_ranks:
          # Consolidate the state on every rank
          self.consolidate_state_dict(recipient_rank=-1)
  
      # Unify the shard states and the state that pytorch would expect, given the model.
      # Indexation needs several redirections, since each shard only knows a limited scope of the model
      # - get the pytorch compliant parameter indexing
      state_dict = super().state_dict()
  
      # - go through the per-shard states, which are all indexed locally
      for rank, s in enumerate(self._all_states):
          # -- match the local indexing and the global partition, update the corresponding saved state globally
          for local_pg, global_pg in zip(
              s["param_groups"], self.partition_parameters()[rank]
          ):
              local_index_to_param_id = {
                  i_param: id(global_pg["params"][i])
                  for i, i_param in enumerate(local_pg["params"])
              }
  
              for local_param_index in local_pg["params"]:
                  # Update the state, if any
                  if local_param_index in s["state"].keys():
                      global_id = self._param_to_index[
                          local_index_to_param_id[local_param_index]
                      ]
                      state_dict["state"][global_id] = s["state"][
                          local_param_index
                      ]
  
      # Make sure that the parameters are sorted in the state, as expected for a pytorch dict
      state_dict["state"] = dict(sorted(state_dict["state"].items()))
  
      return state_dict
  
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 6 for calling the OSS obj's underlying member method code with example usage: obj.state_dict(all_ranks) and returns: state_dict