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: Make all params which are on the same device and tied to the same rank views of a single buffer.
  This is used at construction time, and anytime parameter trainability is changed (frozen or unfrozen) and
  `refresh_trainability` is called.
docstr_tok: ["Make","all","params","which","are","on","the","same","device","and","tied","to","the","same","rank","views","of","a","single","buffer",".","This","is","used","at","construction","time",",","and","anytime","parameter","trainability","is","changed","(","frozen","or","unfrozen",")","and","`","refresh_trainability","`","is","called","."]
code: def _setup_flat_buffers(self) -> None:
      """Make all params which are on the same device and tied to the same rank views of a single buffer.
      This is used at construction time, and anytime parameter trainability is changed (frozen or unfrozen) and
      `refresh_trainability` is called.
      """
  
      for device, per_rank_params in self._per_device_params.items():
          # Only wipe the existing buckets if there are none
          # (could be that this is called twice, when trainability changes)
          if device not in self.buckets.keys():
              self.buckets[device] = {}
  
          # Make parameters a view of the bucket
          for dst_rank, params in enumerate(per_rank_params):
              if len(params) > 0:
                  # Clone the non-trainable params, if in a bucket it will get destroyed
                  for param in filter(
                      lambda x: not x.requires_grad, params
                  ):
                      param.data = param.data.detach().clone()
  
                  # Merge all the trainable params in a single bucket
                  trainable_params = list(
                      filter(lambda x: x.requires_grad, params)
                  )
                  if trainable_params:
                      buffer_size = sum(
                          map(lambda x: x.numel(), trainable_params)
                      )
                      bucket = ParamBucket(
                          size=buffer_size,
                          dtype=trainable_params[0].dtype,
                          device=device,
                      )
  
                      for param in trainable_params:
                          bucket.add_param(param)
  
                      self.buckets[device][dst_rank] = bucket
  
      # Clear the buffer keys which are not in use anymore (could be that the devices changed)
      devices_in_use = list(self._per_device_params.keys())
      devices_to_pop = list(
          filter(lambda x: x not in devices_in_use, self.buckets.keys())
      )
      for d in devices_to_pop:
          self.buckets.pop(d)
  
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 17 for calling the OSS obj's underlying member method code with example usage: obj._setup_flat_buffers() without return types