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

use gpc to get rank

pull/717/head
Jie Zhu 2022-04-11 17:03:57 +08:00
parent a1d7ab041d
commit 21f7bc7d73
No known key found for this signature in database
GPG Key ID: 58B7CA9DB6DC5BD9
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/utils/profiler/mem_profiler.py
View File

@ -36,7 +36,7 @@ class MemProfiler(BaseProfiler):
self._engine.remove_hook(self._mem_tracer) self._engine.remove_hook(self._mem_tracer)
def to_tensorboard(self, log_dir: Path) -> None: def to_tensorboard(self, log_dir: Path) -> None:
rank = self._mem_tracer._rank rank = gpc.get_global_rank()
stats = self._mem_tracer.async_mem_monitor.state_dict['mem_stats'] stats = self._mem_tracer.async_mem_monitor.state_dict['mem_stats']
data = { data = {
@ -47,7 +47,7 @@ class MemProfiler(BaseProfiler):
"cuda_usage": stats "cuda_usage": stats
} }
} }
rank = gpc.get_global_rank()
with open(log_dir.joinpath(f"worker{rank}.memory.json"), "w") as f: with open(log_dir.joinpath(f"worker{rank}.memory.json"), "w") as f:
json.dump(data, f) json.dump(data, f)