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_utils/test_gradient_accumluation.py

@@ -1,21 +1,19 @@
-import colossalai
 import os
+from functools import partial
+from pathlib import Path
+
+import colossalai
 import pytest
 import torch
 import torch.multiprocessing as mp
 import torch.nn as nn
-
-from functools import partial
-from pathlib import Path
-from torchvision import transforms
-from torch.optim import Adam
 from colossalai.core import global_context as gpc
 from colossalai.logging import get_dist_logger
-from colossalai.utils import report_memory_usage, get_dataloader
-from colossalai.initialize import get_default_parser
-from torchvision.models import resnet18
+from colossalai.utils import free_port, get_dataloader
+from torch.optim import Adam
+from torchvision import transforms
 from torchvision.datasets import CIFAR10
-
+from torchvision.models import resnet18
 
 # Config
 BATCH_SIZE = 16
@@ -32,7 +30,7 @@ CONFIG = dict(
 )
 
 
-def run_no_pipeline(rank, world_size):
+def run_no_pipeline(rank, world_size, port):
 
     # init dist env
     colossalai.launch(
@@ -40,7 +38,7 @@ def run_no_pipeline(rank, world_size):
         rank=rank,
         world_size=world_size,
         host='localhost',
-        port=29500,
+        port=port,
         backend='nccl'
     )
 
@@ -110,7 +108,7 @@ def run_no_pipeline(rank, world_size):
 @pytest.mark.dist
 def test_engine():
     world_size = 4
-    func = partial(run_no_pipeline, world_size=world_size)
+    func = partial(run_no_pipeline, world_size=world_size, port=free_port())
     mp.spawn(func, nprocs=world_size)