class_docstr: A base class for offloading optimizer states.
  
  Args:
      params: parameters
      defaults (dict): default dict
      nvme_offload_fraction (float, optional): Fraction of params to be offloaded to NVMe. Defaults to 0.0.
      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.
  
  Raises:
      ImportError: Raise if ``tensornvme`` is not installed.
class_docstr_tok: ["A","base","class","for","offloading","optimizer","states",".","Args",":","params",":","parameters","defaults","(","dict",")",":","default","dict","nvme_offload_fraction","(","float",",","optional",")",":","Fraction","of","params","to","be","offloaded","to","NVMe",".","Defaults","to","0.0.","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",".","Raises",":","ImportError",":","Raise","if","``","tensornvme","``","is","not","installed","."]
docstr: Performs a single optimization step (parameter update).
  
  Example:
  
      >>> self._pre_step('exp_avg', 'exp_avg_sq')
      >>> for group in self.param_groups:
      >>>     for p in group['params']:
      >>>         if p.grad is None:
      >>>             continue
      >>>         state = self.state[p]
      >>>         if len(state) == 0:
      >>>             state['exp_avg'] = ...
      >>>             state['exp_avg_sq'] = ...
      >>>             self._post_state_init(p)
      >>>         if p.device.type == 'cpu':
      >>>             self._pre_update(p, 'exp_avg', 'exp_avg_sq')
      >>>             adam()
      >>>             self._post_update(p, 'exp_avg', 'exp_avg_sq')
      >>>         else:
      >>>             ...
      >>> self._post_step()
  
  Args:
      closure (Optional[Callable[[], float]], optional): A closure that reevaluates the model and
          returns the loss. Optional for most optimizers.
docstr_tok: ["Performs","a","single","optimization","step","(","parameter","update",")",".","Example",":",">",">",">","self._pre_step","(","'exp_avg","'",",","'exp_avg_sq","'",")",">",">",">","for","group","in","self.param_groups",":",">",">",">","for","p","in","group","[","'params","'","]",":",">",">",">","if","p.grad","is","None",":",">",">",">","continue",">",">",">","state","=","self.state","[","p","]",">",">",">","if","len","(","state",")","==","0",":",">",">",">","state","[","'exp_avg","'","]","=","...",">",">",">","state","[","'exp_avg_sq","'","]","=","...",">",">",">","self._post_state_init","(","p",")",">",">",">","if","p.device.type","==","'cpu","'",":",">",">",">","self._pre_update","(","p",",","'exp_avg","'",",","'exp_avg_sq","'",")",">",">",">","adam","(",")",">",">",">","self._post_update","(","p",",","'exp_avg","'",",","'exp_avg_sq","'",")",">",">",">","else",":",">",">",">","...",">",">",">","self._post_step","(",")","Args",":","closure","(","Optional","[","Callable","[","[","]",",","float","]","]",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss",".","Optional","for","most","optimizers","."]
code: def step(
      self, closure: Optional[Callable[[], float]] = ...
  ) -> Optional[float]:
      """Performs a single optimization step (parameter update).
  
      Example:
  
          >>> self._pre_step('exp_avg', 'exp_avg_sq')
          >>> for group in self.param_groups:
          >>>     for p in group['params']:
          >>>         if p.grad is None:
          >>>             continue
          >>>         state = self.state[p]
          >>>         if len(state) == 0:
          >>>             state['exp_avg'] = ...
          >>>             state['exp_avg_sq'] = ...
          >>>             self._post_state_init(p)
          >>>         if p.device.type == 'cpu':
          >>>             self._pre_update(p, 'exp_avg', 'exp_avg_sq')
          >>>             adam()
          >>>             self._post_update(p, 'exp_avg', 'exp_avg_sq')
          >>>         else:
          >>>             ...
          >>> self._post_step()
  
      Args:
          closure (Optional[Callable[[], float]], optional): A closure that reevaluates the model and
              returns the loss. Optional for most optimizers.
      """
      raise NotImplementedError
  
file_path: colossalai/colossalai/nn/optimizer/nvme_optimizer.py
imports: import math
  import os
  import tempfile
  from typing import Callable, Dict, List, Optional
  import torch
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 9 for calling the NVMeOptimizer obj's underlying member method code with example usage: obj.step(closure) without return types