github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Gemini] revert ZeROInitCtx related tracer (#2138)

pull/2139/head^2
Jiarui Fang 2022-12-16 12:37:06 +08:00 committed by GitHub
parent 077a66dd81
commit ee287620f0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/gemini/memory_tracer/memstats_collector.py
View File

@ -73,10 +73,15 @@ class MemStatsCollector:
# deprecated
def record_model_data_volume(self) -> None:
"""Sampling model data statistics.
"""
Sampling model data statistics.
"""
if self._start_flag and not self.use_outside_memstats:
raise NotImplementedError("MemStatsCollector has not implemented record_model_data_volume")
# The following code work for ZeroInitContext, which is deprecated in v0.1.12
cuda_mem = StatefulTensor.GST_MGR.total_mem['cuda']
cpu_mem = StatefulTensor.GST_MGR.total_mem['cpu']
self._memstats.append_model_data('cuda', cuda_mem)
self._memstats.append_model_data('cpu', cpu_mem)
def sample_overall_data(self) -> None:
"""