Revert "[zero] update sharded optim and fix zero init ctx" (#456)

* Revert "polish code"

This reverts commit 8cf7ff08cf.

* Revert "rename variables"

This reverts commit e99af94ab8.

* Revert "remove surplus imports"

This reverts commit 46add4a5c5.

* Revert "update sharded optim and fix zero init ctx"

This reverts commit 57567ee768.

tests/test_zero_data_parallel/common.py

@@ -2,10 +2,11 @@ from functools import partial
 
 import torch
 import torch.distributed as dist
+
 from colossalai.logging import get_dist_logger
 from colossalai.utils import checkpoint
-from colossalai.zero.shard_utils import TensorShardStrategy
 from colossalai.zero.sharded_model import ShardedModelV2
+from colossalai.nn.optimizer import CPUAdam
 
 LOGGER = get_dist_logger('zero_test')
 
@@ -15,18 +16,20 @@ _ZERO_MODEL_CONFIG = dict(reduce_scatter_bucket_size_mb=25,
                           fp32_reduce_scatter=False,
                           offload_config=None,
                           gradient_predivide_factor=1.0,
-                          use_memory_tracer=False,
-                          shard_strategy=TensorShardStrategy)
+                          shard_param=True,
+                          use_memory_tracer=False)
 
-_ZERO_OPTIMIZER_CONFIG = dict(cpu_offload=False,
-                              initial_scale=2**5,
-                              min_scale=1,
-                              growth_factor=2,
-                              backoff_factor=0.5,
-                              growth_interval=1000,
-                              hysteresis=2,
-                              max_scale=2**32,
-                              lr=1e-3)
+_ZERO_OPTIMIZER_CONFIG = dict(
+    optimizer_class=torch.optim.Adam,    #CPUAdam
+    cpu_offload=False,
+    initial_scale=2**5,
+    min_scale=1,
+    growth_factor=2,
+    backoff_factor=0.5,
+    growth_interval=1000,
+    hysteresis=2,
+    max_scale=2**32,
+    lr=1e-3)
 
 ZERO_PARALLEL_CONFIG = dict(fp16=dict(mode=None,),
                             zero=dict(

tests/test_zero_data_parallel/test_shard_model_v2.py

@@ -1,13 +1,15 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
+import copy
+from asyncio.log import logger
 from functools import partial
 
 import colossalai
 import pytest
 import torch
 import torch.multiprocessing as mp
-from colossalai.testing import parameterize
+from colossalai.logging import get_dist_logger
 from colossalai.utils import free_port
 from colossalai.zero.init_ctx import ZeroInitContext
 from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
@@ -18,30 +20,36 @@ from tests.components_to_test.registry import non_distributed_component_funcs
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from common import CONFIG, check_grads_padding, run_fwd_bwd
+from colossalai.testing import parameterize
 
 
 @parameterize("enable_autocast", [True])
-@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
-def run_model_test(enable_autocast, shard_strategy_class):
+@parameterize("use_zero_init_ctx", [True])
+@parameterize("shard_strategy", [TensorShardStrategy, BucketTensorShardStrategy])
+def run_model_test(enable_autocast, use_zero_init_ctx, shard_strategy, logger):
     test_models = ['repeated_computed_layers', 'resnet18', 'bert']
-    shard_strategy = shard_strategy_class()
+    shard_strategy = shard_strategy()
     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()
 
         rm_torch_payload_on_the_fly = False
 
-        with ZeroInitContext(convert_fp16=True,
-                             target_device=torch.cuda.current_device(),
-                             shard_strategy=shard_strategy,
-                             shard_param=True,
-                             rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly):
-            zero_model = model_builder(checkpoint=True)
-        zero_model = ShardedModelV2(zero_model, shard_strategy, use_memory_tracer=True)
+        if use_zero_init_ctx:
+            with ZeroInitContext(convert_fp16=True,
+                                 target_device=torch.device(f'cpu:0'),
+                                 shard_strategy=shard_strategy,
+                                 shard_param=True,
+                                 rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly):
+                zero_model = model_builder(checkpoint=True)
+            zero_model = ShardedModelV2(zero_model, shard_strategy, use_memory_tracer=True)
 
-        model = model_builder(checkpoint=True).half()
-        col_model_deepcopy(zero_model, model)
-        model = model.cuda()
+            model = model_builder(checkpoint=True).half()
+            col_model_deepcopy(zero_model, model)
+            model = model.cuda()
+        else:
+            model = model_builder(checkpoint=True).half().cuda()
+            zero_model = ShardedModelV2(copy.deepcopy(model), shard_strategy)
 
         model = DDP(model)
 
@@ -55,10 +63,15 @@ def run_model_test(enable_autocast, shard_strategy_class):
 
             check_grads_padding(model, zero_model, loose=True)
 
+        # logger.debug('overall cuda ', zero_model._memstats_collector._overall_cuda)
+        # logger.debug('model cuda ', zero_model._memstats_collector._model_data_cuda)
+
 
 def run_dist(rank, world_size, port):
     colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
-    run_model_test()
+    logger = get_dist_logger()
+    logger.set_level('DEBUG')
+    run_model_test(logger=logger)
 
 
 @pytest.mark.dist

tests/test_zero_data_parallel/test_sharded_optim_v2.py

@@ -1,3 +1,4 @@
+import copy
 from functools import partial
 
 import colossalai
@@ -5,18 +6,15 @@ import pytest
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from colossalai.nn.optimizer import CPUAdam
-from colossalai.testing import parameterize
 from colossalai.utils import free_port
-from colossalai.zero.init_ctx import ZeroInitContext
 from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
 from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_model.utils import col_model_deepcopy
 from colossalai.zero.sharded_optim import ShardedOptimizerV2
-from colossalai.zero.sharded_optim._utils import has_inf_or_nan
 from tests.components_to_test.registry import non_distributed_component_funcs
 from torch.nn.parallel import DistributedDataParallel as DDP
-
+from colossalai.nn.optimizer import CPUAdam
+from colossalai.zero.sharded_optim._utils import has_inf_or_nan
+from colossalai.testing import parameterize
 from common import CONFIG, check_sharded_params_padding
 
 
@@ -40,42 +38,36 @@ def _run_step(model, optimizer, data, label, criterion, enable_autocast=False):
 
 @parameterize("cpu_offload", [True, False])
 @parameterize("use_cpuadam", [True, False])
-@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
-def _run_test_sharded_optim_v2(cpu_offload, shard_strategy_class, use_cpuadam):
+@parameterize("shard_strategy", [TensorShardStrategy, BucketTensorShardStrategy])
+def _run_test_sharded_optim_v2(cpu_offload, shard_strategy, use_cpuadam):
     test_models = ['repeated_computed_layers', 'resnet18', 'bert']
-    shard_strategy = shard_strategy_class()
+    shard_strategy = shard_strategy()
 
     if use_cpuadam and cpu_offload is False:
         return
 
     for model_name in test_models:
         get_components_func = non_distributed_component_funcs.get_callable(model_name)
-        model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
-
-        with ZeroInitContext(convert_fp16=True,
-                             target_device=torch.device(f'cpu:0'),
-                             shard_strategy=shard_strategy,
-                             shard_param=True,
-                             rm_torch_payload_on_the_fly=False):
-            zero_model = model_builder(checkpoint=True)
-        zero_model = ShardedModelV2(zero_model,
+        model, train_dataloader, _, optimizer_class, criterion = get_components_func()
+        model = model(checkpoint=True).cuda()
+        zero_model = ShardedModelV2(copy.deepcopy(model),
                                     shard_strategy,
                                     offload_config=dict(device='cpu') if cpu_offload else None)
-
-        model = model_builder(checkpoint=True).half()
-        col_model_deepcopy(zero_model, model)
-        model = model.cuda().float()
         if dist.get_world_size() > 1:
             model = DDP(model)
-
+        lr = 1e-3
         if use_cpuadam:
-            optimizer_class = CPUAdam
-        optim = optimizer_class(model.parameters(), lr=1e-3)
-        sharded_optim = optimizer_class(zero_model.parameters(), lr=1e-3)
-        sharded_optim = ShardedOptimizerV2(zero_model, sharded_optim, cpu_offload=cpu_offload, initial_scale=2**5)
-
+            optim = torch.optim.Adam(model.parameters(), lr=lr)
+            sharded_optim = ShardedOptimizerV2(zero_model, CPUAdam, cpu_offload=cpu_offload, initial_scale=2**5, lr=lr)
+        else:
+            optim = optimizer_class(model.parameters(), lr=lr)
+            sharded_optim = ShardedOptimizerV2(zero_model,
+                                               optimizer_class,
+                                               cpu_offload=cpu_offload,
+                                               initial_scale=2**5,
+                                               lr=lr)
         for i, (data, label) in enumerate(train_dataloader):
-            # FIXME() if i > 5, the unittest will fail
+            #FIXME() if i > 5, the unittest will fail
             if i > 3:
                 break
             data, label = data.cuda(), label.cuda()

tests/test_zero_data_parallel/test_sharded_optim_with_sync_bn.py

@@ -6,12 +6,12 @@ from functools import partial
 import colossalai
 import pytest
 import torch
-import torch.distributed as dist
 import torch.multiprocessing as mp
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
 from colossalai.utils import free_port
+from colossalai.core import global_context as gpc
+from colossalai.context.parallel_mode import ParallelMode
 from torchvision.models import resnet50
+import torch.distributed as dist
 
 
 def run_dist(rank, world_size, port):
@@ -64,10 +64,6 @@ def run_dist(rank, world_size, port):
         'expected the output from different ranks to be the same, but got different values'
 
 
-# FIXME: enable this test in next PR
-
-
-@pytest.mark.skip
 @pytest.mark.dist
 def test_sharded_optim_with_sync_bn():
     """

tests/test_zero_data_parallel/test_state_dict.py

@@ -8,37 +8,24 @@ import colossalai
 import pytest
 import torch
 import torch.multiprocessing as mp
-from colossalai.testing import parameterize
 from colossalai.utils import free_port
-from colossalai.zero.init_ctx import ZeroInitContext
 from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
 from colossalai.zero.sharded_model import ShardedModelV2
-from colossalai.zero.sharded_model.utils import col_model_deepcopy
 from tests.components_to_test.registry import non_distributed_component_funcs
-
+from colossalai.testing import parameterize
 from common import CONFIG
 
 
-@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
-def run_zero_state_dict(shard_strategy_class):
+@parameterize("shard_strategy", [TensorShardStrategy, BucketTensorShardStrategy])
+def run_zero_state_dict(shard_strategy):
     test_models = ['repeated_computed_layers', 'resnet18']
-    shard_strategy = shard_strategy_class()
+    shard_strategy = shard_strategy()
     for model_name in test_models:
         get_components_func = non_distributed_component_funcs.get_callable(model_name)
         model_builder, train_dataloader, test_dataloader, optimizer, criterion = get_components_func()
-
-        with ZeroInitContext(convert_fp16=True,
-                             target_device=torch.cuda.current_device(),
-                             shard_strategy=shard_strategy,
-                             shard_param=True,
-                             rm_torch_payload_on_the_fly=False):
-            zero_model = model_builder(checkpoint=True)
-        zero_model = ShardedModelV2(zero_model, shard_strategy)
-
-        model = model_builder(checkpoint=True).half()
-        col_model_deepcopy(zero_model, model)
-        model = model.cuda()
-
+        model = model_builder()
+        model = model.half().cuda()
+        zero_model = ShardedModelV2(deepcopy(model), shard_strategy)
         zero_state_dict = zero_model.state_dict()
         for key, val in model.state_dict().items():
             assert torch.equal(val, zero_state_dict[key])

tests/test_zero_data_parallel/test_zero_engine.py

@@ -1,24 +1,21 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
+import copy
 from functools import partial
+from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
+import pytest
 
 import colossalai
-import pytest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
 from colossalai.utils import free_port
-from colossalai.zero.init_ctx import ZeroInitContext
-from colossalai.zero.sharded_model.utils import col_model_deepcopy
 from colossalai.zero.sharded_optim._utils import has_inf_or_nan
-from tests.components_to_test.registry import non_distributed_component_funcs
+
+import torch.multiprocessing as mp
+import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from common import (MP_PARALLEL_CONFIG, ZERO_PARALLEL_CONFIG, check_params,
-                    check_sharded_params_padding)
+from tests.components_to_test.registry import non_distributed_component_funcs
+from common import check_sharded_params_padding, ZERO_PARALLEL_CONFIG, MP_PARALLEL_CONFIG, check_params
 
 
 def run_dist(rank, world_size, port, parallel_config):
@@ -33,16 +30,10 @@ def run_dist(rank, world_size, port, parallel_config):
     for model_name in test_models:
         get_components_func = non_distributed_component_funcs.get_callable(model_name)
         model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
-        with ZeroInitContext(convert_fp16=hasattr(gpc.config, 'fp16'),
-                             target_device=torch.cuda.current_device(),
-                             shard_strategy=gpc.config.zero.model_config.shared_strategy(
-                                 gpc.get_group(ParallelMode.DATA)),
-                             shard_param=True):
-            colo_model = model_builder(checkpoint=True)
 
-        torch_model = model_builder(checkpoint=True).half()
-        col_model_deepcopy(colo_model, torch_model)
-        torch_model = torch_model.cuda().float()
+        colo_model = model_builder(checkpoint=True)
+        torch_model = copy.deepcopy(colo_model).cuda()
+        torch_model.train()
         engine, train_dataloader, _, _ = colossalai.initialize(colo_model,
                                                                optimizer=optimizer_class,
                                                                criterion=criterion,
@@ -91,10 +82,6 @@ def run_dist(rank, world_size, port, parallel_config):
             check_sharded_params_padding(torch_model, colo_model, loose=True)
 
 
-# FIXME: enable this test in next PR
-
-
-@pytest.mark.skip
 @pytest.mark.dist
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_mp_engine(world_size):
@@ -102,7 +89,6 @@ def test_mp_engine(world_size):
     mp.spawn(run_func, nprocs=world_size)
 
 
-@pytest.mark.skip
 @pytest.mark.dist
 @pytest.mark.parametrize("world_size", [1, 2])
 def test_zero_engine(world_size):