class_docstr: Implements Adam algorithm.
  
  Supports parameters updating on both GPU and CPU, depending on the device of parameters.
  But the parameters and gradients should on the same device:
    * Parameters on CPU and gradients on CPU is allowed.
    * Parameters on GPU and gradients on GPU is allowed.
    * Parameters on GPU and gradients on CPU is **not** allowed.
  
  `CPUAdam` requires CUDA extensions which can be built during installation or runtime.
  
  This version of CPU Adam accelerates parameters updating on CPU with SIMD.
  Support of AVX2 or AVX512 is required.
  
  The GPU part is implemented in an naive way.
  
  CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.
  
  :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,
  or ``torch.optim.Adam`` with ``adamw_mode=False``
  
  Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.
  
  Arguments:
      model_params (iterable): iterable of parameters of 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 yet in CPUAdam!
      adamw_mode (boolean, optional): Apply L2 regularization or weight decay
          True for decoupled weight decay(also known as AdamW) (default: True)
      simd_log (boolean, optional): whether to show if you are using SIMD to
          accelerate. (default: False)
      nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.
      nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.
          If it's ``None``, a random temporary directory will be used. Defaults to None.
  
  .. _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
class_docstr_tok: ["Implements","Adam","algorithm",".","Supports","parameters","updating","on","both","GPU","and","CPU",",","depending","on","the","device","of","parameters",".","But","the","parameters","and","gradients","should","on","the","same","device",":","*","Parameters","on","CPU","and","gradients","on","CPU","is","allowed",".","*","Parameters","on","GPU","and","gradients","on","GPU","is","allowed",".","*","Parameters","on","GPU","and","gradients","on","CPU","is","*","*","not","*","*","allowed",".","`","CPUAdam","`","requires","CUDA","extensions","which","can","be","built","during","installation","or","runtime",".","This","version","of","CPU","Adam","accelerates","parameters","updating","on","CPU","with","SIMD",".","Support","of","AVX2","or","AVX512","is","required",".","The","GPU","part","is","implemented","in","an","naive","way",".","CPU","Adam","also","supports","the","hybrid","precision","calculation",",","eg",".","fp32","parameters","and","fp16","gradients",".",":","class",":","`","colossalai.nn.optimizer.CPUAdam","`","may","be","used","as","a","drop-in","replacement","for","``","torch.optim.AdamW","``",",","or","``","torch.optim.Adam","``","with","``","adamw_mode=False","``","Adam","was","been","proposed","in","`","Adam",":","A","Method","for","Stochastic","Optimization","`","_",".","Arguments",":","model_params","(","iterable",")",":","iterable","of","parameters","of","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","yet","in","CPUAdam","!","adamw_mode","(","boolean",",","optional",")",":","Apply","L2","regularization","or","weight","decay","True","for","decoupled","weight","decay","(","also","known","as","AdamW",")","(","default",":","True",")","simd_log","(","boolean",",","optional",")",":","whether","to","show","if","you","are","using","SIMD","to","accelerate",".","(","default",":","False",")","nvme_offload_fraction","(","float",",","optional",")",":","Fraction","of","optimizer","states","to","be","offloaded","to","NVMe",".","Defaults","to","0.0.","nvme_offload_dir","(","Optional","[","str","]",",","optional",")",":","Directory","to","save","NVMe","offload","files",".","If","it","'s","``","None","``",",","a","random","temporary","directory","will","be","used",".","Defaults","to","None",".","..","_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"]
code: def torch_adam_update(
      self,
      data,
      grad,
      exp_avg,
      exp_avg_sq,
      lr,
      beta1,
      beta2,
      eps,
      weight_decay,
      bias_correction1,
      bias_correction2,
      use_adamw=False,
  ):
      grad = grad.to(data.dtype)
  
      if weight_decay != 0:
          if use_adamw:
              data.mul_(1 - lr * weight_decay)
          else:
              grad = grad.add(data, alpha=weight_decay)
  
      # Decay the first and second moment running average coefficient
      exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
      exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
  
      # TODO(jiaruifang) dose not support amsgrad
      denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(
          eps
      )
  
      step_size = lr / bias_correction1
  
      data.addcdiv_(exp_avg, denom, value=-step_size)
  
file_path: ColossalAI/colossalai/nn/optimizer/cpu_adam.py
imports: import math
  import platform
  from typing import Optional
  import torch
  from colossalai.kernel.op_builder import ArmCPUAdamBuilder, CPUAdamBuilder
  from .nvme_optimizer import NVMeOptimizer
label_desc: Use image node_id 2 for calling the CPUAdam obj's underlying member method code with example usage: obj.torch_adam_update(data, grad, exp_avg, exp_avg_sq, lr, beta1, beta2, eps, weight_decay, bias_correction1, bias_correction2, use_adamw) without return types