[refactor] refactor ColoTensor's unit tests (#1340)

tests/test_tensor/model/test_gpt2.py

@@ -0,0 +1,119 @@
+import pytest
+
+from functools import partial
+from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, set_seed
+
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.multiprocessing as mp
+
+import colossalai
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils.cuda import get_current_device
+from colossalai.utils import free_port
+from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.tensor import ShardSpec, ComputePattern, ComputeSpec, ProcessGroup, ColoTensor, ColoTensorSpec
+from colossalai.nn.parallel.data_parallel import ColoDDP
+from tests.components_to_test.registry import non_distributed_component_funcs
+
+
+def init_1d_row_spec(model, pg: ProcessGroup):
+    tensor_spec = (ShardSpec([0], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
+    for n, p in model.named_parameters():
+        p.set_process_group(pg)
+        if 'weight' in n and 'ln' not in n:
+            p.set_tensor_spec(*tensor_spec)
+
+
+def init_1d_col_spec(model, pg: ProcessGroup):
+    spec = (ShardSpec([-1], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
+
+    for n, p in model.named_parameters():
+        p.set_process_group(pg)
+        if 'ln' not in n and ('weight' in n or 'bias' in n):
+            p.set_tensor_spec(*spec)
+
+
+def check_param_equal(model, torch_model, pg: ProcessGroup):
+    for p, torch_p in zip(model.parameters(), torch_model.parameters()):
+        assert pg.tp_local_rank() is not None, f"{pg.rank()} {pg.tp_world_size()} {pg._tp_degree} {pg.tp_local_rank()}1"
+        assert pg.tp_world_size() is not None
+        assert tensor_shard_equal(torch_p, p, pg.tp_local_rank(), pg.tp_world_size())
+
+
+def check_grad_equal(model, torch_model, pg: ProcessGroup):
+    for p, torch_p in zip(model.parameters(), torch_model.parameters()):
+        assert tensor_shard_equal(torch_p.grad, p.grad, pg.tp_local_rank(), pg.tp_world_size())
+
+
+def run_gpt(init_spec_func, use_ddp):
+    world_size = torch.distributed.get_world_size()
+
+    # build a PG with TP and DP hybrid
+    pg = ProcessGroup(dp_degree=(2 if (use_ddp and world_size >= 2) else 1))
+
+    # set seed make processes of the same tp group use the same seed
+    # set_seed(pg.tp_local_rank())
+
+    get_components_func = non_distributed_component_funcs.get_callable('gpt2')
+    model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
+
+    # make sure torch_model and model has the same parameter values
+    with ColoInitContext(device=get_current_device()):
+        model = model_builder()
+    model = model.cuda()
+    torch_model = model_builder().cuda()
+
+    if use_ddp:
+        torch_model = DDP(torch_model, device_ids=[pg.rank()], process_group=pg.dp_process_group())
+        model = ColoDDP(model, process_group=pg)
+
+    for torch_p, p in zip(torch_model.parameters(), model.parameters()):
+        torch_p.data.copy_(p)
+
+    init_spec_func(model, pg)
+
+    check_param_equal(model, torch_model, pg)
+
+    # close the dropout in eval mode
+    model.eval()
+    torch_model.eval()
+    set_seed(pg.dp_local_rank())
+    torch.distributed.barrier()
+    for i, (input_ids, attn_mask) in enumerate(train_dataloader):
+        colo_input = ColoTensor.from_torch_tensor(input_ids, ColoTensorSpec(pg))
+        logits = model(colo_input, attn_mask)
+        torch_logits = torch_model(input_ids, attn_mask)
+        assert tensor_equal(torch_logits, logits), f"{torch_logits - logits}"
+        loss = criterion(logits, input_ids)
+        torch_loss = criterion(torch_logits, input_ids)
+        if use_ddp:
+            model.backward(loss)
+        else:
+            loss.backward()
+        torch_loss.backward()
+        check_grad_equal(model, torch_model, pg)
+        if i > 0:
+            break
+    set_seed(313)
+
+
+def run_dist(rank, world_size, port, use_ddp):
+    if use_ddp and world_size == 1:
+        return
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    run_gpt(init_1d_row_spec, use_ddp)
+    run_gpt(init_1d_col_spec, use_ddp)
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 4])
+@pytest.mark.parametrize('use_ddp', [False, True])
+@rerun_if_address_is_in_use()
+def test_gpt(world_size, use_ddp):
+    run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp)
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_gpt(4, use_ddp=True)

tests/test_tensor/model/test_model.py

@@ -0,0 +1,335 @@
+import pytest
+from functools import partial
+import torch
+import torch.multiprocessing as mp
+
+from colossalai.tensor.colo_parameter import ColoParameter
+import colossalai
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils.cuda import get_current_device
+from colossalai.utils import free_port
+from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.tensor import ColoTensor, ProcessGroup
+from colossalai.nn.optimizer import ColossalaiOptimizer
+
+from tests.components_to_test.registry import non_distributed_component_funcs
+from tests.test_tensor.common_utils import tensor_shard_equal, check_equal, set_seed, \
+    split_param_row_tp1d, split_param_col_tp1d
+
+
+def run_1d_hybrid_tp(model_name):
+    # A simple net with two stacked nn.Linear
+    get_components_func = non_distributed_component_funcs.get_callable(model_name)
+    model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
+
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    set_seed(1)
+    with ColoInitContext(device=get_current_device()):
+        model = model_builder(checkpoint=True)
+
+    if rank == 0:
+        model_torch = model_builder(checkpoint=True)
+        model_torch = model_torch.cuda()
+
+        optimizer_torch = ColossalaiOptimizer(torch.optim.SGD(model_torch.parameters(), lr=0.1))
+
+        # Make two models have the same init params
+        for p1, p2 in zip(model.parameters(), model_torch.parameters()):
+            p2.data.copy_(p1.data)
+    else:
+        model_torch = None
+        optimizer_torch = None
+
+    pg = ProcessGroup(tp_degree=world_size)
+    if 'bert' == model_name:
+        for name, p in model.named_parameters():
+            if not isinstance(p, ColoTensor):
+                continue
+
+            # num_class = type_vocab_size = 2 | (8, 2)
+            if 'classifier' in name and 'weight' in name:
+                split_param_col_tp1d(p, pg)
+            # num_class = vocab_size = 30524 | (30524, 8)
+            elif 'word_embeddings' in name and 'weight' in name:
+                split_param_row_tp1d(p, pg)
+            # num_class = seq_len = 512 | (512, 8)
+            elif 'position_embeddings' in name and 'weight' in name:
+                split_param_row_tp1d(p, pg)
+            # num_class = type_vocab_size = 2 | (2, 8)
+            elif 'token_type_embeddings' in name and 'weight' in name:
+                split_param_col_tp1d(p, pg)
+
+    elif "simple_net" == model_name:
+        # A naive way to set spec for all weights in Linear
+        for name, p in model.named_parameters():
+            if not isinstance(p, ColoTensor):
+                continue
+            if 'embed' in name and 'weight' in name:
+                split_param_col_tp1d(p, pg)
+            if 'proj1' in name and ('weight' in name or 'bias' in name):
+                split_param_row_tp1d(p, pg)
+            if 'proj2' in name and 'weight' in name:
+                split_param_col_tp1d(p, pg)
+            if 'classifier' in name and ('weight' in name or 'bias' in name):
+                split_param_row_tp1d(p, pg)
+
+    model = model.cuda()
+    model.eval()
+    if rank == 0:
+        model_torch.eval()
+
+    colo_optimizer = ColossalaiOptimizer(torch.optim.SGD(model.parameters(), lr=0.1))
+
+    for i, (data, label) in enumerate(train_dataloader):
+
+        # Zero grad
+        colo_optimizer.zero_grad()
+        if rank == 0:
+            optimizer_torch.zero_grad()
+        torch.distributed.barrier()
+
+        data = data.to(get_current_device())
+        label = label.to(get_current_device())
+
+        torch.distributed.broadcast(data, 0, group=pg.tp_process_group())
+        torch.distributed.broadcast(label, 0, group=pg.tp_process_group())
+
+        # Bcast rank0 data to all processes
+        if criterion:
+            output = model(data)
+            loss = criterion(output, label)
+        else:
+            output = model(data, label)
+            loss = output
+
+        # Test output
+        if rank == 0:
+            if criterion:
+                output_torch = model_torch(data)
+                loss_torch = criterion(output_torch, label)
+            else:
+                output_torch = model_torch(data, label)
+                loss_torch = output_torch
+            assert torch.allclose(loss, loss_torch, rtol=1e-2)
+        torch.distributed.barrier()
+
+        loss.backward()
+        colo_optimizer.step()
+
+        if rank == 0:
+            loss_torch.backward()
+            optimizer_torch.step()
+
+            with torch.no_grad():
+                # check param
+                for p, torch_p in zip(model.parameters(), model_torch.parameters()):
+                    assert tensor_shard_equal(torch_p, p, pg.tp_local_rank(), pg.tp_world_size())
+        torch.distributed.barrier()
+        if i > 5:
+            break
+
+
+# Test the overrided parameters() and named_parameters() member functions
+def test_model_parameters():
+    colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl')
+
+    # build a module with 2 Linear, 4 parameters in total.
+    class Net(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.fcs = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.Linear(3, 2))
+            self.extra_param = torch.nn.Parameter(torch.randn(2))
+
+    with ColoInitContext(device=get_current_device()):
+        model = Net()
+
+    param_cnt = 0
+    for name, p in model.named_parameters():
+        param_cnt += 1
+    assert param_cnt == 5
+
+    for name, colo_p in model.named_parameters():
+        assert colo_p.is_model_data()
+
+    param_cnt = 0
+    for name, p in model.named_parameters(recurse=False):
+        param_cnt += 1
+    assert param_cnt == 1
+
+    param_cnt = 0
+    for p in model.fcs[0].parameters(recurse=False):
+        param_cnt += 1
+    assert param_cnt == 2
+
+
+def test_colo_optimizer():
+    get_components_func = non_distributed_component_funcs.get_callable('simple_net')
+    model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
+    set_seed(1)
+    with ColoInitContext(lazy_memory_allocate=False, device=get_current_device()):
+        model = model_builder(checkpoint=True)
+
+    colo_optimizer = ColossalaiOptimizer(torch.optim.SGD(model.parameters(), lr=0.1))
+    for i, (data, label) in enumerate(train_dataloader):
+        colo_optimizer.zero_grad()
+        data = data.to(get_current_device())
+        label = label.to(get_current_device())
+
+        # Bcast rank0 data to all processes
+        if criterion:
+            output = model(data)
+            loss = criterion(output, label)
+        else:
+            output = model(data, label)
+            loss = output
+
+        loss.backward()
+        colo_optimizer.step()
+
+        if i > 5:
+            break
+
+
+def run_1d_row_tp(model_name: str):
+    # A simple net with two stacked nn.Linear
+    get_components_func = non_distributed_component_funcs.get_callable(model_name)
+    model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
+    rank = torch.distributed.get_rank()
+
+    set_seed(1)
+    with ColoInitContext(device=get_current_device()):
+        model = model_builder(checkpoint=True)
+
+    world_size = torch.distributed.get_world_size()
+    pg = ProcessGroup(tp_degree=world_size)
+
+    set_seed(1)
+    if rank == 0:
+        model_torch = model_builder(checkpoint=True)
+        model_torch = model_torch.cuda()
+
+    # A naive way to set spec for all weights in Linear
+    for mo_name, module in model.named_modules():
+        # print(mo_name)
+        for pa_name, param in module.named_parameters(recurse=False):
+            # print('\t', pa_name, param.shape)
+            if not isinstance(param, ColoTensor):
+                continue
+            if 'weight' in pa_name:
+                if 'embed' in mo_name and 'token' not in mo_name and 'LayerNorm' not in mo_name:
+                    split_param_row_tp1d(param, pg)
+                elif 'LayerNorm' not in mo_name and 'ln' not in mo_name:
+                    split_param_col_tp1d(param, pg)
+
+    model = model.cuda()
+
+    for i, (data, label) in enumerate(train_dataloader):
+        data = data.to(get_current_device())
+        label = label.to(get_current_device())
+
+        torch.distributed.broadcast(data, 0, group=pg.tp_process_group())
+        torch.distributed.broadcast(label, 0, group=pg.tp_process_group())
+
+        # Bcast rank0 data to all processes
+        if criterion:
+            output = model(data)
+            loss = criterion(output, label)
+        else:
+            output = model(data, label)
+            loss = output
+
+        # For reference
+        if rank == 0:
+            if criterion:
+                output_torch = model_torch(data)
+                loss_torch = criterion(output_torch, label)
+            else:
+                output_torch = model_torch(data, label)
+                loss_torch = output_torch
+            assert torch.allclose(loss, loss_torch, rtol=1e-2)
+        torch.distributed.barrier()
+
+        loss.backward()
+
+        if rank == 0:
+            loss_torch.backward()
+        torch.distributed.barrier()
+
+        if i > 5:
+            break
+
+
+def _run_pretrain_load():
+    from transformers import BertForMaskedLM
+    set_seed(1)
+    model_pretrained = BertForMaskedLM.from_pretrained('bert-base-uncased')
+    with ColoInitContext(lazy_memory_allocate=False, device=get_current_device()):
+        model = BertForMaskedLM.from_pretrained('bert-base-uncased')
+
+    model_pretrained = model_pretrained.cuda()
+    model = model.cuda()
+
+    dict_pretrained = {}
+    dict_col = {}
+    c_ref = 0
+    for name, param in model_pretrained.named_parameters():
+        dict_pretrained[name] = param
+        c_ref += 1
+    c1 = 0
+    c2 = 0
+    for name, param in model.named_parameters():
+        if isinstance(param, ColoParameter):
+            c1 += 1
+        else:
+            c2 += 1
+        dict_col[name] = param
+    assert c_ref == c1
+    assert c2 == 0
+    if model_pretrained.cls.predictions.decoder.bias is model_pretrained.cls.predictions.bias:
+        assert model.cls.predictions.decoder.bias is model.cls.predictions.bias
+
+    for name, param in dict_pretrained.items():
+        check_equal(param, dict_col[name])
+
+
+def run_model_dist(rank, world_size, port):
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    # Comment below test for speed consideration
+    # for name in ['bert', 'simple_net']:
+    #     run_1d_row_tp(name)
+    for name in ['bert', 'simple_net']:
+        run_1d_hybrid_tp(name)
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 4])
+@rerun_if_address_is_in_use()
+def test_model(world_size):
+    run_func = partial(run_model_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+def run_pretrain_load_dist(rank, world_size, port):
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    _run_pretrain_load()
+
+
+# The test case has to download huggingface pretrained models from the internet
+# So we manually trigger the test.
+@pytest.mark.skip
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 4])
+@rerun_if_address_is_in_use()
+def test_pretrain_load(world_size):
+    run_func = partial(run_pretrain_load_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    # test_model_parameters()
+    # test_colo_optgimizer()
+    test_model(4)
+    # test_pretrain_load(4)

tests/test_tensor/model/test_module_spec.py

@@ -0,0 +1,229 @@
+from copy import deepcopy
+import pytest
+from functools import partial
+
+import torch
+import torch.multiprocessing as mp
+
+from colossalai.tensor import ColoTensor, ComputePattern, ComputeSpec, ShardSpec, ColoTensorSpec
+from colossalai.nn.parallel.layers import init_colo_module, check_colo_module
+from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, set_seed
+
+import colossalai
+from colossalai.utils.cuda import get_current_device
+from colossalai.utils.model.colo_init_context import ColoInitContext
+
+from colossalai.tensor import distspec, ProcessGroup, ReplicaSpec
+
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils import free_port
+
+from tests.components_to_test.registry import non_distributed_component_funcs
+
+
+def run_model_with_spec(mode, model_name):
+    get_components_func = non_distributed_component_funcs.get_callable(model_name)
+    model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
+    world_size = torch.distributed.get_world_size()
+    pg = ProcessGroup(tp_degree=world_size)
+    rank = pg.rank()
+
+    set_seed(1)
+    with ColoInitContext(device=get_current_device()):
+        model = model_builder(checkpoint=False)
+
+    if rank == 0:
+        model_seq = model_builder(checkpoint=False)
+        model_seq = model_seq.cuda()
+
+        # Make two models have the same init params
+        for p1, p2 in zip(model.parameters(), model_seq.parameters()):
+            p2.data.copy_(p1.data)
+
+    compute_spec = ComputeSpec(ComputePattern.TP1D)
+    # Not all layers in Bert can be mod by 4.
+    # e.g. row shard for all layers is invalid because the first dim of some layer is the classification type size 2.
+    if 'bert' == model_name:
+        if 'col' == mode:
+            init_colo_module(model.bert.embeddings, compute_spec, pg=pg, recursive=True, mode=mode)
+            init_colo_module(model.bert.encoder, compute_spec, pg=pg, recursive=True, mode=mode)
+            init_colo_module(model.classifier, compute_spec, pg=pg, recursive=True, mode='row')
+        elif 'row' == mode:
+            init_colo_module(model.bert.embeddings, compute_spec, pg=pg, recursive=True, mode='col')
+            init_colo_module(model.bert.encoder, compute_spec, pg=pg, recursive=True, mode=mode)
+            init_colo_module(model.classifier, compute_spec, pg=pg, recursive=True, mode=mode)
+    elif 'simple_net' == model_name:
+        init_colo_module(model, compute_spec, pg=pg, recursive=True, mode=mode)
+
+    model = model.cuda()
+    for i, (data, label) in enumerate(train_dataloader):
+        data = data.to(get_current_device())
+        label = label.to(get_current_device())
+
+        torch.distributed.broadcast(data, 0, group=pg.tp_process_group())
+        torch.distributed.broadcast(label, 0, group=pg.tp_process_group())
+
+        if criterion:
+            output = model(data)
+            loss = criterion(output, label)
+        else:
+            output = model(data, label)
+            loss = output
+
+        # For reference
+        if rank == 0:
+            if criterion:
+                output_seq = model_seq(data)
+                loss_seq = criterion(output_seq, label)
+            else:
+                output_seq = model_seq(data, label)
+                loss_seq = output_seq
+
+        if rank == 0:
+            with torch.no_grad():
+                assert torch.allclose(loss, loss_seq, rtol=1e-2)
+
+        loss.backward()
+
+        if rank == 0:
+            loss_seq.backward()
+
+            with torch.no_grad():
+                # check param
+                for p1, p2 in zip(model.parameters(), model_seq.parameters()):
+                    if p1.size() == p2.size():
+                        assert torch.allclose(p1, p2)
+                    else:
+                        if p1.size(-1) < p2.size(-1):    # col
+                            world_size = p2.size(-1) // p1.size(-1)
+                            split_p2 = torch.chunk(p2, world_size, dim=-1)[0]
+
+                        elif p1.size(0) < p2.size(0):    # row
+                            world_size = p2.size(0) // p1.size(0)
+                            split_p2 = torch.chunk(p2, world_size, dim=0)[0]
+
+                        assert torch.allclose(p1, split_p2)
+
+        if i > 3:
+            break
+
+
+def run_linear_with_spec(mode):
+    with ColoInitContext(device=get_current_device()):
+        model = torch.nn.Linear(4, 8)
+
+    model_handy = deepcopy(model)
+    world_size = torch.distributed.get_world_size()
+    pg = ProcessGroup(tp_degree=world_size)
+    compute_spec = ComputeSpec(ComputePattern.TP1D)
+    init_colo_module(model, compute_spec, pg=pg, recursive=True, mode=mode)
+
+    x = torch.rand(2, 4).cuda()
+    colo_x = ColoTensor.from_torch_tensor(x, ColoTensorSpec(pg))
+
+    out = model(x)
+    colo_out = model_handy(colo_x)
+    assert tensor_equal(out, colo_out)
+
+    grad = torch.rand_like(out)
+    out.backward(grad)
+    colo_out.backward(grad)
+
+    assert tensor_shard_equal(model_handy.weight.grad, model.weight.grad, pg.tp_local_rank(), pg.tp_world_size())
+    assert tensor_shard_equal(model_handy.bias.grad, model.bias.grad, pg.tp_local_rank(), pg.tp_world_size())
+
+
+def run_check_shared_param():
+    from transformers import BertForMaskedLM, BertConfig
+    hidden_dim = 8
+    num_head = 4
+    sequence_length = 12
+    num_layer = 2
+    vocab_size = 24
+
+    world_size = torch.distributed.get_world_size()
+    pg = ProcessGroup(tp_degree=world_size)
+    rank = pg.rank()
+
+    config = BertConfig(vocab_size=vocab_size,
+                        hidden_size=hidden_dim,
+                        intermediate_size=hidden_dim * 4,
+                        num_attention_heads=num_head,
+                        max_position_embeddings=sequence_length,
+                        num_hidden_layers=num_layer,
+                        hidden_dropout_prob=0.,
+                        attention_probs_dropout_prob=0.)
+    with ColoInitContext(lazy_memory_allocate=False, device=get_current_device()):
+        model = BertForMaskedLM(config)
+
+    model = model.cuda()
+    compute_spec = ComputeSpec(ComputePattern.TP1D)
+    # model.cls.predictions.decoder and model.cls.predictions share the bias, so they should have the same spec
+    assert len(model.cls.predictions.decoder.bias.shared_param_modules) == 2
+    # They are all Linear, so both row is allowed. This should pass check.
+    init_colo_module(model, compute_spec, pg=pg, recursive=True, mode='row')
+    # This should be detected by check because you can not set weight as row while set bias as col.
+    col_spec = (ShardSpec([0], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
+
+    # TODO(jiaruifang) optimize this line
+    if not model.cls.predictions.bias.has_initialized:
+        model.cls.predictions.bias.pg = pg
+        model.cls.predictions.bias.dist_spec = ReplicaSpec()
+        model.cls.predictions.bias.has_initialized = True
+    model.cls.predictions.bias.set_tensor_spec(*col_spec)
+    try:
+        check_colo_module(model.cls.predictions.decoder, pg=pg, recursive=False)
+    except Exception as e:
+        assert 'incorrectly sharded' in str(e)
+
+
+def run_dist(rank, world_size, port):
+    config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
+    colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    run_linear_with_spec('col')
+    run_linear_with_spec('row')
+
+
+def run_dist_model(rank, world_size, port):
+    config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
+    colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    for model_name in ['simple_net', 'bert']:
+        run_model_with_spec('col', model_name)
+        run_model_with_spec('row', model_name)
+
+
+def run_dist_check(rank, world_size, port):
+    config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
+    colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    run_check_shared_param()
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 4])
+@pytest.mark.skip("for higher testing speed")
+@rerun_if_address_is_in_use()
+def test_module_linear_1d(world_size):
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 4])
+@pytest.mark.skip("for higher testing speed")
+@rerun_if_address_is_in_use()
+def test_module_model(world_size):
+    run_func = partial(run_dist_model, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 2])
+@pytest.mark.skip("for higher testing speed")
+@rerun_if_address_is_in_use()
+def test_module_check(world_size):
+    run_func = partial(run_dist_check, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_module_linear_1d(4)