github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-16 06:30:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] refactor memstats_collector (#746)

Browse Source
This commit is contained in:
HELSON
2022-04-14 12:01:12 +08:00
committed by GitHub
parent b8899e0905
commit 84c6700b2a
10 changed files with 179 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
colossalai/utils/memory_tracer/memstats_collector.py
View File

@@ -1,6 +1,6 @@
from colossalai.utils.memory_tracer.model_data_memtracer import GLOBAL_MODEL_DATA_TRACER
from colossalai.utils.memory import colo_device_memory_used
from colossalai.utils.memory_tracer import AsyncMemoryMonitor
from colossalai.utils.memory_tracer import SyncCudaMemoryMonitor
import torch
import time
from typing import List
@@ -19,7 +19,7 @@ class MemStatsCollector:
"""
def __init__(self) -> None:
self._mem_monitor = AsyncMemoryMonitor()
self._mem_monitor = SyncCudaMemoryMonitor()
self._model_data_cuda_list = []
self._overall_cuda_list = []
@@ -31,9 +31,10 @@ class MemStatsCollector:
self._sampling_time = []
self._start_flag = False
self._period_idx = 0
self._step_idx = 0
self._step_total = 0
def overall_mem_stats(self, device_type: str):
def overall_mem_stats(self, device_type: str) -> List[int]:
if device_type == 'cuda':
return self._overall_cuda_list
elif device_type == 'cpu':
@@ -41,47 +42,23 @@ class MemStatsCollector:
else:
raise TypeError
def model_data_list(self, device_type: str, unit: str = 'B') -> List[int]:
if unit == 'GB':
scale = 1e9
elif unit == 'MB':
scale = 1e6
elif unit == 'KB':
scale = 1e3
elif unit == 'B':
scale = 1
else:
raise TypeError
def model_data_list(self, device_type: str) -> List[int]:
if device_type == 'cuda':
return [elem / scale for elem in self._model_data_cuda_list]
return self._model_data_cuda_list
elif device_type == 'cpu':
return [elem / scale for elem in self._model_data_cpu_list]
else:
raise TypeError
def non_model_data_list(self, device_type: str, unit: str = 'B') -> List[int]:
"""Non model data stats
"""
if unit == 'GB':
scale = 1e9
elif unit == 'MB':
scale = 1e6
elif unit == 'KB':
scale = 1e3
elif unit == 'B':
scale = 1
return self._model_data_cpu_list
else:
raise TypeError
def non_model_data_list(self, device_type: str) -> List[int]:
if device_type == 'cuda':
return [elem / scale for elem in self._non_model_data_cuda_list]
return self._non_model_data_cuda_list
elif device_type == 'cpu':
return [elem / scale for elem in self._non_model_data_cpu_list]
return self._non_model_data_cpu_list
else:
raise TypeError
def max_non_model_data(self, device_type: str) -> int:
def next_period_non_model_data_usage(self, device_type: str) -> int:
"""Get max non model data memory usage of current sampling period
Args:
@@ -91,12 +68,10 @@ class MemStatsCollector:
int: max non model data memory usage of current sampling period
"""
assert not self._start_flag, 'Cannot get mem stats info during collection phase.'
assert len(self._sampling_time) > 0, 'Cannot get mem stats info before collection phase.'
next_period_idx = (self._period_idx + 1) % len(self._sampling_time)
current_non_model_data = self.non_model_data_list(device_type)[self._period_idx]
next_non_model_data = self.non_model_data_list(device_type)[next_period_idx]
self._period_idx = next_period_idx
return max(current_non_model_data, next_non_model_data)
assert self._step_total > 0, 'Cannot get mem stats info before collection phase.'
next_non_model_data = self.non_model_data_list(device_type)[self._step_idx]
self._step_idx = (self._step_idx + 1) % self._step_total
return next_non_model_data
@property
def sampling_time(self):
@@ -107,9 +82,37 @@ class MemStatsCollector:
self._mem_monitor.start()
def finish_collection(self):
self.sample_overall_data()
self._step_total = len(self._sampling_time)
self._start_flag = False
self._mem_monitor.finish()
def sample_model_data(self) -> None:
"""Sampling model data statistics.
"""
if self._start_flag:
cuda_mem, cpu_mem = GLOBAL_MODEL_DATA_TRACER.both_mem_usage
self._model_data_cuda_list.append(cuda_mem)
self._model_data_cpu_list.append(cpu_mem)
def sample_overall_data(self) -> None:
"""Sampling non model data statistics.
"""
if self._start_flag:
# overall data recording is after model data recording
if len(self._model_data_cuda_list) == 0:
return
self._overall_cuda_list.append(self._mem_monitor.finish())
self._overall_cpu_list.append(colo_device_memory_used(torch.device('cpu')))
assert len(self._model_data_cuda_list) == len(self._overall_cuda_list)
self._non_model_data_cuda_list.append(self._overall_cuda_list[-1] - self._model_data_cuda_list[-1])
self._non_model_data_cpu_list.append(self._overall_cpu_list[-1] - self._model_data_cpu_list[-1])
self._sampling_time.append(time.time())
self._mem_monitor.start()
def sample_memstats(self) -> None:
"""
Sampling memory statistics.
@@ -119,7 +122,7 @@ class MemStatsCollector:
if self._start_flag:
self._model_data_cuda_list.append(GLOBAL_MODEL_DATA_TRACER.cuda_usage)
self._overall_cuda_list.append(self._mem_monitor.finish())
self._non_model_data_cuda_list.append(self._model_data_cuda_list[-1] - self._overall_cuda_list[-1])
self._non_model_data_cuda_list.append(self._overall_cuda_list[-1] - self._model_data_cuda_list[-1])
self._model_data_cpu_list.append(GLOBAL_MODEL_DATA_TRACER.cpu_usage)
# FIXME(jiaruifang) cpu sys used should also return from self._mem_monitor()
@@ -136,4 +139,5 @@ class MemStatsCollector:
self._overall_cpu_list = []
self._start_flag = False
self._period_idx = 0
self._step_idx = 0
self._step_total = 0

4
colossalai/utils/memory_tracer/model_data_memtracer.py
View File

@@ -101,5 +101,9 @@ class ModelDataTracer(metaclass=SingletonMeta):
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()