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","."]
code: def _setup_prefetch_params(self) -> List[Parameter]:
      if self.offloader is None:
          return
      assert (
          len(self.prefetch_params) == 0
          and len(self.param_to_prefetch_idx) == 0
      )
      for group in self.param_groups:
          for p in group["params"]:
              if p.grad is None:
                  continue
              if len(self.state[p]) > 0 and self.is_on_nvme[p]:
                  assert p.device.type == "cpu"
                  self.param_to_prefetch_idx[p] = len(
                      self.prefetch_params
                  )
                  self.prefetch_params.append(p)
  
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 4 for calling the NVMeOptimizer obj's underlying member method code with example usage: obj._setup_prefetch_params() without return types