Skip to content

GitLab

Unverified Commit ee287620 authored by Jiarui Fang's avatar Jiarui Fang Committed by GitHub
Browse files

[Gemini] revert ZeROInitCtx related tracer (#2138)

parent 077a66dd
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 2 deletions
colossalai/gemini/memory_tracer/memstats_collector.py
View file @ ee287620
...@@ -73,10 +73,15 @@ class MemStatsCollector: ...@@ -73,10 +73,15 @@ class MemStatsCollector:
# deprecated # deprecated
def record_model_data_volume(self) -> None: 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: 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: def sample_overall_data(self) -> None:
""" """
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment