[legacy] clean up legacy code (#4743)

* [legacy] remove outdated codes of pipeline (#4692)

* [legacy] remove cli of benchmark and update optim (#4690)

* [legacy] remove cli of benchmark and update optim

* [doc] fix cli doc test

* [legacy] fix engine clip grad norm

* [legacy] remove outdated colo tensor (#4694)

* [legacy] remove outdated colo tensor

* [test] fix test import

* [legacy] move outdated zero to legacy (#4696)

* [legacy] clean up utils (#4700)

* [legacy] clean up utils

* [example] update examples

* [legacy] clean up amp

* [legacy] fix amp module

* [legacy] clean up gpc (#4742)

* [legacy] clean up context

* [legacy] clean core, constants and global vars

* [legacy] refactor initialize

* [example] fix examples ci

* [example] fix examples ci

* [legacy] fix tests

* [example] fix gpt example

* [example] fix examples ci

* [devops] fix ci installation

* [example] fix examples ci

tests/test_utils/test_commons.py → tests/test_legacy/test_zero/test_commons.py

 import torch
 
 import colossalai
+from colossalai.legacy.zero.gemini.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
+from colossalai.legacy.zero.sharded_param import ShardedTensor
 from colossalai.testing import rerun_if_address_is_in_use, spawn
-from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
-from colossalai.zero.legacy.sharded_param import ShardedTensor
 
 
 def run_tensor_move(rank, world_size, port):
-    colossalai.launch(config={}, rank=0, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.legacy.launch(config={}, rank=0, world_size=world_size, host='localhost', port=port, backend='nccl')
 
     src_t = torch.ones(2, 3).cuda()
     tgt_t = torch.zeros(2, 3)

tests/test_moe/test_kernel.py

@@ -3,9 +3,9 @@ import torch
 import torch.nn as nn
 
 import colossalai
-from colossalai.context import ParallelMode
 from colossalai.context.moe_context import MOE_CONTEXT
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.nn.layer.moe import Experts, MoeLayer, Top1Router, Top2Router
 from colossalai.testing import rerun_if_address_is_in_use, spawn
 from colossalai.utils import get_current_device

tests/test_moe/test_moe_zero_optim.py

@@ -2,8 +2,8 @@ import pytest
 import torch
 
 import colossalai
-from colossalai.amp import convert_to_apex_amp
 from colossalai.context import MOE_CONTEXT
+from colossalai.legacy.amp import convert_to_apex_amp
 from colossalai.legacy.engine.gradient_handler import MoeGradientHandler
 from colossalai.nn import MoeLoss
 from colossalai.nn.optimizer import CPUAdam

tests/test_tensor/test_comm_spec_apply.py

 import pytest
 import torch
+import torch.distributed as dist
 
-from colossalai.core import global_context as gpc
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
@@ -184,7 +184,7 @@ def check_comm(rank, world_size, port):
     launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
 
     physical_mesh_id = torch.arange(0, 4)
-    assert rank == gpc.get_global_rank()
+    assert rank == dist.get_rank()
 
     mesh_shape = (2, 2)
     # [[0, 1,
@@ -205,7 +205,6 @@ def check_comm(rank, world_size, port):
 
     # test all reduce in 1D flatten device mesh
     check_all_reduce_in_flatten_device_mesh(device_mesh, rank)
-    gpc.destroy()
 
 
 @pytest.mark.dist

tests/test_tensor/test_dtensor/test_comm_spec.py

 import pytest
 import torch
+import torch.distributed as dist
 
-from colossalai.core import global_context as gpc
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
@@ -127,7 +127,7 @@ def check_comm(rank, world_size, port):
     launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
 
     physical_mesh_id = torch.arange(0, 4)
-    assert rank == gpc.get_global_rank()
+    assert rank == dist.get_rank()
 
     mesh_shape = (2, 2)
     # [[0, 1,
@@ -149,8 +149,6 @@ def check_comm(rank, world_size, port):
     check_all_reduce_fwd(process_group_dict, rank)
     check_all_reduce_bwd(process_group_dict, rank)
 
-    gpc.destroy()
-
 
 @pytest.mark.dist
 @rerun_if_address_is_in_use()

tests/test_tensor/test_mix_gather.py

 import pytest
 import torch
+import torch.distributed as dist
 
-from colossalai.core import global_context as gpc
 from colossalai.device.device_mesh import DeviceMesh
 from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
@@ -295,7 +295,7 @@ def check_comm(rank, world_size, port):
     launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
 
     physical_mesh_id = torch.arange(0, 8)
-    assert rank == gpc.get_global_rank()
+    assert rank == dist.get_rank()
 
     mesh_shape = (2, 4)
     # [[0, 1, 2, 3],

tests/test_utils/test_zero_gradient_clippling.py

-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-from functools import partial
-
-import pytest
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.nn.utils import clip_grad_norm_
-
-import colossalai
-from colossalai.logging import disable_existing_loggers
-from colossalai.testing import rerun_if_address_is_in_use, spawn
-from colossalai.utils import checkpoint, clip_grad_norm_fp32
-from colossalai.zero.legacy.shard_utils.tensor_shard_strategy import TensorShardStrategy
-from colossalai.zero.legacy.sharded_model.sharded_model_v2 import ShardedModelV2
-
-
-def checkpoint_wrapper(module, enable=True):
-    if enable:
-        module.forward = partial(checkpoint, module.forward, False)
-    return module
-
-
-class Net(nn.Module):
-
-    def __init__(self, checkpoint=False) -> None:
-        super().__init__()
-        self.fc1 = nn.Linear(5, 5)
-        self.fc2 = nn.Linear(5, 5)
-        self.fc3 = nn.Linear(5, 1)
-        if checkpoint:
-            self.fc1 = checkpoint_wrapper(self.fc1)
-        self.layers = [self.fc1, self.fc2, self.fc1, self.fc2, self.fc3]
-
-    def forward(self, x):
-        for layer in self.layers:
-            x = layer(x)
-        return x
-
-
-def run_step(model, optimizer, x, enable_autocast=False, norm_type=2.0):
-    model.train()
-    optimizer.zero_grad()
-    with torch.cuda.amp.autocast(enabled=enable_autocast):
-        y = model(x)
-        loss = y.sum()
-    loss = loss.float()
-    loss.backward()
-    clip_grad(model, norm_type)
-    optimizer.step()
-
-
-def clip_grad(model, norm_type):
-    if isinstance(model, DDP):
-        clip_grad_norm_(model.parameters(), max_norm=1.0, norm_type=norm_type)
-    else:
-        clip_grad_norm_fp32(model.parameters(), max_norm=1.0, norm_type=norm_type)
-
-
-def allclose(tensor_a: torch.Tensor, tensor_b: torch.Tensor, loose=False) -> bool:
-    if loose:
-        return torch.allclose(tensor_a, tensor_b, atol=1e-3, rtol=1e-3)
-    return torch.allclose(tensor_a, tensor_b)
-
-
-def check_grads(model, zero_model, loose=False):
-    rank = dist.get_rank()
-    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
-        zero_grad = zero_p.grad.clone().to(p.device)
-        chunks = torch.flatten(p.grad).chunk(4)
-        if rank >= len(chunks):
-            continue
-        grad = chunks[rank]
-        if zero_p.zero_shard_padding > 0:
-            zero_grad = zero_grad[:-zero_p.zero_shard_padding]
-        assert grad.dtype == zero_grad.dtype
-        assert allclose(grad, zero_grad, loose=loose)
-
-
-def check_params(model, zero_model, loose=False):
-    rank = dist.get_rank()
-    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
-        zero_shard_padding = zero_p.zero_shard_padding
-        zero_p = zero_p.clone().to(p.device)
-        chunks = torch.flatten(p).chunk(4)
-        if rank >= len(chunks):
-            continue
-        p = chunks[rank]
-        if zero_shard_padding > 0:
-            zero_p = zero_p[:-zero_shard_padding]
-        assert p.dtype == zero_p.dtype
-        assert allclose(p, zero_p, loose=loose)
-
-
-def run_dist(rank, world_size, port):
-    disable_existing_loggers()
-    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
-
-
-@pytest.mark.dist
-@rerun_if_address_is_in_use()
-def test_zero_clip_grad():
-    world_size = 4
-    spawn(run_dist, world_size)
-
-
-if __name__ == '__main__':
-    test_zero_clip_grad()

tests/test_zero/test_gemini/test_chunk_mgrv2.py

@@ -6,7 +6,6 @@ import colossalai
 from colossalai.tensor import ColoTensor
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
 from colossalai.zero.gemini.chunk import ChunkManager
-from tests.test_tensor.common_utils import debug_print
 
 CUDA_MEM_0 = {False: 512, True: 1024}
 CUDA_MEM_1 = {False: 0, True: 1024}
@@ -16,7 +15,6 @@ CPU_MEM = {True: {True: 0, False: 0}, False: {True: 512, False: 0}}
 @parameterize('keep_gathered', [True, False])
 @parameterize('pin_memory', [True, False])
 def exam_chunk_memory(keep_gathered, pin_memory):
-    debug_print([0], "keep_gathered: {}, pin_memory: {}".format(keep_gathered, pin_memory))
 
     params = [ColoTensor(torch.rand(8, 8)) for _ in range(3)]
     config = {2: dict(chunk_size=128, keep_gathered=keep_gathered)}

tests/test_zero/test_gemini/test_fwd_bwd.py

@@ -5,15 +5,15 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.testing import assert_close
 
 import colossalai
-from colossalai.amp import convert_to_apex_amp
+from colossalai.legacy.amp import convert_to_apex_amp
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.utils.cuda import get_current_device
 from colossalai.zero import GeminiDDP, GeminiOptimizer
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 PLACEMENT_CONFIGS = [
     {

tests/test_zero/test_gemini/test_gemini_use_rmt.py

@@ -4,12 +4,12 @@ import torch.distributed as dist
 
 import colossalai
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.zero import GeminiDDP
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from colossalai.zero.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 # run gemini use the runtime memory tracer
 

tests/test_zero/test_gemini/test_grad_clip.py

@@ -5,14 +5,14 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.testing import assert_close
 
 import colossalai
-from colossalai.amp import convert_to_apex_amp
+from colossalai.legacy.amp import convert_to_apex_amp
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.zero import GeminiDDP, GeminiOptimizer
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 PLACEMENT_CONFIGS = [
     {

tests/test_zero/test_gemini/test_inference.py

@@ -7,15 +7,15 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.testing import assert_close
 
 import colossalai
-from colossalai.amp import convert_to_apex_amp
+from colossalai.legacy.amp import convert_to_apex_amp
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.utils.cuda import get_current_device
 from colossalai.zero import GeminiDDP, GeminiOptimizer
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 PLACEMENT_CONFIGS = [
     {

tests/test_zero/test_gemini/test_optim.py

@@ -5,15 +5,15 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.testing import assert_close
 
 import colossalai
-from colossalai.amp import convert_to_apex_amp
+from colossalai.legacy.amp import convert_to_apex_amp
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.utils.cuda import get_current_device
 from colossalai.zero import GeminiDDP, GeminiOptimizer
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from tests.components_to_test import run_fwd_bwd
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 PLACEMENT_CONFIGS = [
     {

tests/test_zero/test_gemini/test_zeroddp_state_dict.py

@@ -4,10 +4,10 @@ from torch.testing import assert_close
 
 import colossalai
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.zero import GeminiDDP
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 PLACEMENT_CONFIGS = [
     {

tests/test_zero/test_gemini/test_zerooptim_state_dict.py

@@ -5,10 +5,10 @@ import torch.distributed as dist
 import colossalai
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils import set_seed
 from colossalai.zero import GeminiDDP, GeminiOptimizer
 from colossalai.zero.gemini.chunk import search_chunk_configuration
 from tests.components_to_test.registry import non_distributed_component_funcs
-from tests.test_tensor.common_utils import set_seed
 
 PLACEMENT_CONFIGS = [
     {

tests/test_zero/test_low_level/test_zero_tp.py

-import pytest
-import torch
-import torch.nn as nn
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.testing import assert_close
-
-import colossalai
-from colossalai.tensor import ProcessGroup
-from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
-from colossalai.utils import get_current_device
-from colossalai.zero import ColoInitContext, LowLevelZeroOptimizer
-from tests.test_tensor.common_utils import set_seed, split_param_col_tp1d, split_param_row_tp1d, tensor_shard_equal
-
-
-def strict_shard_equal(tensor, shard, tp_pg, rtol=1e-3, atol=1e-4):
-    return tensor_shard_equal(tensor, shard, tp_pg.tp_local_rank(), tp_pg.tp_world_size(), rtol, atol)
-
-
-class MlpModel(nn.Module):
-
-    def __init__(self):
-        super(MlpModel, self).__init__()
-        self.linear1 = nn.Linear(32, 128)
-        self.act = nn.GELU()
-        self.linear2 = nn.Linear(128, 32)
-
-    def forward(self, x):
-        y = self.linear1(x)
-        y = self.act(y)
-        y = self.linear2(y)
-        return x + y
-
-
-@parameterize("overlap_flag", [False, True])
-@parameterize("partition_flag", [False, True])
-def exam_zero_with_tp(overlap_flag, partition_flag):
-    set_seed(233010)
-    tp_pg = ProcessGroup(tp_degree=2)
-
-    with ColoInitContext(device=get_current_device(), default_pg=tp_pg):
-        hybrid_model = MlpModel()
-    torch_model = MlpModel().cuda()
-    for pt, ph in zip(torch_model.parameters(), hybrid_model.parameters()):
-        pt.data.copy_(ph.data)
-
-    for name, param in hybrid_model.named_parameters():
-        if 'linear1' in name:
-            split_param_row_tp1d(param, tp_pg)
-            param.compute_spec.set_output_replicate(False)
-        if 'linear2.weight' in name:
-            split_param_col_tp1d(param, tp_pg)
-
-    torch_model = DDP(torch_model, device_ids=[tp_pg.rank()], process_group=tp_pg.dp_process_group())
-    torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-2)    # set to 1e-2 for torch-1.11
-    hybrid_optim = torch.optim.Adam(hybrid_model.parameters(), lr=1e-2)
-    hybrid_optim = LowLevelZeroOptimizer(hybrid_optim,
-                                         initial_scale=2,
-                                         clip_grad_norm=1.0,
-                                         overlap_communication=overlap_flag,
-                                         partition_grad=partition_flag,
-                                         dp_process_group=tp_pg.dp_process_group(),
-                                         tp_process_group=tp_pg.tp_process_group())
-
-    dp_local_rank = tp_pg.dp_local_rank()
-    set_seed(255 + dp_local_rank)
-
-    data = torch.randn(8, 32, device=get_current_device())
-    torch_loss = torch_model(data).sum()
-    hybrid_loss = hybrid_model(data).sum()
-    assert_close(torch_loss, hybrid_loss)
-
-    torch_loss.backward()
-    torch.nn.utils.clip_grad_norm_(torch_model.parameters(), 1.0)
-    hybrid_optim.backward(hybrid_loss)
-
-    torch_optim.step()
-    hybrid_optim.step()
-
-    for (name, pt), ph in zip(torch_model.named_parameters(), hybrid_model.parameters()):
-        assert strict_shard_equal(pt.data, ph.data, tp_pg)
-
-
-def run_dist(rank, world_size, port):
-    colossalai.launch(config={}, rank=rank, world_size=world_size, port=port, host='localhost')
-    exam_zero_with_tp()
-
-
-@pytest.mark.skip('this will be rewritten by shardformer')
-@pytest.mark.dist
-@rerun_if_address_is_in_use()
-def test_zero_with_tp():
-    spawn(run_dist, 4)
-
-
-if __name__ == '__main__':
-    test_zero_with_tp()