github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-05-05 12:24:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Hotfix/Colossalai layers (#92)

Browse Source 
* 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>
This commit is contained in:
アマデウス
2021-12-29 23:32:10 +08:00
committed by GitHub
parent 0fedef4f3c
commit 01a80cd86d
71 changed files with 1033 additions and 773 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/test_zero_tensor_parallel/test_vit_2d_level_2.py
View File

@@ -13,7 +13,7 @@ import torch.multiprocessing as mp
from colossalai.core import global_context as gpc
from colossalai.logging import get_dist_logger
from colossalai.nn import CrossEntropyLoss
from colossalai.utils import get_dataloader
from colossalai.utils import free_port, get_dataloader
from model_zoo.vit import vit_lite_depth7_patch4_32
from torchvision import transforms
from torchvision.datasets import CIFAR10
@@ -40,11 +40,11 @@ def train_epoch(engine, train_dataloader):
return avg_loss
def run_2d_parallel_vision_transformer_level_2(rank, world_size):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=29950, backend='nccl')
def run_2d_parallel_vision_transformer_level_2(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
# build model
model = vit_lite_depth7_patch4_32(tensor_parallel='2d')
model = vit_lite_depth7_patch4_32()
# build dataloader# build dataloaders
train_dataset = CIFAR10(root=Path(os.environ['DATA']),
@@ -62,7 +62,7 @@ def run_2d_parallel_vision_transformer_level_2(rank, world_size):
# build optimizer and loss
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = CrossEntropyLoss(tensor_parallel='2d')
criterion = CrossEntropyLoss()
engine, train_dataloader, *args = colossalai.initialize(model=model,
optimizer=optimizer,
@@ -90,7 +90,7 @@ def run_2d_parallel_vision_transformer_level_2(rank, world_size):
@pytest.mark.dist
def test_2d_vit_zero_level_2():
world_size = 8
run_func = partial(run_2d_parallel_vision_transformer_level_2, world_size=world_size)
run_func = partial(run_2d_parallel_vision_transformer_level_2, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)

12
tests/test_zero_tensor_parallel/test_vit_2d_level_3.py
View File

@@ -13,7 +13,7 @@ import torch.multiprocessing as mp
from colossalai.core import global_context as gpc
from colossalai.logging import get_dist_logger
from colossalai.nn import CrossEntropyLoss
from colossalai.utils import get_dataloader
from colossalai.utils import free_port, get_dataloader
from model_zoo.vit import vit_lite_depth7_patch4_32
from torchvision import transforms
from torchvision.datasets import CIFAR10
@@ -40,11 +40,11 @@ def train_epoch(engine, train_dataloader):
return avg_loss
def run_2d_parallel_vision_transformer_level_3(rank, world_size):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=29951, backend='nccl')
def run_2d_parallel_vision_transformer_level_3(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
# build model
model = vit_lite_depth7_patch4_32(tensor_parallel='2d')
model = vit_lite_depth7_patch4_32()
# build dataloader# build dataloaders
train_dataset = CIFAR10(root=Path(os.environ['DATA']),
@@ -62,7 +62,7 @@ def run_2d_parallel_vision_transformer_level_3(rank, world_size):
# build optimizer and loss
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = CrossEntropyLoss(tensor_parallel='2d')
criterion = CrossEntropyLoss()
engine, train_dataloader, *args = colossalai.initialize(model=model,
optimizer=optimizer,
@@ -91,7 +91,7 @@ def run_2d_parallel_vision_transformer_level_3(rank, world_size):
@pytest.mark.skip("Level 3 has unknown bug so skip this test for now")
def test_3d_vit_zero_level_3():
world_size = 8
run_func = partial(run_2d_parallel_vision_transformer_level_3, world_size=world_size)
run_func = partial(run_2d_parallel_vision_transformer_level_3, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)