code: def step(self, closure=None, grad_scaler=None):
      loss = None
      if closure is not None:
          loss = closure()
  
      # The full precision params are set up in the first step of the optimizer
      # instead of in the constructor because the full precision params will get out
      # out of sync with the model params if DDP syncs the model params across devices
      # after the optimizer is constructed.
      if len(self.param_groups_full_precision) == 0:
          self._setup_full_precision_params()
  
      # create separate grad lists for params
      grad_list = []
      for gid, group in enumerate(self.param_groups):
          for pid, p in enumerate(group["params"]):
              assert (
                  group["params"][0].dtype == p.dtype
              ), "Error: Parameters are not of the identical type: {} != {}".format(
                  group["params"][0].dtype, p.dtype
              )
              if p.grad is None:
                  continue
              grad_list.append(p.grad)
  
      # Overflow check of gradients
      device = self.param_groups[0]["params"][0].device
      found_inf = (
          grad_scaler._check_inf_per_device(self)[device]
          if grad_scaler is not None
          else torch.zeros((1,), device=device)
      )
      self._dummy_overflow_buf.copy_(found_inf)
  
      # Get unscale scale factor
      scale, inv_scale = None, None
      if grad_scaler:
          scale = grad_scaler._get_scale_async()
          inv_scale = scale.double().reciprocal().float()
      else:
          scale = torch.ones((1,), device=device)
          inv_scale = torch.ones((1,), device=device)
  
      # grad_norm is of scaled gradients.
      # So, multiply `max_grad_norm` by scale.
      max_grad_norm = self.defaults["max_grad_norm"] * scale
      grad_norm = multi_tensor_applier(
          self.multi_tensor_l2norm,
          self._dummy_overflow_buf,
          [grad_list],
          False,
      )[0]
  
      # Run LAMB optimization math
      for gid, (group, group_full) in enumerate(
          zip(self.param_groups, self.param_groups_full_precision)
      ):
          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
          group["step"] += (self._dummy_overflow_buf != 1).to(torch.int)
  
          state_lists = [
              [],  # (0) grads
              [],  # (1) params
              [],  # (2) momentum state
              [],  # (3) velocity state
          ]
          if self.reduced_precision_dtype is not None:
              state_lists.append([])  # (4) params reduced_dtype
  
          for p, p_full in zip(group["params"], group_full["params"]):
              if p.grad is None:
                  continue
              assert not p.grad.is_sparse
  
              state = self.state[p]
              # State initialization
              if len(state) == 0:
                  dtype = p.dtype
                  if (
                      self.reduced_precision_dtype is not None
                      and p.dtype == self.reduced_precision_dtype
                  ):
                      dtype = torch.float32
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(
                      p.data, dtype=dtype
                  )
                  # Exponential moving average of gradient values
                  state["exp_avg_sq"] = torch.zeros_like(
                      p.data, dtype=dtype
                  )
  
              if self.reduced_precision_dtype is not None:
                  state_lists[0].append(p.grad.data)
                  state_lists[1].append(p_full.data)
                  state_lists[2].append(state["exp_avg"])
                  state_lists[3].append(state["exp_avg_sq"])
                  state_lists[4].append(p.data)
              else:
                  state_lists[0].append(p.grad.data)
                  state_lists[1].append(p.data)
                  state_lists[2].append(state["exp_avg"])
                  state_lists[3].append(state["exp_avg_sq"])
  
          multi_tensor_applier(
              self.multi_tensor_lamb,
              self._dummy_overflow_buf,
              state_lists,
              group["lr"],
              beta1,
              beta2,
              group["eps"],
              group["step"],
              bias_correction,
              group["weight_decay"],
              grad_averaging,
              self.adam_w_mode,
              grad_norm,
              max_grad_norm,
              self.use_nvlamb,
              found_inf,
              inv_scale,
          )
  
      return loss
  
file_path: apex/apex/optimizers/fused_mixed_precision_lamb.py
imports: import torch
  from copy import deepcopy
  from itertools import chain
  from collections import defaultdict, abc
  from apex.multi_tensor_apply import multi_tensor_applier
label_desc: Use image node_id 5 for calling the FusedMixedPrecisionLamb obj's underlying member method code with example usage: obj.step(closure, grad_scaler) and returns: loss