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

[Gemini] update the non model data record method in runtime memory tracer (#2128)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-13 17:11:31 +08:00
committed by GitHub
parent deee317b0f
commit 2938edf446
9 changed files with 65 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/test_gemini/update/test_gemini_use_rmt.py
View File

@@ -47,7 +47,13 @@ def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_
runtime_tracer_non_model_data = runtime_mem_tracer._memstats._non_model_data_cuda_list
print('runtime tracer non model data points: ', len(runtime_tracer_non_model_data))
print('runtime tracer: ', runtime_tracer_non_model_data)
print([memstats.param_used_timestep(p) for p in model.parameters()])
print([memstats.param_used_step(p) for p in model.parameters()])
if model_name == 'repeated_computed_layers':
for idx, p in enumerate(model.parameters()):
step_list = memstats.param_used_step(p)
if idx < 4:
assert len(step_list) == 4
if model_name == 'repeated_computed_layers':
for idx, p in enumerate(model.parameters()):