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

[Gemini] rename hooks related to runtime mem tracer (#2076)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-05 15:00:03 +08:00
committed by GitHub
parent 40b7d55bf3
commit a7adad9ccb
5 changed files with 15 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_gemini/test_runtime_mem_tracer.py
View File

@@ -29,7 +29,7 @@ def test_runtime_mem_tracer():
model_builder, train_dataloader, _, _, criterion = get_components_func()
with ColoInitContext(device=torch.device('cpu')):
model = model_builder(checkpoint=True)
model = model_builder(checkpoint=False)
model_bk = deepcopy(model)
runtime_mem_tracer = RuntimeMemTracer(model)
@@ -47,7 +47,7 @@ def test_runtime_mem_tracer():
cuda_non_model_data_list = np.array(GLOBAL_CUDA_MEM_INFO.non_model_data_list) / 1024**2
print("cuda_non_model_data_list", len(cuda_non_model_data_list))
# print(GLOBAL_CUDA_MEM_INFO.non_model_data_list)
print(GLOBAL_CUDA_MEM_INFO.non_model_data_list)
del model