[refactory] add nn.parallel module (#1068)

tests/test_tensor/test_zero_optim.py

@@ -6,11 +6,11 @@ from colossalai.context.parallel_mode import ParallelMode
 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 import ColoInitContext
+from colossalai.utils.model.colo_init_context import ColoInitContext
 from colossalai.tensor import ChunkManager
 from colossalai.core import global_context as gpc
 from functools import partial
-from _utils import tensor_equal, tensor_shard_equal, set_seed
+from _utils import tensor_equal, set_seed
 from tests.components_to_test.registry import non_distributed_component_funcs
 from torch.nn.parallel import DistributedDataParallel as DDP
 from colossalai.nn.parallel import ColoDDPV2

tests/test_utils/test_pipelinable.py

-import os.path as osp
-
-import pytest
 import torch
 import torch.multiprocessing as mp
 
 from colossalai.utils.model.pipelinable import PipelinableContext
 
-from functools import partial
-from colossalai.utils import free_port
 from colossalai.testing import rerun_on_exception
 
 NUM_CHUNKS = 1