[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_data/test_data_parallel_sampler.py → tests/test_legacy/test_data/test_data_parallel_sampler.py

@@ -10,10 +10,11 @@ import torch.distributed as dist
 from torchvision import datasets, transforms
 
 import colossalai
-from colossalai.context import Config, ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.context import Config
+from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.utils import get_dataloader
 from colossalai.testing import rerun_if_address_is_in_use, spawn
-from colossalai.utils import get_dataloader
 
 CONFIG = Config(dict(
     parallel=dict(
@@ -26,7 +27,7 @@ CONFIG = Config(dict(
 
 def run_data_sampler(rank, world_size, port):
     dist_args = dict(config=CONFIG, rank=rank, world_size=world_size, backend='gloo', port=port, host='localhost')
-    colossalai.launch(**dist_args)
+    colossalai.legacy.launch(**dist_args)
     print('finished initialization')
 
     # build dataset

tests/test_legacy/test_data/test_deterministic_dataloader.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+import torch.distributed as dist
+from torchvision import datasets, transforms
+
+import colossalai
+from colossalai.context import Config
+from colossalai.legacy.context import ParallelMode
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.utils import get_dataloader
+from colossalai.testing import rerun_if_address_is_in_use, spawn
+
+CONFIG = Config(
+    dict(
+        train_data=dict(
+            dataset=dict(
+                type='CIFAR10',
+                root=Path(os.environ['DATA']),
+                train=True,
+                download=True,
+            ),
+            dataloader=dict(num_workers=2, batch_size=2, shuffle=True),
+        ),
+        parallel=dict(
+            pipeline=dict(size=1),
+            tensor=dict(size=1, mode=None),
+        ),
+        seed=1024,
+    ))
+
+
+def run_data_sampler(rank, world_size, port):
+    dist_args = dict(config=CONFIG, rank=rank, world_size=world_size, backend='gloo', port=port, host='localhost')
+    colossalai.legacy.launch(**dist_args)
+
+    # build dataset
+    transform_pipeline = [transforms.ToTensor(), transforms.RandomCrop(size=32, padding=4)]
+    transform_pipeline = transforms.Compose(transform_pipeline)
+    dataset = datasets.CIFAR10(root=Path(os.environ['DATA']), train=True, download=True, transform=transform_pipeline)
+
+    # build dataloader
+    dataloader = get_dataloader(dataset, batch_size=8, add_sampler=False)
+
+    data_iter = iter(dataloader)
+    img, label = data_iter.next()
+    img = img[0]
+
+    if gpc.get_local_rank(ParallelMode.DATA) != 0:
+        img_to_compare = img.clone()
+    else:
+        img_to_compare = img
+    dist.broadcast(img_to_compare, src=0, group=gpc.get_group(ParallelMode.DATA))
+
+    if gpc.get_local_rank(ParallelMode.DATA) != 0:
+        # this is without sampler
+        # this should be false if data parallel sampler to given to the dataloader
+        assert torch.equal(img,
+                           img_to_compare), 'Same image was distributed across ranks and expected it to be the same'
+    torch.cuda.empty_cache()
+
+
+@rerun_if_address_is_in_use()
+def test_data_sampler():
+    spawn(run_data_sampler, 4)
+
+
+if __name__ == '__main__':
+    test_data_sampler()

tests/test_legacy/test_engine/test_engine.py

 import pytest
 
 import colossalai
-from colossalai.amp import AMP_TYPE
-from colossalai.core import global_context as gpc
+from colossalai.legacy.amp import AMP_TYPE
+from colossalai.legacy.core import global_context as gpc
 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
 from tests.components_to_test.registry import non_distributed_component_funcs
 
@@ -20,8 +20,9 @@ def run_train(model_name, amp_mode):
     model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
 
     model = model_builder(checkpoint=False)
-    engine, train_dataloader, *args = colossalai.initialize(model=model,
-                                                            optimizer=optimizer_class(model.parameters(), lr=1e-3),
+    engine, train_dataloader, *args = colossalai.legacy.initialize(model=model,
+                                                                   optimizer=optimizer_class(model.parameters(),
+                                                                                             lr=1e-3),
                                                                    criterion=criterion,
                                                                    train_dataloader=train_dataloader)
 
@@ -48,7 +49,12 @@ def run_train(model_name, amp_mode):
 
 def run_engine(rank, world_size, port):
     # init dist env
-    colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.legacy.launch(config=CONFIG,
+                             rank=rank,
+                             world_size=world_size,
+                             host='localhost',
+                             port=port,
+                             backend='nccl')
     run_train()
 
 

tests/test_legacy/test_engine/test_gradient_accumluation.py

@@ -10,10 +10,10 @@ from torchvision.datasets import CIFAR10
 from torchvision.models import resnet18
 
 import colossalai
-from colossalai.core import global_context as gpc
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.utils import get_dataloader
 from colossalai.logging import get_dist_logger
 from colossalai.testing import rerun_if_address_is_in_use, spawn
-from colossalai.utils import get_dataloader
 
 # Config
 BATCH_SIZE = 2
@@ -27,7 +27,12 @@ CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=1, mode=None
 def run_no_pipeline(rank, world_size, port):
 
     # init dist env
-    colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.legacy.launch(config=CONFIG,
+                             rank=rank,
+                             world_size=world_size,
+                             host='localhost',
+                             port=port,
+                             backend='nccl')
 
     # build model
     model = resnet18(num_classes=10)
@@ -49,7 +54,7 @@ def run_no_pipeline(rank, world_size, port):
     optimizer = Adam(model.parameters(), lr=0.001)
     criterion = nn.CrossEntropyLoss()
 
-    engine, train_dataloader, *args = colossalai.initialize(model=model,
+    engine, train_dataloader, *args = colossalai.legacy.initialize(model=model,
                                                                    optimizer=optimizer,
                                                                    criterion=criterion,
                                                                    train_dataloader=train_dataloader)

tests/test_legacy/test_layers/test_1d/checks_1d/check_layer_1d.py

@@ -2,9 +2,9 @@ import torch
 import torch.distributed as dist
 from torch.nn import Parameter
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
-from colossalai.global_variables import tensor_parallel_env as env
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.global_variables import tensor_parallel_env as env
 from colossalai.legacy.nn import (
     Classifier1D,
     Embedding1D,
@@ -15,7 +15,8 @@ from colossalai.legacy.nn import (
     VocabParallelCrossEntropyLoss1D,
     VocabParallelEmbedding1D,
 )
-from colossalai.utils import get_current_device, print_rank_0
+from colossalai.legacy.utils import print_rank_0
+from colossalai.utils import get_current_device
 
 from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, NUM_CLASSES, SEQ_LENGTH, VOCAB_SIZE, check_equal
 

tests/test_legacy/test_layers/test_1d/test_1d.py

@@ -5,8 +5,8 @@ import pytest
 import torch
 from checks_1d.check_layer_1d import *
 
-from colossalai.core import global_context as gpc
-from colossalai.initialize import launch
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.testing import rerun_if_address_is_in_use, spawn
 

tests/test_legacy/test_layers/test_2d/checks_2d/check_layer_2d.py

 import torch
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.legacy.nn import (
     Classifier2D,
     CrossEntropyLoss2D,
@@ -15,7 +15,8 @@ from colossalai.legacy.nn import (
     VocabParallelCrossEntropyLoss2D,
     VocabParallelEmbedding2D,
 )
-from colossalai.utils import get_current_device, print_rank_0
+from colossalai.legacy.utils import print_rank_0
+from colossalai.utils import get_current_device
 
 from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, IMG_SIZE, NUM_CLASSES, SEQ_LENGTH, VOCAB_SIZE, check_equal
 

tests/test_legacy/test_layers/test_2d/checks_2d/check_operation_2d.py

@@ -3,10 +3,11 @@
 
 import torch
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.legacy.nn.layer.parallel_2d._operation import Matmul_AB_2D, Matmul_ABT_2D, Matmul_ATB_2D
-from colossalai.utils import get_current_device, print_rank_0
+from colossalai.legacy.utils import print_rank_0
+from colossalai.utils import get_current_device
 
 from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, SEQ_LENGTH, check_equal
 

tests/test_legacy/test_layers/test_2d/test_2d.py

@@ -18,8 +18,8 @@ from checks_2d.check_layer_2d import (
 )
 from checks_2d.check_operation_2d import check_AB, check_ABT, check_ATB
 
-from colossalai.core import global_context as gpc
-from colossalai.initialize import launch
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.testing import rerun_if_address_is_in_use, spawn
 

tests/test_legacy/test_layers/test_2p5d/checks_2p5d/check_layer_2p5d.py

 import torch
 from torch.nn import Parameter
 
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.core import global_context as gpc
+from colossalai.legacy.context.parallel_mode import ParallelMode
+from colossalai.legacy.core import global_context as gpc
 from colossalai.legacy.nn import (
     Classifier2p5D,
     CrossEntropyLoss2p5D,
@@ -16,7 +16,8 @@ from colossalai.legacy.nn import (
     VocabParallelCrossEntropyLoss2p5D,
     VocabParallelEmbedding2p5D,
 )
-from colossalai.utils import get_current_device, print_rank_0
+from colossalai.legacy.utils import print_rank_0
+from colossalai.utils import get_current_device
 
 from .common import *
 

tests/test_legacy/test_layers/test_2p5d/checks_2p5d/check_operation_2p5d.py

 import torch
 
-from colossalai.context import ParallelMode
-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.legacy.nn.layer.parallel_2p5d._operation import Matmul_AB_2p5D, Matmul_ABT_2p5D, Matmul_ATB_2p5D
-from colossalai.utils import get_current_device, print_rank_0
+from colossalai.legacy.utils import print_rank_0
+from colossalai.utils import get_current_device
 
 from .common import *
 

tests/test_legacy/test_layers/test_2p5d/test_2p5d.py

@@ -3,8 +3,8 @@ import torch
 from checks_2p5d.check_layer_2p5d import *
 from checks_2p5d.check_operation_2p5d import check_AB, check_ABT, check_ATB
 
-from colossalai.core import global_context as gpc
-from colossalai.initialize import launch
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.testing import rerun_if_address_is_in_use, spawn
 

tests/test_legacy/test_layers/test_3d/checks_3d/check_layer_3d.py

@@ -5,8 +5,8 @@ import time
 
 import torch
 
-from colossalai.constants import INPUT_GROUP_3D, OUTPUT_GROUP_3D, WEIGHT_GROUP_3D
-from colossalai.core import global_context
+from colossalai.legacy.constants import INPUT_GROUP_3D, OUTPUT_GROUP_3D, WEIGHT_GROUP_3D
+from colossalai.legacy.core import global_context
 from colossalai.legacy.nn import (
     Classifier3D,
     CrossEntropyLoss3D,
@@ -21,8 +21,9 @@ from colossalai.legacy.nn import (
     VocabParallelEmbedding3D,
 )
 from colossalai.legacy.nn.layer.parallel_3d._utils import get_parallel_mode_from_env
+from colossalai.legacy.utils import print_rank_0
 from colossalai.logging import get_dist_logger
-from colossalai.utils import get_current_device, print_rank_0
+from colossalai.utils import get_current_device
 
 from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, IMG_SIZE, NUM_CLASSES, SEQ_LENGTH, VOCAB_SIZE, check_equal
 

tests/test_legacy/test_layers/test_3d/test_3d.py

@@ -15,8 +15,8 @@ from checks_3d.check_layer_3d import (
     check_vocab_parallel_loss,
 )
 
-from colossalai.core import global_context as gpc
-from colossalai.initialize import launch
+from colossalai.legacy.core import global_context as gpc
+from colossalai.legacy.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn
 

tests/test_legacy/test_layers/test_cache_embedding.py

@@ -14,7 +14,8 @@ from colossalai.legacy.nn.parallel.layers import (
     ParallelCachedEmbeddingBagTablewise,
     TablewiseEmbeddingBagConfig,
 )
-from colossalai.tensor import ColoTensor, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec
+from colossalai.legacy.tensor import ComputePattern, ComputeSpec, ProcessGroup, ShardSpec
+from colossalai.tensor import ColoTensor
 from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
 
 NUM_EMBED, EMBED_DIM = 10, 8
@@ -359,7 +360,7 @@ def run_parallel_freq_aware_embed_columnwise(rank, world_size):
 
 
 def run_dist(rank, world_size, port):
-    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    colossalai.legacy.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     # run_parallel_freq_aware_embed_columnwise(rank, world_size)
     run_parallel_freq_aware_embed_tablewise(rank, world_size)
 

tests/test_legacy/test_layers/test_sequence/checks_seq/check_layer_seq.py

 import torch
 
-from colossalai.context import ParallelMode
-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.legacy.nn import TransformerSelfAttentionRing
 from colossalai.utils import get_current_device
 

tests/test_legacy/test_layers/test_sequence/test_sequence.py

@@ -3,8 +3,8 @@ import torch
 import torch.distributed as dist
 
 import colossalai
-from colossalai.context import ParallelMode
-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.legacy.nn.layer.parallel_sequence import RingAV, RingQK
 from colossalai.testing import rerun_if_address_is_in_use, spawn
 
@@ -120,7 +120,7 @@ def check_ring_av(rank, world_size):
 
 
 def run_test(rank, world_size, port):
-    colossalai.launch(rank=rank, world_size=world_size, config=CONFIG, host='localhost', port=port)
+    colossalai.legacy.launch(rank=rank, world_size=world_size, config=CONFIG, host='localhost', port=port)
 
     # check_ring_qk(rank, world_size)
     check_ring_av(rank, world_size)

tests/test_pipeline/rpc_test_utils.py → tests/test_legacy/test_pipeline/rpc_test_utils.py

@@ -10,9 +10,9 @@ from torch import nn
 from torch._C._distributed_rpc import _is_current_rpc_agent_set
 from torch.optim import SGD, Adam, Optimizer, RMSprop
 
-from colossalai import launch
+from colossalai.legacy import launch
+from colossalai.legacy.pipeline.pipeline_process_group import ppg
 from colossalai.logging import disable_existing_loggers
-from colossalai.pipeline.pipeline_process_group import ppg
 
 rpc_is_initialized = _is_current_rpc_agent_set
 

tests/test_pipeline/test_cuda_rpc_chimera.py → tests/test_legacy/test_pipeline/test_cuda_rpc_chimera.py

 import torch
-from torch import nn
 import torch.autograd as autograd
+from rpc_test_utils import RpcTestModel, parse_args, rpc_run
+from torch import nn
 
-from colossalai.pipeline.rpc import ChimeraPipelineEngine
+from colossalai.legacy.pipeline.rpc import ChimeraPipelineEngine
 from colossalai.testing import assert_close
-from rpc_test_utils import rpc_run, parse_args, RpcTestModel
 
 # global variable for model created
 feat_num = 100

tests/test_pipeline/test_cuda_rpc_optimizer.py → tests/test_legacy/test_pipeline/test_cuda_rpc_optimizer.py

 import torch
-from torch import nn
-from torch import autograd
-from torch.optim import SGD, Adam, RMSprop, Optimizer
+from rpc_test_utils import RpcTestModel, parse_args, rpc_run
+from torch import autograd, nn
+from torch.optim import SGD, Adam, Optimizer, RMSprop
 
-from colossalai.pipeline.rpc._pipeline_schedule import FillDrainPipelineEngine, OneFOneBPipelineEngine
+from colossalai.legacy.pipeline.rpc._pipeline_schedule import FillDrainPipelineEngine, OneFOneBPipelineEngine
 from colossalai.testing import assert_close
-from rpc_test_utils import rpc_run, parse_args, RpcTestModel
 
 # global variable for model created
 feat_num = 100