class_docstr: Implements Shampoo algorithm.
  
  See details in:
  - https://arxiv.org/pdf/1802.09568.pdf
  - https://arxiv.org/pdf/2002.09018.pdf
  
  If root_inv_dist = True, assigns each parameter's preconditioners to different GPUs in a
  round-robin fashion.
  
  Uses infinity norm to evaluate residuals and errors.
  
  Args:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups
      lr (float): learning rate (Default: 1e-2)
      betas (Tuple[float, float]): coefficients used for computing running averages
          of gradient and its square (Default: (0.9, 1.0))
      epsilon (float): term added to the denominator to improve numerical stability (Default: 1e-12)
      use_bias_correction (bool): flag for using bias correction (Default: True)
      adam_w_mode (bool): Flag for using AdamW-style weight decay (Default: True)
      weight_decay (float): weight decay (L2 penalty) (Default: 0)
      update_freq (int): frequency for updating inverse preconditioner (Default: 100)
      init_delay (int): initial delay before starting to compute root inverse (Default: 1000)
      threshold (int): threshold for switching to diagonal preconditioner (Default: 1024)
      preconditioner_dtype (torch.dtype): data type for preconditioner (Default: torch.float)
      large_dim_method (LargeDimMethod): method for handling large scale tensors. (Default: LargeDimMethod.BLOCKING)
      root_inv_dist (bool): distributes root inverse computation across multiple GPU workers (Default: True)
      use_merge_dims (bool): merge dimensions if possible while respecting threshold. (Default: True)
      grafting_type (GraftingType): Selects grafting method. (Default: GraftingType.ADAGRAD)
      grafting_epsilon (float): Epsilon for grafting method. (Default: 1e-3)
      grafting_beta2 (float): Exponential moving average factor for grafting method. (Default: 1.0)
docstr: Performs a single optimization step.
  
  Args:
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Args",":","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure=None):
      """Performs a single optimization step.
  
      Args:
          closure (callable, optional): A closure that reevaluates the model
              and returns the loss.
  
      """
      loss = None
      if closure is not None:
          with torch.enable_grad():
              loss = closure()
  
      self.iter += 1
  
      # update preconditioners
      self._update_preconditioners()
  
      # compute root inverse if delay is 0
      if (
          self.iter % self.update_freq == 0
          and self.iter >= self.init_delay
      ):
          self._compute_root_inverse()
          if self.root_inv_dist:
              self._broadcast_inv_preconditioners()
  
      # perform update
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
  
              # Initialize gradient, states, and dim for parameter
              grad = p.grad
              state = self.state[p]
              beta1, _ = group["betas"]
              weight_decay = group["weight_decay"]
              lr = group["lr"]
  
              # TODO: Sparse case still not supported.
              if p.grad.is_sparse:
                  raise Exception(
                      "Sparse parameters are not currently supported by Shampoo."
                  )
  
              # Dense case
              else:
                  # incorporate momentum
                  if beta1 != 0:
                      # compute bias corrections if necessary
                      bias_correction1 = 1.0
                      if self.use_bias_correction and beta1 < 1:
                          bias_correction1 -= beta1**self.iter
  
                      # modify grad with momentum term
                      if state["exp_avg"] is None:
                          state["exp_avg"] = torch.zeros_like(
                              grad, memory_format=torch.preserve_format
                          )
                      buf = state["exp_avg"]
                      buf.mul_(beta1).add_(grad, alpha=1 - beta1)
                      grad.copy_(buf / bias_correction1)
  
                  # perform AdamW weight decay
                  if self.adam_w_mode and weight_decay != 0:
                      p.mul_(1 - lr * weight_decay)
  
                  # compute preconditioned gradient and update parameters
                  state["preconditioners"].precondition_and_update(
                      p, grad, lr
                  )
  
      return loss
  
file_path: bert-ultrafast/training/cramming/backend/optimizers/shampoo/shampoo.py
imports: import logging
  from typing import Tuple
  import torch
  import torch.distributed
  from torch.optim.optimizer import Optimizer
  from .shampoo_utils import BlockShampooPreconditioner, AdagradPreconditioner, ShampooPreconditioner, LargeDimMethod, GraftingType
label_desc: Use image node_id 6 for calling the Shampoo obj's underlying member method code with example usage: obj.step(closure) and returns: loss