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: Clip all gradients at this point in time. The norm is computed over all gradients together, as if they were
  concatenated into a single vector. Gradients are modified in-place.
  
  Arguments:
      max_norm (float or int): max norm of the gradients
      norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm.
  
  Returns:
      Total norm of the parameters (viewed as a single vector).
  
  .. note: This is analogous to `torch.nn.utils.clip_grad_norm_` but handles the partitioning and multiple devices per rank
      under the hood. The default torch util is not applicable here, because each rank only has a partial view of all the grads
      in the model, so calling it in the OSS context would lead to different scaling being applied per subset of model parameters
  
  .. warning: This needs to be called on all ranks, since synchronization primitives will be used
docstr_tok: ["Clip","all","gradients","at","this","point","in","time",".","The","norm","is","computed","over","all","gradients","together",",","as","if","they","were","concatenated","into","a","single","vector",".","Gradients","are","modified","in-place",".","Arguments",":","max_norm","(","float","or","int",")",":","max","norm","of","the","gradients","norm_type","(","float","or","int",")",":","type","of","the","used","p-norm",".","Can","be","``","'inf","'","``","for","infinity","norm",".","Returns",":","Total","norm","of","the","parameters","(","viewed","as","a","single","vector",")",".","..","note",":","This","is","analogous","to","`","torch.nn.utils.clip_grad_norm_","`","but","handles","the","partitioning","and","multiple","devices","per","rank","under","the","hood",".","The","default","torch","util","is","not","applicable","here",",","because","each","rank","only","has","a","partial","view","of","all","the","grads","in","the","model",",","so","calling","it","in","the","OSS","context","would","lead","to","different","scaling","being","applied","per","subset","of","model","parameters","..","warning",":","This","needs","to","be","called","on","all","ranks",",","since","synchronization","primitives","will","be","used"]
code: def clip_grad_norm(
      self,
      max_norm: Union[float, int],
      norm_type: Union[float, int] = 2.0,
      filter_params_fn: Callable[[Any], Any] = None,
  ) -> torch.Tensor:
      """
      Clip all gradients at this point in time. The norm is computed over all gradients together, as if they were
      concatenated into a single vector. Gradients are modified in-place.
  
      Arguments:
          max_norm (float or int): max norm of the gradients
          norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm.
  
      Returns:
          Total norm of the parameters (viewed as a single vector).
  
      .. note: This is analogous to `torch.nn.utils.clip_grad_norm_` but handles the partitioning and multiple devices per rank
          under the hood. The default torch util is not applicable here, because each rank only has a partial view of all the grads
          in the model, so calling it in the OSS context would lead to different scaling being applied per subset of model parameters
  
      .. warning: This needs to be called on all ranks, since synchronization primitives will be used
  
      """
  
      # Compute the max norm for this shards's worth of gradients
      max_norm = float(max_norm)
      norm_type = float(norm_type)
  
      with profiler.record_function("fairscale::oss::clip_grad_norm"):
          # Option to filter parameters from the grad_norm calculation. This is useful for model parallelism.
          # To avoid double counting, only consider parameters on rank zero + anything marked 'model_parallel'
          # 'model_parallel' flag is set in Megatron-LM:
          # https://github.com/NVIDIA/Megatron-LM/blob/19301985dd31c8b612095cbad15bd903e8ddd497/megatron/mpu/layers.py#L54
          local_params = (
              filter_params_fn(self._local_params)
              if filter_params_fn is not None
              else self._local_params
          )
  
          local_norm = calc_grad_norm(local_params, norm_type).to(
              self._default_device
          )
          # Compute the norm on this grad set,
          # then sync all the norms from all ranks
          if norm_type == inf:
              total_norm = local_norm
              # all reduce over data parallel and model parallel workers
              dist.all_reduce(
                  total_norm,
                  op=torch.distributed.ReduceOp.MAX,
                  group=dist.group.WORLD,
              )
          else:
              # local norm result can be accumulated with the remote ones if put to the right power
              # n_i = sum_rank(a^p)^1/p
              # -> n_total = all_reduce(n_i^p)^(1/p) = sum_i(n_i^p)^1/p = sum_i(sum_rank(a^p))^1/p
              # all reduce over data parallel and model parallel workers
              total_norm = local_norm**norm_type
              dist.all_reduce(total_norm)
              total_norm = total_norm ** (1.0 / norm_type)
  
          clip_coef = torch.tensor(
              max_norm, dtype=total_norm.dtype, device=total_norm.device
          ) / (total_norm + 1e-6)
          if clip_coef < 1:
              for (
                  device,
                  device_params,
              ) in self._per_device_params.items():
                  for p in filter(
                      lambda x: x.grad is not None,
                      device_params[self.rank],
                  ):
                      p.grad.detach().mul_(clip_coef.to(device))
  
      return total_norm
  
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 4 for calling the OSS obj's underlying member method code with example usage: obj.clip_grad_norm(max_norm, norm_type, filter_params_fn) and returns: total_norm