github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[format] polish name format for MOE (#481)

Browse Source
This commit is contained in:
Jiarui Fang
2022-03-21 23:19:47 +08:00
committed by GitHub
parent 8d3250d74b
commit 65c0f380c2
6 changed files with 45 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_moe/test_grad_handler.py
View File

@@ -36,7 +36,7 @@ def run_test(rank, world_size, port):
model = model.to(get_current_device())
sync_moe_model_param(model)
dist_dict = MOE_CONTEXT.information
dist_dict = MOE_CONTEXT.parallel_info_dict
assert_equal_in_group(layer_list[0].experts.experts[0].weight.data, dist_dict[1].dp_group)
assert_equal_in_group(layer_list[1].experts.experts[0].weight.data, dist_dict[2].dp_group)
# MoE model synchronization passed

30
tests/test_moe/test_moe_group.py
View File

@@ -1,6 +1,5 @@
from functools import partial
import pytest
import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
@@ -16,7 +15,8 @@ D_FF = 8
CONFIG = dict()
def run_test(rank, world_size, port):
def run_test(rank, port):
world_size = 4
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
expert_module = nn.Linear
expert_factor = dict(in_features=D_MODEL, out_features=D_FF, device=get_current_device())
@@ -33,36 +33,36 @@ def run_test(rank, world_size, port):
assert exp3.num_local_experts == 2
# experts deployment passed
dist_dict = MOE_CONTEXT.information
parallel_info_dict = MOE_CONTEXT.parallel_info_dict
rank = dist.get_rank()
assert len(dist_dict) == 3
assert dist.get_rank(dist_dict[4].ep_group) == rank
assert dist.get_rank(dist_dict[2].ep_group) == rank % 2
assert dist.get_rank(dist_dict[1].ep_group) == 0
assert len(parallel_info_dict) == 3
assert dist.get_rank(parallel_info_dict[4].ep_group) == rank
assert dist.get_rank(parallel_info_dict[2].ep_group) == rank % 2
assert dist.get_rank(parallel_info_dict[1].ep_group) == 0
assert dist.get_rank(dist_dict[4].dp_group) == 0
assert dist.get_rank(dist_dict[2].dp_group) == rank // 2
assert dist.get_rank(dist_dict[1].dp_group) == rank
assert dist.get_rank(parallel_info_dict[4].dp_group) == 0
assert dist.get_rank(parallel_info_dict[2].dp_group) == rank // 2
assert dist.get_rank(parallel_info_dict[1].dp_group) == rank
# group creation passed
model = nn.ModuleList([exp0, exp1, exp2, exp3])
model = model.to(get_current_device())
sync_moe_model_param(model)
assert_equal_in_group(exp0.experts[0].weight.data, dist_dict[1].dp_group)
assert_equal_in_group(exp0.experts[0].bias.data, dist_dict[1].dp_group)
assert_equal_in_group(exp0.experts[0].weight.data, parallel_info_dict[1].dp_group)
assert_equal_in_group(exp0.experts[0].bias.data, parallel_info_dict[1].dp_group)
# MOE experts layout success when ep_size = 1
assert_equal_in_group(exp1.experts[0].weight.data, dist_dict[2].dp_group)
assert_equal_in_group(exp1.experts[0].bias.data, dist_dict[2].dp_group)
assert_equal_in_group(exp1.experts[0].weight.data, parallel_info_dict[2].dp_group)
assert_equal_in_group(exp1.experts[0].bias.data, parallel_info_dict[2].dp_group)
# MOE experts layout success when ep_size = 2
@pytest.mark.dist
def test_moe_initialization():
world_size = 4
run_func = partial(run_test, world_size=world_size, port=free_port())
run_func = partial(run_test, port=free_port())
mp.spawn(run_func, nprocs=world_size)