class_docstr: Manage Embedding Weights on CPU and CUDA memory uses a software cache.
  CPU maintains the entire original weight.
  CUDA maintains a fraction of the weights used in the upcoming computation. The row number in CUDA is controlled by `cuda_row_num`.
  During training, GPU needs to transmit embedding rows between CPU and GPU.
  Args:
      weight (torch.Tensor): the weight of the Embedding layer.
      cuda_row_num (int, optional): the number of rows cached in CUDA memory. Defaults to 0.
      buffer_size (int, optional): the number of rows in a data transmitter buffer. Defaults to 50_000.
      pin_weight (bool, optional): use pin memory to store the cpu weight. If set `True`, the cpu memory usage will increase largely. Defaults to False.
      evict_strategy (EvictionStrategy, optional): the eviction strategy. There are two options.
      `EvictionStrategy.LFU`: use the least frequently used cache.
      `EvictionStrategy.DATASET`: use the stats collected from the target dataset. It usually leads to less cpu-gpu communication volume.
      Defaults to EvictionStrategy.DATASET.
class_docstr_tok: ["Manage","Embedding","Weights","on","CPU","and","CUDA","memory","uses","a","software","cache",".","CPU","maintains","the","entire","original","weight",".","CUDA","maintains","a","fraction","of","the","weights","used","in","the","upcoming","computation",".","The","row","number","in","CUDA","is","controlled","by","`","cuda_row_num","`",".","During","training",",","GPU","needs","to","transmit","embedding","rows","between","CPU","and","GPU",".","Args",":","weight","(","torch.Tensor",")",":","the","weight","of","the","Embedding","layer",".","cuda_row_num","(","int",",","optional",")",":","the","number","of","rows","cached","in","CUDA","memory",".","Defaults","to","0.","buffer_size","(","int",",","optional",")",":","the","number","of","rows","in","a","data","transmitter","buffer",".","Defaults","to","50_000",".","pin_weight","(","bool",",","optional",")",":","use","pin","memory","to","store","the","cpu","weight",".","If","set","`","True","`",",","the","cpu","memory","usage","will","increase","largely",".","Defaults","to","False",".","evict_strategy","(","EvictionStrategy",",","optional",")",":","the","eviction","strategy",".","There","are","two","options",".","`","EvictionStrategy.LFU","`",":","use","the","least","frequently","used","cache",".","`","EvictionStrategy.DATASET","`",":","use","the","stats","collected","from","the","target","dataset",".","It","usually","leads","to","less","cpu-gpu","communication","volume",".","Defaults","to","EvictionStrategy.DATASET","."]
docstr: move the cpu embedding rows w.r.t. ids into CUDA memory
  Args:
      ids (torch.Tensor): the ids to be computed
  Returns:
      torch.Tensor: indices on the cuda_cached_weight.
docstr_tok: ["move","the","cpu","embedding","rows","w.r.t",".","ids","into","CUDA","memory","Args",":","ids","(","torch.Tensor",")",":","the","ids","to","be","computed","Returns",":","torch.Tensor",":","indices","on","the","cuda_cached_weight","."]
code: def prepare_ids(self, ids: torch.Tensor) -> torch.Tensor:
      """
      move the cpu embedding rows w.r.t. ids into CUDA memory
      Args:
          ids (torch.Tensor): the ids to be computed
      Returns:
          torch.Tensor: indices on the cuda_cached_weight.
      """
      torch.cuda.synchronize()
      with self.timer("cache_op") as gtimer:
          # identify cpu rows to cache
          with self.timer("1_identify_cpu_row_idxs") as timer:
              with record_function("(cache) get unique indices"):
                  if self._evict_strategy == EvictionStrategy.LFU:
                      cpu_row_idxs, repeat_times = torch.unique(
                          ids, return_counts=True
                      )
                  else:
                      cpu_row_idxs, repeat_times = torch.unique(
                          self.idx_map.index_select(0, ids),
                          return_counts=True,
                      )
  
                  assert len(cpu_row_idxs) <= self.cuda_row_num, (
                      f"You move {len(cpu_row_idxs)} embedding rows from CPU to CUDA. "
                      f"It is larger than the capacity of the cache, which at most contains {self.cuda_row_num} rows, "
                      f"Please increase cuda_row_num or decrease the training batch size."
                  )
                  self.evict_backlist = cpu_row_idxs
                  tmp = torch.isin(
                      cpu_row_idxs, self.cached_idx_map, invert=True
                  )
                  comm_cpu_row_idxs = cpu_row_idxs[tmp]
  
                  if self._show_cache_miss:
                      self._cache_miss += torch.sum(repeat_times[tmp])
                      self._total_cache += ids.numel()
  
          self.num_hits_history.append(
              len(cpu_row_idxs) - len(comm_cpu_row_idxs)
          )
          self.num_miss_history.append(len(comm_cpu_row_idxs))
          self.num_write_back_history.append(0)
  
          # move sure the cuda rows will not be evicted!
          with record_function("(cache) prepare_rows_on_cuda"):
              with self.timer("prepare_rows_on_cuda") as timer:
                  self._prepare_rows_on_cuda(comm_cpu_row_idxs)
  
          self.evict_backlist = torch.tensor(
              [], device=cpu_row_idxs.device, dtype=cpu_row_idxs.dtype
          )
  
          with self.timer("6_update_cache") as timer:
              with record_function("6_update_cache"):
                  gpu_row_idxs = self._id_to_cached_cuda_id(ids)
  
              # update for LFU.
              if self._evict_strategy == EvictionStrategy.LFU:
                  unique_gpu_row_idxs = self.inverted_cached_idx[
                      cpu_row_idxs
                  ]
                  self.freq_cnter.scatter_add_(
                      0, unique_gpu_row_idxs, repeat_times
                  )
  
      return gpu_row_idxs
  
imports: import sys
  from contextlib import contextmanager
  from enum import Enum
  from typing import List, Optional
  import numpy
  import torch
  from contexttimer import Timer
  from torch.profiler import record_function
  from .copyer import LimitBuffIndexCopyer
label_desc: Use image node_id 12 for calling the CachedParamMgr obj's underlying member method code with example usage: obj.prepare_ids(ids) and returns: gpu_row_idxs