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_2d/test_2d.py

@@ -1,16 +1,17 @@
 #!/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 colossalai.utils import free_port
+
 from checks_2d.check_layer_2d import *
 from checks_2d.check_operation_2d import *
-from functools import partial
-
 
 CONFIG = dict(
     parallel=dict(
@@ -34,12 +35,12 @@ def check_layer():
     check_layernorm()
     check_classifier()
 
-def check_layer_and_operation(rank, world_size):
+def check_layer_and_operation(rank, world_size, port):
     launch(config=CONFIG,
            rank=rank,
            world_size=world_size,
            host='localhost',
-           port=29921,
+           port=port,
            backend='nccl')
 
     # check_operations()
@@ -51,7 +52,7 @@ def check_layer_and_operation(rank, world_size):
 @pytest.mark.dist
 def test_2d():
     world_size = 4
-    run_func = partial(check_layer_and_operation, world_size=world_size)
+    run_func = partial(check_layer_and_operation, world_size=world_size, port=free_port())
     mp.spawn(run_func, nprocs=world_size)