docstr: Args:
          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 coefficient (default: 1e-2)
  
          # Any Precision specific
          use_kahan_summation = creates auxiliary buffer to ensure high precision
          model param updates (default: False)
          momentum_dtype = dtype for momentum  (default: BFloat32)
          variance_dtype = dtype for uncentered variance (default: BFloat16)
          compensation_buffer_dtype  = dtype for Kahan summation
                                       buffer (default: BFloat16)
  
          # Usage
          This optimizer implements optimizer states, and Kahan summation
          for high precision updates, all in user controlled dtypes.
          Defaults are variance in BF16, Momentum in FP32.
          This can be run in FSDP mixed precision, amp, or full precision,
          depending on what training pipeline you wish to work with.
  
          Setting to use_kahan_summation = False, and changing momentum and
          variance dtypes to FP32, reverts this to a standard AdamW optimizer.
docstr_tok: ["Args",":","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","coefficient","(","default",":","1e-2",")","#","Any","Precision","specific","use_kahan_summation","=","creates","auxiliary","buffer","to","ensure","high","precision","model","param","updates","(","default",":","False",")","momentum_dtype","=","dtype","for","momentum","(","default",":","BFloat32",")","variance_dtype","=","dtype","for","uncentered","variance","(","default",":","BFloat16",")","compensation_buffer_dtype","=","dtype","for","Kahan","summation","buffer","(","default",":","BFloat16",")","#","Usage","This","optimizer","implements","optimizer","states",",","and","Kahan","summation","for","high","precision","updates",",","all","in","user","controlled","dtypes",".","Defaults","are","variance","in","BF16",",","Momentum","in","FP32",".","This","can","be","run","in","FSDP","mixed","precision",",","amp",",","or","full","precision",",","depending","on","what","training","pipeline","you","wish","to","work","with",".","Setting","to","use_kahan_summation","=","False",",","and","changing","momentum","and","variance","dtypes","to","FP32",",","reverts","this","to","a","standard","AdamW","optimizer","."]
code: def __init__(
      self,
      params,
      lr=1e-3,
      betas=(0.9, 0.999),
      eps=1e-8,
      weight_decay=0.0,
      use_kahan_summation=False,
      momentum_dtype=torch.bfloat16,
      variance_dtype=torch.bfloat16,
      compensation_buffer_dtype=torch.bfloat16,
  ):
      """
      Args:
              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 coefficient (default: 1e-2)
  
              # Any Precision specific
              use_kahan_summation = creates auxiliary buffer to ensure high precision
              model param updates (default: False)
              momentum_dtype = dtype for momentum  (default: BFloat32)
              variance_dtype = dtype for uncentered variance (default: BFloat16)
              compensation_buffer_dtype  = dtype for Kahan summation
                                           buffer (default: BFloat16)
  
              # Usage
              This optimizer implements optimizer states, and Kahan summation
              for high precision updates, all in user controlled dtypes.
              Defaults are variance in BF16, Momentum in FP32.
              This can be run in FSDP mixed precision, amp, or full precision,
              depending on what training pipeline you wish to work with.
  
              Setting to use_kahan_summation = False, and changing momentum and
              variance dtypes to FP32, reverts this to a standard AdamW optimizer.
  
      """
      defaults = dict(
          lr=lr,
          betas=betas,
          eps=eps,
          weight_decay=weight_decay,
          use_kahan_summation=use_kahan_summation,
          momentum_dtype=momentum_dtype,
          variance_dtype=variance_dtype,
          compensation_buffer_dtype=compensation_buffer_dtype,
      )
  
      super().__init__(params, defaults)
  
imports: import torch
  from torch.optim.optimizer import Optimizer
label_desc: Use image node_id 1 to create a new AnyPrecisionAdamW object from inherited base classes: Optimizer with example: obj = AnyPrecisionAdamW(params, lr, betas, eps, weight_decay, use_kahan_summation, momentum_dtype, variance_dtype, compensation_buffer_dtype)