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"]
code: def step(self):
      weight_decays = []
      with torch.no_grad():
          # Compute the parameter norms and gradient norms that are
          # required to find the adaptative_lr of each parameter group
          if self.distributed_norm:
              param_norms, grad_norms = self._compute_distributed_norms(
                  self.optim.param_groups
              )
          else:
              param_norms, grad_norms = self._compute_local_norms(
                  self.optim.param_groups
              )
  
          param_group_count = 0
          for group in self.optim.param_groups:
              # absorb weight decay control from optimizer
              weight_decay = (
                  group["weight_decay"]
                  if "weight_decay" in group
                  else 0
              )
              weight_decays.append(weight_decay)
              group["weight_decay"] = 0
  
              for p in group["params"]:
                  if p.grad is None:
                      continue
  
                  param_norm = param_norms[param_group_count]
                  grad_norm = grad_norms[param_group_count]
                  param_group_count += 1
  
                  if param_norm != 0 and grad_norm != 0:
                      # calculate adaptive lr + weight decay
                      adaptive_lr = (
                          self.trust_coefficient
                          * (param_norm)
                          / (
                              grad_norm
                              + param_norm * weight_decay
                              + self.eps
                          )
                      )
  
                      # clip learning rate for LARC
                      if self.clip:
                          # calculation of adaptive_lr so that when multiplied by lr
                          # it is equal to `min(adaptive_lr, lr)`
                          adaptive_lr = min(
                              adaptive_lr / group["lr"], 1
                          )
  
                      p.grad.data += weight_decay * p.data
                      p.grad.data *= adaptive_lr
  
      self.optim.step()
  
      # return weight decay control to optimizer
      for i, group in enumerate(self.optim.param_groups):
          group["weight_decay"] = weight_decays[i]
  
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 12 for calling the LARC_FSDP obj's underlying member method code with example usage: obj.step() without return types