class_docstr: Implements LAMB algorithm.
  
  `FusedLAMB` requires CUDA extensions which can be built during installation or runtime.
  
  This version of fused LAMB implements 2 fusions.
  
    * Fusion of the LAMB update's elementwise operations
    * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
  
  :class:`colossalai.nn.optimizer.FusedLAMB`'s usage is identical to any ordinary Pytorch optimizer
  
  :class:`colossalai.nn.optimizer.FusedLAMB` may be used with or without Amp.
  
  LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups.
      lr (float, optional): learning rate. (default: 1e-3)
      betas (Tuple[float, float], optional): coefficients used for computing
          running averages of gradient and its norm. (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve
          numerical stability. (default: 1e-6)
      weight_decay (float, optional): weight decay (L2 penalty) (default: 0.01)
      amsgrad (boolean, optional): whether to use the AMSGrad variant of this
          algorithm from the paper `On the Convergence of Adam and Beyond`_
          NOT SUPPORTED now! (default: False)
      adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
          True for decoupled weight decay(also known as AdamW) (default: True)
      grad_averaging (bool, optional): whether apply (1-beta2) to grad when
          calculating running averages of gradient. (default: True)
      set_grad_none (bool, optional): whether set grad to None when zero_grad()
          method is called. (default: True)
      max_grad_norm (float, optional): value used to clip global grad norm
          (default: 1.0)
      use_nvlamb (boolean, optional): Apply adaptive learning rate to 0.0
          weight decay parameter (default: False)
  
  .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
      https://arxiv.org/abs/1904.00962
  .. _On the Convergence of Adam and Beyond:
      https://openreview.net/forum?id=ryQu7f-RZ
docstr: Performs a single optimization step.
  
  Arguments:
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure=None):
      """Performs a single optimization step.
  
      Arguments:
          closure (callable, optional): A closure that reevaluates the model
              and returns the loss.
      """
      loss = None
      if closure is not None:
          loss = closure()
  
      # create separate grad lists for fp32 and fp16 params
      g_all_32, g_all_16 = [], []
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
              if p.dtype == torch.float32:
                  g_all_32.append(p.grad.data)
              elif p.dtype == torch.float16:
                  g_all_16.append(p.grad.data)
              else:
                  raise RuntimeError(
                      "FusedLAMB only support fp16 and fp32."
                  )
  
      device = self.param_groups[0]["params"][0].device
      g_norm_32, g_norm_16 = torch.zeros(1, device=device), torch.zeros(
          1, device=device
      )
      # compute grad norm for two lists
      if len(g_all_32) > 0:
          g_norm_32 = multi_tensor_applier(
              self.multi_tensor_l2norm,
              self._dummy_overflow_buf,
              [g_all_32],
              False,
          )[0]
      if len(g_all_16) > 0:
          g_norm_16 = multi_tensor_applier(
              self.multi_tensor_l2norm,
              self._dummy_overflow_buf,
              [g_all_16],
              False,
          )[0]
  
      # blend two grad norms to get global grad norm
      global_grad_norm = multi_tensor_applier(
          self.multi_tensor_l2norm,
          self._dummy_overflow_buf,
          [[g_norm_32, g_norm_16]],
          False,
      )[0]
      max_grad_norm = self.defaults["max_grad_norm"]
  
      for group in self.param_groups:
          bias_correction = 1 if group["bias_correction"] else 0
          beta1, beta2 = group["betas"]
          grad_averaging = 1 if group["grad_averaging"] else 0
  
          # assume same step across group now to simplify things
          # per parameter step can be easily support by making it tensor, or pass list into kernel
          if "step" in group:
              group["step"] += 1
          else:
              group["step"] = 1
  
          # create lists for multi-tensor apply
          g_16, p_16, m_16, v_16 = [], [], [], []
          g_32, p_32, m_32, v_32 = [], [], [], []
  
          for p in group["params"]:
              if p.grad is None:
                  continue
              if p.grad.data.is_sparse:
                  raise RuntimeError(
                      "FusedLAMB does not support sparse gradients, please consider SparseAdam instead"
                  )
  
              state = self.state[p]
              # State initialization
              if len(state) == 0:
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p)
                  # Exponential moving average of gradient values
                  state["exp_avg_sq"] = torch.zeros_like(p)
  
              if p.dtype == torch.float16:
                  g_16.append(p.grad.data)
                  p_16.append(p.data)
                  m_16.append(state["exp_avg"])
                  v_16.append(state["exp_avg_sq"])
              elif p.dtype == torch.float32:
                  g_32.append(p.grad.data)
                  p_32.append(p.data)
                  m_32.append(state["exp_avg"])
                  v_32.append(state["exp_avg_sq"])
              else:
                  raise RuntimeError(
                      "FusedLAMB only support fp16 and fp32."
                  )
  
          if len(g_16) > 0:
              multi_tensor_applier(
                  self.multi_tensor_lamb,
                  self._dummy_overflow_buf,
                  [g_16, p_16, m_16, v_16],
                  group["lr"],
                  beta1,
                  beta2,
                  group["eps"],
                  group["step"],
                  bias_correction,
                  group["weight_decay"],
                  grad_averaging,
                  self.adam_w_mode,
                  global_grad_norm,
                  max_grad_norm,
                  self.use_nvlamb,
              )
          if len(g_32) > 0:
              multi_tensor_applier(
                  self.multi_tensor_lamb,
                  self._dummy_overflow_buf,
                  [g_32, p_32, m_32, v_32],
                  group["lr"],
                  beta1,
                  beta2,
                  group["eps"],
                  group["step"],
                  bias_correction,
                  group["weight_decay"],
                  grad_averaging,
                  self.adam_w_mode,
                  global_grad_norm,
                  max_grad_norm,
                  self.use_nvlamb,
              )
  
      return loss
  
imports: import torch
  from colossalai.utils import multi_tensor_applier
label_desc: Use image node_id 3 for calling the FusedLAMB obj's underlying member method code with example usage: obj.step(closure) and returns: loss