Hotfix/Colossalai layers (#92)

* optimized 1d layer apis; reorganized nn.layer modules; fixed tests

* fixed 2.5d runtime issue

* reworked split batch, now called in trainer.schedule.load_batch

Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>

tests/test_layers/test_1d/test_1d.py

@@ -1,13 +1,15 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
+from functools import partial
+
 import pytest
 import torch
 import torch.multiprocessing as mp
-
 from colossalai.core import global_context as gpc
 from colossalai.initialize import launch
-from functools import partial
+from colossalai.utils import free_port
+
 from checks_1d.check_layer_1d import *
 
 CONFIG = dict(
@@ -21,12 +23,12 @@ CONFIG = dict(
 )
 
 
-def check_layer(rank, world_size):
+def check_layer(rank, world_size, port):
     launch(config=CONFIG,
            rank=rank,
            world_size=world_size,
            host='localhost',
-           port=29920,
+           port=port,
            backend='nccl')
 
     check_linear_col()
@@ -39,7 +41,7 @@ def check_layer(rank, world_size):
 @pytest.mark.dist
 def test_1d():
     world_size = 4
-    run_func = partial(check_layer, world_size=world_size)
+    run_func = partial(check_layer, world_size=world_size, port=free_port())
     mp.spawn(run_func, nprocs=world_size)