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

[Gemini] open grad checkpoint when model building (#1984)

pull/1987/head^2
Jiarui Fang 2022-11-18 16:32:54 +08:00 committed by GitHub
parent c26f21d365
commit 5bec3b2168
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tests/test_gemini/test_mem_tracer.py
View File

@ -22,9 +22,10 @@ def run_fwd_bwd(model, data, label, criterion, enable_autocast=False):
model.backward(loss)
def run_tracer(rank, world_size, port, grad_check=True):
def run_tracer(rank, world_size, port, use_grad_check=True):
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
test_models = ['repeated_computed_layers', 'resnet18', 'no_leaf_module', 'bert']
# test_models = ['bert']
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, _, _, criterion = get_components_func()
@ -32,7 +33,7 @@ def run_tracer(rank, world_size, port, grad_check=True):
# init model on cpu
# TODO() memtrace hook can not handle buff registered on a non-leaf module (for example the BertEmbedding).
# a simple method is that always puts buff on cuda and viewed them as non-model data.
model = MemtracerWrapper(model_builder(grad_check))
model = MemtracerWrapper(model_builder(checkpoint=use_grad_check))
for n, buff in model.named_buffers():
buff.data = buff.data.cuda()
@ -44,14 +45,15 @@ def run_tracer(rank, world_size, port, grad_check=True):
run_fwd_bwd(model, data, label, criterion, False)
# model._ophook_list[0].print_non_model_data()
model._ophook_list[0].print_non_model_data()
@pytest.mark.dist
@pytest.mark.parametrize("world_size", [1])
@pytest.mark.parametrize("use_grad_check", [True, False])
@rerun_if_address_is_in_use()
def test_tracer(world_size):
run_func = partial(run_tracer, world_size=world_size, port=free_port())
def test_tracer(world_size, use_grad_check):
run_func = partial(run_tracer, world_size=world_size, port=free_port(), use_grad_check=use_grad_check)
mp.spawn(run_func, nprocs=world_size)