class_docstr: This AdamW implementation is copied from Huggingface.
  We modified it with Adagrad scaling by rms of a weight tensor
  Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay
  Regularization](https://arxiv.org/abs/1711.05101).
  Parameters:
      params (`Iterable[nn.parameter.Parameter]`):
          Iterable of parameters to optimize or dictionaries defining parameter groups.
      lr (`float`, *optional*, defaults to 1e-3):
          The learning rate to use.
      betas (`Tuple[float,float]`, *optional*, defaults to (0.9, 0.999)):
          Adam's betas parameters (b1, b2).
      eps (`float`, *optional*, defaults to 1e-6):
          Adam's epsilon for numerical stability.
      weight_decay (`float`, *optional*, defaults to 0):
          Decoupled weight decay to apply.
      correct_bias (`bool`, *optional*, defaults to `True`):
          Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`).
      no_deprecation_warning (`bool`, *optional*, defaults to `False`):
          A flag used to disable the deprecation warning (set to `True` to disable the warning).
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()
  
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
              grad = p.grad.data
              if grad.is_sparse:
                  raise RuntimeError(
                      "Adam does not support sparse gradients, please consider SparseAdam instead"
                  )
  
              state = self.state[p]
              beta1, beta2 = group["betas"]
  
              # State initialization
              if len(state) == 0:
                  state["step"] = 0
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p.data)
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = torch.zeros_like(p.data)
  
              exp_avg, exp_avg_sq = (
                  state["exp_avg"],
                  state["exp_avg_sq"],
              )
  
              state["step"] += 1
  
              # Decay the first and second moment running average coefficient
              # In-place operations to update the averages at the same time
              exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
              exp_avg_sq.mul_(beta2).addcmul_(
                  grad, grad, value=1.0 - beta2
              )
              denom = exp_avg_sq.sqrt().add_(group["eps"])
  
              step_size = group["lr"]
              if group["correct_bias"]:  # No bias correction for Bert
                  bias_correction1 = 1.0 - beta1 ** state["step"]
                  bias_correction2 = 1.0 - beta2 ** state["step"]
                  step_size = (
                      step_size
                      * math.sqrt(bias_correction2)
                      / bias_correction1
                  )
  
              # /Adapt Step from Adagrad
              step_size = step_size * max(1e-3, self._rms(p.data))
              # /Adapt Step from Adagrad
  
              p.data.addcdiv_(exp_avg, denom, value=-step_size)
  
              # Just adding the square of the weights to the loss function is *not*
              # the correct way of using L2 regularization/weight decay with Adam,
              # since that will interact with the m and v parameters in strange ways.
              #
              # Instead we want to decay the weights in a manner that doesn't interact
              # with the m/v parameters. This is equivalent to adding the square
              # of the weights to the loss with plain (non-momentum) SGD.
              # Add weight decay at the end (fixed version)
              if group["weight_decay"] > 0.0:
                  p.data.add_(
                      p.data,
                      alpha=(-group["lr"] * group["weight_decay"]),
                  )
  
      return loss
  
file_path: bert-ultrafast/training/cramming/backend/optimizers/nanoT5_optimizer.py
imports: import torch
  import math
  from typing import Iterable, Tuple
  from torch import nn
label_desc: Use image node_id 3 for calling the AdamWScale obj's underlying member method code with example usage: obj.step(closure) and returns: loss