class_docstr: ## Adam Optimizer for Half Precision Training
  
  We extend [Adam Optimizer](adam.html) but use FP32 to store gradients and moments.
class_docstr_tok: ["#","#","Adam","Optimizer","for","Half","Precision","Training","We","extend","[","Adam","Optimizer","]","(","adam.html",")","but","use","FP32","to","store","gradients","and","moments","."]
docstr: ### Take an update step for a given parameter tensor
  
  * `state` is the optimizer state of the parameter (tensor)
  * `group` stores optimizer attributes of the parameter group
  * `grad` is the current gradient tensor  $g_t$ for the parameter $      heta_{t-1}$
  * `param` is the parameter tensor $     heta_{t-1}$
docstr_tok: ["#","#","#","Take","an","update","step","for","a","given","parameter","tensor","*","`","state","`","is","the","optimizer","state","of","the","parameter","(","tensor",")","*","`","group","`","stores","optimizer","attributes","of","the","parameter","group","*","`","grad","`","is","the","current","gradient","tensor","$","g_t","$","for","the","parameter","$","heta_","{","t-1","}","$","*","`","param","`","is","the","parameter","tensor","$","heta_","{","t-1","}","$"]
code: def step_param(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      grad: torch.Tensor,
      param: torch.nn.Parameter,
  ):
      """
      ### Take an update step for a given parameter tensor
  
      * `state` is the optimizer state of the parameter (tensor)
      * `group` stores optimizer attributes of the parameter group
      * `grad` is the current gradient tensor  $g_t$ for the parameter $\theta_{t-1}$
      * `param` is the parameter tensor $\theta_{t-1}$
      """
  
      # Get the FP32 parameters
      param_fp32 = state["fp32_copy"]
      # Get the FP32 gradients if available
      grad_fp32 = self.grad_fp32.get(param, None)
      if grad_fp32 is not None:
          del self.grad_fp32[param]
          grad = grad_fp32
      else:
          # Otherwise, convert the gradients to FP32
          grad = grad.to(torch.float)
  
      # Calculate weight decay
      grad = self.weight_decay(param_fp32, grad, group)
  
      # Get $m_t$ and $v_t$
      m, v = self.get_mv(state, group, grad)
  
      # Increment $t$ the number of optimizer steps
      state["step"] += 1
  
      # Perform *Adam* update
      self.adam_update(state, group, param_fp32, m, v)
  
      # Set the parameters
      param.data = param_fp32.to(param.dtype)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/adam_fp16.py
imports: from typing import Dict, Tuple, Optional, Any
  import torch
  from torch import nn
  from torch.optim import Optimizer
  from torch.cuda.amp import grad_scaler
  from collections import defaultdict, abc
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.adam import Adam
label_desc: Use image node_id 3 for calling the AdamFP16 obj's underlying member method code with example usage: obj.step_param(state, group, grad, param) without return types