text: import torch
  
  from distseq.training.ops.pytorch.builder import AdamBuilder
  
  fused_adam_cuda = None
  
  
  class LSAdam(torch.optim.Optimizer):
      """
      Modified from Fairseq and Use LightSeq adam kernel.
  
      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-8)
          weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
          amsgrad (boolean, optional): whether to use the AMSGrad variant of this
              algorithm from the paper `On the Convergence of Adam and Beyond`_
              (default: False) NOT SUPPORTED in LSAdam!
          eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
              adds eps to the bias-corrected second moment estimate before
              evaluating square root instead of adding it to the square root of
              second moment estimate as in the original paper. (default: False)
      .. _Adam: A Method for Stochastic Optimization:
          https://arxiv.org/abs/1412.6980
      .. _On the Convergence of Adam and Beyond:
          https://openreview.net/forum?id=ryQu7f-RZ
      """
  
      def __init__(
          self,
          params,
          lr=1e-3,
          bias_correction=True,
          betas=(0.9, 0.999),
          eps=1e-8,
          eps_inside_sqrt=False,
          weight_decay=0.0,
          max_grad_norm=0.0,
          amsgrad=False,
      ):
          global fused_adam_cuda
  
          if fused_adam_cuda is None:
              fused_adam_cuda = AdamBuilder().load()
  
          if amsgrad:
              raise RuntimeError("LSAdam does not support the AMSGrad variant.")
          defaults = {
              "lr": lr,
              "bias_correction": bias_correction,
              "betas": betas,
              "eps": eps,
              "weight_decay": weight_decay,
              "max_grad_norm": max_grad_norm,
          }
          super().__init__(params, defaults)
          self.eps_mode = 0 if eps_inside_sqrt else 1
  
      @property
      def supports_memory_efficient_fp16(self):
          return True
  
      @property
      def supports_flat_params(self):
          return True
  
      @property
      def supports_step_with_scale(self):
          return True
  
      def step(self, closure=None, grads=None, scale=1.0, grad_norms=None):
          """Performs a single optimization step.
          Arguments:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
              grads (list of tensors, optional): weight gradient to use for the
                  optimizer update. If gradients have type torch.half, parameters
                  are expected to be in type torch.float. (default: None)
              output params (list of tensors, optional): A reduced precision copy
                  of the updated weights written out in addition to the regular
                  updated weights. Have to be of same type as gradients. (default: None)
              scale (float, optional): factor to divide gradient tensor values
                  by before applying to weights. (default: 1)
          """
          loss = None
          if closure is not None:
              loss = closure()
  
          if grads is None:
              grads_group = [None] * len(self.param_groups)
          # backward compatibility
          # assuming a list/generator of parameter means single group
          elif isinstance(grads, types.GeneratorType):
              grads_group = [grads]
          elif type(grads[0]) != list:
              grads_group = [grads]
          else:
              grads_group = grads
  
          if grad_norms is None:
              grad_norms = [None] * len(self.param_groups)
  
          for group, grads_this_group, grad_norm in zip(
              self.param_groups, grads_group, grad_norms
          ):
              if grads_this_group is None:
                  grads_this_group = [None] * len(group["params"])
  
              # compute combined scale factor for this group
              combined_scale = scale
              if group.get("max_grad_norm", 0) > 0:
                  # norm is in fact norm*scale
                  clip = ((grad_norm / scale) + 1e-6) / group["max_grad_norm"]
                  if clip > 1:
                      combined_scale = clip * scale
  
              bias_correction = 1 if group.get("bias_correction", 1) else 0
  
              for p, grad in zip(group["params"], grads_this_group):
                  # note: p.grad should not ever be set for correct
                  # operation of mixed precision optimizer that sometimes
                  # sends None gradients
                  if p.grad is None and grad is None:
                      continue
                  if grad is None:
                      grad = p.grad.data
                  if grad.is_sparse:
                      raise RuntimeError(
                          "LSAdam does not support sparse gradients, "
                          "please consider SparseAdam instead"
                      )
  
                  p_data_fp32 = p.data.float()
  
                  state = self.state[p]
  
                  # State initialization
                  if len(state) == 0:
                      state["step"] = 0
                      # Exponential moving average of gradient values
                      state["exp_avg"] = torch.zeros_like(p_data_fp32)
                      # Exponential moving average of squared gradient values
                      state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
                  else:
                      state["exp_avg"] = state["exp_avg"].to(p_data_fp32)
                      state["exp_avg_sq"] = state["exp_avg_sq"].to(p_data_fp32)
  
                  exp_avg = state["exp_avg"]
                  exp_avg_sq = state["exp_avg_sq"]
                  beta1, beta2 = group["betas"]
  
                  state["step"] += 1
  
                  out_p = p.data
                  with torch.cuda.device(p.device):
                      fused_adam_cuda.adam(
                          p_data_fp32,
                          out_p,
                          exp_avg,
                          exp_avg_sq,
                          grad,
                          group["lr"],
                          beta1,
                          beta2,
                          group["eps"],
                          combined_scale,
                          state["step"],
                          self.eps_mode,
                          bias_correction,
                          group["weight_decay"],
                      )
  
          return loss
  