class_docstr: Implements Lamb algorithm.
  
  It has been 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 square (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)
  
  .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
      https://arxiv.org/abs/1904.00962
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, fp16, and bf16 params
      g_all_32, g_all_16, g_all_bf16 = [], [], []
      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)
              elif p.dtype == torch.bfloat16:
                  g_all_bf16.append(p.grad.data)
              else:
                  raise RuntimeError(
                      "FusedLAMB only support fp16, fp32, and bf16."
                  )
  
      device = self.param_groups[0]["params"][0].device
      g_norm_32, g_norm_16, g_norm_bf16 = (
          torch.zeros(1, device=device),
          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]
      if len(g_all_bf16) > 0:
          g_norm_bf16 = multi_tensor_applier(
              self.multi_tensor_l2norm,
              self._dummy_overflow_buf,
              [g_all_bf16],
              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, g_norm_bf16]],
          False,
      )[0]
  
      max_grad_norm = 1.0
      clipped_ratio = max_grad_norm / max(
          global_grad_norm, max_grad_norm
      )
  
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
              p.grad.data *= clipped_ratio
              grad = p.grad.data
              if grad.is_sparse:
                  raise RuntimeError(
                      "Lamb does not support sparse gradients, consider SparseAdam instad."
                  )
  
              state = self.state[p]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["m"] = torch.zeros_like(p.data)
                  # Exponential moving average of squared gradient values
                  state["v"] = torch.zeros_like(p.data)
  
              m_t, v_t = state["m"], state["v"]
              beta1, beta2 = group["betas"]
  
              state["step"] += 1
  
              # m_t = beta1 * m + (1 - beta1) * g_t
              m_t.mul_(beta1).add_(grad, alpha=1 - beta1)
              # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
              if len(g_all_16) > 0:
                  v_t.mul_(beta2)
                  v_t = v_t.to(torch.float32)
                  grad32 = grad.to(torch.float32)
                  v_t.addcmul_(grad32, grad32, value=1 - beta2)
              else:
                  v_t.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
  
              # Debiasing
              m_t_hat = m_t / (1.0 - beta1 ** state["step"])
              v_t_hat = v_t / (1.0 - beta2 ** state["step"])
  
              update = m_t_hat / v_t_hat.sqrt().add(group["eps"])
  
              if group["weight_decay"] != 0:
                  update.add_(p.data, alpha=group["weight_decay"])
  
              trust_ratio = 1.0
              w_norm = p.data.to(torch.float32).pow(2).sum().sqrt()
              g_norm = update.pow(2).sum().sqrt()
              if w_norm > 0 and g_norm > 0:
                  trust_ratio = w_norm / g_norm
  
              state["w_norm"] = w_norm
              state["g_norm"] = g_norm
              state["trust_ratio"] = trust_ratio
  
              step_size = group["lr"]
  
              p.data.add_(update, alpha=-step_size * trust_ratio)
  
      return loss
  
imports: import unittest
  import os
  import torch
  from torch.optim import Optimizer
  import apex
  from apex.multi_tensor_apply import multi_tensor_applier
  from itertools import product
label_desc: Use image node_id 2 for calling the RefLAMB obj's underlying member method code with example usage: obj.step(closure) and returns: loss