docstr: Store the non model data statistics used for Gemini and GeminiOptimizer.
docstr_tok: ["Store","the","non","model","data","statistics","used","for","Gemini","and","GeminiOptimizer","."]
code: def __init__(self) -> None:
      """
      Store the non model data statistics used for Gemini and GeminiOptimizer.
      """
      # (preop_step, List[param])
      self._step_param_dict = dict()
      # (param, List[preop_step])
      self._param_step_dict = dict()
      # (preop_step, non_model_data) non model data used during preop_step ~ (preop_step+1)
      self._step_nmd_dict = dict()
      self._param_runtime_order = OrderedParamGenerator()
  
      self._preop_step = 0
  
      self._prev_overall_cuda = -1
      self._max_overall_cuda = 0
      self._prev_md_cuda = -1
  
      # old version
      self._model_data_cuda_list = []
      self._model_data_cpu_list = []
  
      self._overall_cuda_list = []
      self._overall_cpu_list = []
  
      self._non_model_data_cuda_list = []
      self._non_model_data_cpu_list = []
  
file_path: ColossalAI/colossalai/zero/gemini/memory_tracer/memory_stats.py
imports: from typing import List, Optional
  import torch
  from .param_runtime_order import OrderedParamGenerator
label_desc: Use image node_id 1 to create a new MemStats object from inherited base classes: object with example: obj = MemStats()