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_context/test_2d_init.py

@@ -1,15 +1,16 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
+from functools import partial
+from pathlib import Path
+
 import pytest
 import torch
 import torch.multiprocessing as mp
-
 from colossalai import launch
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
-from functools import partial
-from pathlib import Path
+from colossalai.utils import free_port
 
 CONFIG_PATH = Path(__file__).parent.joinpath('configs/parallel_2d_init.py').absolute()
 
@@ -87,7 +88,7 @@ def test_2d_init():
     test_fn = partial(init_2d,
                       world_size=world_size,
                       backend='gloo',
-                      port='29900',
+                      port=free_port(),
                       host='localhost'
                       )
     mp.spawn(test_fn, nprocs=world_size)

tests/test_context/test_2p5d_init.py

@@ -7,10 +7,10 @@ from pathlib import Path
 import pytest
 import torch
 import torch.multiprocessing as mp
-
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
 from colossalai.initialize import launch
+from colossalai.utils import free_port
 
 CONFIG_PATH = Path(__file__).parent.joinpath('configs/parallel_2p5d_init.py').absolute()
 
@@ -111,7 +111,7 @@ def test_2halfd_init():
     test_fn = partial(init_2halfd,
                       world_size=world_size,
                       backend='gloo',
-                      port='29901',
+                      port=free_port(),
                       host='localhost'
                       )
     mp.spawn(test_fn, nprocs=world_size)

tests/test_context/test_3d_init.py

@@ -7,11 +7,10 @@ from pathlib import Path
 import pytest
 import torch
 import torch.multiprocessing as mp
-
-
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.core import global_context as gpc
 from colossalai.initialize import launch
+from colossalai.utils import free_port
 
 CONFIG_PATH = Path(__file__).parent.joinpath('configs/parallel_3d_init.py').absolute()
 
@@ -104,7 +103,7 @@ def test_3d_init():
     test_fn = partial(init_3d,
                       world_size=world_size,
                       backend='gloo',
-                      port='29902',
+                      port=free_port(),
                       host='localhost'
                       )
     mp.spawn(test_fn, nprocs=world_size)