class_docstr: Implements AdamW algorithm.
  
  This class uses the same functional AdamW implementation as
  :class:`torch.optim.AdamW`. The main difference is that it keeps its state
  in single precision (i.e. ``torch.float32``), and temporarily casts half
  precision parameters and gradients (i.e. ``torch.float16`` and
  ``torch.bfloat16``) to single precision during an optimization step for
  better numerical stability.
class_docstr_tok: ["Implements","AdamW","algorithm",".","This","class","uses","the","same","functional","AdamW","implementation","as",":","class",":","`","torch.optim.AdamW","`",".","The","main","difference","is","that","it","keeps","its","state","in","single","precision","(","i.e",".","``","torch.float32","``",")",",","and","temporarily","casts","half","precision","parameters","and","gradients","(","i.e",".","``","torch.float16","``","and","``","torch.bfloat16","``",")","to","single","precision","during","an","optimization","step","for","better","numerical","stability","."]
code: def _do_step(self) -> None:
      self._cuda_graph_capture_health_check()  # type: ignore[attr-defined]
  
      for pg in self.param_groups:
          params_with_grad: List[Tensor] = []
          fp32_params_with_grad: List[Tensor] = []
          fp32_grads: List[Tensor] = []
          steps: List[Tensor] = []
          exp_avgs: List[Tensor] = []
          exp_avg_sqs: List[Tensor] = []
          max_exp_avg_sqs: List[Tensor] = []
          amsgrad = pg["amsgrad"]
          beta1, beta2 = pg["betas"]
  
          for p in pg["params"]:
              self._init_param(
                  p,
                  pg,
                  params_with_grad,
                  fp32_params_with_grad,
                  fp32_grads,
                  steps,
                  exp_avgs,
                  exp_avg_sqs,
                  max_exp_avg_sqs,
                  amsgrad,
              )
  
          kwargs: Dict[str, Any] = {}
  
          if pg["differentiable"]:
              kwargs["differentiable"] = True
  
          if (impl := pg["impl"]) != "auto":
              kwargs[impl] = True
  
              for attr in ["grad_scale", "found_inf"]:
                  try:
                      kwargs[attr] = getattr(self, attr)
                  except AttributeError:
                      pass
  
          # Mitigates a shape issue that is specific to PyTorch 2.0.1.
          try:
              kwargs["found_inf"] = kwargs["found_inf"].squeeze()
          except KeyError:
              pass
  
          adamw(
              fp32_params_with_grad,
              fp32_grads,
              exp_avgs,
              exp_avg_sqs,
              max_exp_avg_sqs,
              steps,
              amsgrad=amsgrad,
              beta1=beta1,
              beta2=beta2,
              lr=pg["lr"],
              weight_decay=pg["weight_decay"],
              eps=pg["eps"],
              maximize=pg["maximize"],
              capturable=pg["capturable"],
              **kwargs,
          )
  
          # Cast parameters back to fp16/bf16.
          for p, fp32_p in zip(params_with_grad, fp32_params_with_grad):
              if p.dtype == torch.float16 or p.dtype == torch.bfloat16:
                  p.copy_(fp32_p)
  
file_path: fairseq2/src/fairseq2/optim/adamw.py
imports: from itertools import chain
  from typing import Any, Dict, Iterable, List, Literal, Tuple, Union, cast, final
  import torch
  from torch import Tensor
  from torch.optim.adamw import adamw
  from fairseq2.optim.optimizer_base import OptimizerBase
  from fairseq2.typing import finaloverride
label_desc: Use image node_id 3 for calling the AdamW obj's underlying member method code with example usage: obj._do_step() without return types