[Gemini] remove GLOBAL_MODEL_DATA_TRACER (#2091)

colossalai/gemini/memory_tracer/__init__.py

@@ -1,11 +1,10 @@
 from .memory_monitor import AsyncMemoryMonitor, SyncCudaMemoryMonitor    # isort:skip
 from .memstats_collector import MemStatsCollector    # isort:skip
-from .model_data_memtracer import GLOBAL_MODEL_DATA_TRACER    # isort:skip
 from .chunk_memstats_collector import ChunkMemStatsCollector    # isort:skip
 from .static_memstats_collector import StaticMemStatsCollector    # isort:skip
 from .memory_stats import MemStats
 
 __all__ = [
     'AsyncMemoryMonitor', 'SyncCudaMemoryMonitor', 'MemStatsCollector', 'ChunkMemStatsCollector',
-    'StaticMemStatsCollector', 'GLOBAL_MODEL_DATA_TRACER', 'MemStats'
+    'StaticMemStatsCollector', 'MemStats'
 ]

colossalai/gemini/memory_tracer/utils.py

@@ -2,9 +2,6 @@ from typing import Optional, Tuple
 
 import torch
 
-from colossalai.context.singleton_meta import SingletonMeta
-from colossalai.logging import DistributedLogger
-
 
 def colo_model_optimizer_usage(optim) -> Tuple[int, int]:
     """Trace the optimizer memory usage
@@ -60,52 +57,3 @@ def colo_model_mem_usage(model: torch.nn.Module) -> Tuple[int, int]:
             cpu_mem_usage += t_cpu
 
     return cuda_mem_usage, cpu_mem_usage
-
-
-class ModelDataTracer(metaclass=SingletonMeta):
-    """
-    A tracer singleton to trace model data usage during runtime.
-    You have to register a model on the singleton first.
-    """
-
-    def __init__(self) -> None:
-        self._logger = DistributedLogger("ModelDataTracer")
-        self._model = None
-        self._opitimizer = None
-
-    def _get_mem_usage(self) -> Tuple[int, int]:
-        """
-        get the memory usage of the model registered.
-        Returns:
-            Tuple[int, int]: cuda, cpu mem usage
-        """
-        cuda_use_opt, cpu_use_opt = colo_model_optimizer_usage(self._opitimizer)
-        cuda_use_model, cpu_use_model = colo_model_mem_usage(self._model)
-        return cuda_use_opt + cuda_use_model, cpu_use_opt + cpu_use_model
-
-    def register_model(self, model) -> None:
-        if self._model is not None:
-            self._logger.warning("ModelDataTracer has already registered a model")
-        self._model = model
-
-    def register_optimizer(self, optimizer) -> None:
-        if self._opitimizer is not None:
-            self._logger.warning("ModelDataTracer has already registered an optimizer")
-        self._opitimizer = optimizer
-
-    @property
-    def cpu_usage(self):
-        _, cpu_usage = self._get_mem_usage()
-        return cpu_usage
-
-    @property
-    def cuda_usage(self):
-        cuda_usage, _ = self._get_mem_usage()
-        return cuda_usage
-
-    @property
-    def both_mem_usage(self):
-        return self._get_mem_usage()
-
-
-GLOBAL_MODEL_DATA_TRACER = ModelDataTracer()