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

[zero] refactor memstats collector (#706)

Browse Source 
* refactor memstats collector

* fix disposable

* polish code
This commit is contained in:
ver217
2022-04-11 10:46:08 +08:00
committed by GitHub
parent 3fc8a204dc
commit ab8c6b4a0e
8 changed files with 44 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tests/test_zero_data_parallel/test_stateful_tensor_mgr.py
View File

@@ -55,7 +55,6 @@ def run_stm():
apply_adjust(model, model.p1, [model.p1, model.p2], stateful_tensor_mgr)
mem_collector.sample_memstats()
mem_collector.finish_collection()
mem_collector.reset_sampling_cnter()
stateful_tensor_mgr.reset()
# warmup done