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

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)