From 5bec3b21683f97defd7e38f1d68a86bb7141fc0f Mon Sep 17 00:00:00 2001 From: Jiarui Fang Date: Fri, 18 Nov 2022 16:32:54 +0800 Subject: [PATCH] [Gemini] open grad checkpoint when model building (#1984) --- tests/test_gemini/test_mem_tracer.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_gemini/test_mem_tracer.py b/tests/test_gemini/test_mem_tracer.py index 05da462a4..7e524765b 100644 --- a/tests/test_gemini/test_mem_tracer.py +++ b/tests/test_gemini/test_mem_tracer.py @@ -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)