class_docstr: A version of the LARC optimizer which works with FSDP:
  https://github.com/NVIDIA/apex/blob/master/apex/parallel/LARC.py
  
  Args:
      optimizer: Pytorch optimizer to wrap and modify learning rate for.
      distributed_norm:
          if True compute the norm in a distributed fashion
          if False, revert to the same computations as APEX LARC
      trust_coefficient:
          Trust coefficient for calculating the lr.
          See https://arxiv.org/abs/1708.03888
      clip: Decides between clipping or scaling mode of LARC.
          If `clip=True` the learning rate is set to `min(optimizer_lr, local_lr)` for each parameter.
          If `clip=False` the learning rate is set to `local_lr*optimizer_lr`.
      eps: epsilon kludge to help with numerical stability while calculating adaptive_lr
class_docstr_tok: ["A","version","of","the","LARC","optimizer","which","works","with","FSDP",":","https",":","\/\/github.com\/NVIDIA\/apex\/blob\/master\/apex\/parallel\/LARC.py","Args",":","optimizer",":","Pytorch","optimizer","to","wrap","and","modify","learning","rate","for",".","distributed_norm",":","if","True","compute","the","norm","in","a","distributed","fashion","if","False",",","revert","to","the","same","computations","as","APEX","LARC","trust_coefficient",":","Trust","coefficient","for","calculating","the","lr",".","See","https",":","\/\/arxiv.org\/abs\/1708.03888","clip",":","Decides","between","clipping","or","scaling","mode","of","LARC",".","If","`","clip=True","`","the","learning","rate","is","set","to","`","min","(","optimizer_lr",",","local_lr",")","`","for","each","parameter",".","If","`","clip=False","`","the","learning","rate","is","set","to","`","local_lr","*","optimizer_lr","`",".","eps",":","epsilon","kludge","to","help","with","numerical","stability","while","calculating","adaptive_lr"]
docstr: Compute the parameter and gradient norms for each parameter group
  
  If the parameters are sharded, this require reassembling the parameters
  of each parameter group to compute the norm.
  
  The algorithm below makes use of all_reduce instead of all_gather,
  profiting from the fact that the norm can be decomposed into:
  - a sum of squares locally
  - all_reduced to a global sum of squares
  - and then locally apply a square root
  
  Due to the non-associativity of float operations, this computation will
  lead to do different norms compared to the native implementation, but
  has a much lower impact on communication:
  
  ```
  def all_gather_norm(data):
      tensor_list = [torch.zeros_like(data) for _ in range(get_global_group().size())]
      dist.all_gather(tensor_list, data)
      tensor_list = torch.cat(tensor_list, dim=0)
      return torch.norm(tensor_list)
  
  param_norms = []
  grad_norms = []
  for group in param_groups:
      for p in group["params"]:
          if p.grad is not None:
              param_norms.append(all_gather_norm(p.data))
              grad_norms.append(all_gather_norm(p.grad.data))
  return param_norms, grad_norms
  ```
docstr_tok: ["Compute","the","parameter","and","gradient","norms","for","each","parameter","group","If","the","parameters","are","sharded",",","this","require","reassembling","the","parameters","of","each","parameter","group","to","compute","the","norm",".","The","algorithm","below","makes","use","of","all_reduce","instead","of","all_gather",",","profiting","from","the","fact","that","the","norm","can","be","decomposed","into",":","-","a","sum","of","squares","locally","-","all_reduced","to","a","global","sum","of","squares","-","and","then","locally","apply","a","square","root","Due","to","the","non-associativity","of","float","operations",",","this","computation","will","lead","to","do","different","norms","compared","to","the","native","implementation",",","but","has","a","much","lower","impact","on","communication",":","``","`","def","all_gather_norm","(","data",")",":","tensor_list","=","[","torch.zeros_like","(","data",")","for","_","in","range","(","get_global_group","(",")",".size","(",")",")","]","dist.all_gather","(","tensor_list",",","data",")","tensor_list","=","torch.cat","(","tensor_list",",","dim=0",")","return","torch.norm","(","tensor_list",")","param_norms","=","[","]","grad_norms","=","[","]","for","group","in","param_groups",":","for","p","in","group","[","``","params","''","]",":","if","p.grad","is","not","None",":","param_norms.append","(","all_gather_norm","(","p.data",")",")","grad_norms.append","(","all_gather_norm","(","p.grad.data",")",")","return","param_norms",",","grad_norms","``","`"]
code: def _compute_distributed_norms(param_groups):
      """
      Compute the parameter and gradient norms for each parameter group
  
      If the parameters are sharded, this require reassembling the parameters
      of each parameter group to compute the norm.
  
      The algorithm below makes use of all_reduce instead of all_gather,
      profiting from the fact that the norm can be decomposed into:
      - a sum of squares locally
      - all_reduced to a global sum of squares
      - and then locally apply a square root
  
      Due to the non-associativity of float operations, this computation will
      lead to do different norms compared to the native implementation, but
      has a much lower impact on communication:
  
      ```
      def all_gather_norm(data):
          tensor_list = [torch.zeros_like(data) for _ in range(get_global_group().size())]
          dist.all_gather(tensor_list, data)
          tensor_list = torch.cat(tensor_list, dim=0)
          return torch.norm(tensor_list)
  
      param_norms = []
      grad_norms = []
      for group in param_groups:
          for p in group["params"]:
              if p.grad is not None:
                  param_norms.append(all_gather_norm(p.data))
                  grad_norms.append(all_gather_norm(p.grad.data))
      return param_norms, grad_norms
      ```
      """
      param_squares = []
      grad_squares = []
      for group in param_groups:
          for p in group["params"]:
              if p.grad is not None:
                  param_squares.append((p.data**2).sum())
                  grad_squares.append((p.grad.data**2).sum())
      all_squared = torch.stack(param_squares + grad_squares)
      dist.all_reduce(all_squared, group=get_global_group())
      all_squared = all_squared**0.5
      param_norms = all_squared[: len(param_squares)]
      grad_norms = all_squared[len(param_squares) :]
      return param_norms, grad_norms
  
imports: from typing import Any, Dict
  import torch
  import torch.distributed
  from classy_vision.optim import register_optimizer, SGD
  from vissl.utils.fsdp_utils import get_global_group
label_desc: Use image node_id 14 for calling the LARC_FSDP obj's underlying member method code with example usage: obj._compute_distributed_norms(param_groups) and returns: param_norms,  grad_norms