mirror of https://github.com/hpcaitech/ColossalAI
81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
|
|
import colossalai
|
|
from colossalai.accelerator import get_accelerator
|
|
from colossalai.moe import SparseMLP
|
|
from colossalai.moe.manager import MOE_MANAGER
|
|
from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn
|
|
from tests.test_moe.moe_utils import MoeGradientHandler
|
|
|
|
BATCH_SIZE = 4
|
|
DIM = 16
|
|
|
|
|
|
def run_test(rank, world_size, port):
|
|
colossalai.launch(
|
|
config=dict(),
|
|
rank=rank,
|
|
world_size=world_size,
|
|
host="localhost",
|
|
port=port,
|
|
backend="nccl",
|
|
)
|
|
|
|
MOE_MANAGER.setup(parallel="EP") # MOE initialization
|
|
num_experts_list = [1, 2, 4]
|
|
layer_list = []
|
|
for num_experts in num_experts_list:
|
|
moe_layer = SparseMLP(
|
|
hidden_size=DIM,
|
|
intermediate_size=DIM * 4,
|
|
num_experts=num_experts,
|
|
router_top_k=1,
|
|
router_noisy_policy="Jitter",
|
|
)
|
|
layer_list.append(moe_layer)
|
|
|
|
model = nn.ModuleList(layer_list)
|
|
model = model.to(get_accelerator().get_current_device())
|
|
dist_dict = MOE_MANAGER.parallel_info_dict
|
|
assert_equal_in_group(layer_list[0].experts.wi.data, dist_dict[1].dp_group)
|
|
assert_equal_in_group(layer_list[0].experts.wo.data, dist_dict[1].dp_group)
|
|
assert_equal_in_group(layer_list[1].experts.wi.data, dist_dict[2].dp_group)
|
|
assert_equal_in_group(layer_list[1].experts.wo.data, dist_dict[2].dp_group)
|
|
assert_equal_in_group(layer_list[2].experts.wi.data, dist_dict[4].dp_group)
|
|
assert_equal_in_group(layer_list[2].experts.wo.data, dist_dict[4].dp_group)
|
|
# MoE model synchronization passed
|
|
|
|
grad_handler = MoeGradientHandler(model, 0)
|
|
|
|
rank = dist.get_rank()
|
|
torch.cuda.manual_seed(78 + rank)
|
|
data = torch.randn(BATCH_SIZE, DIM, device=get_accelerator().get_current_device())
|
|
grad = torch.randn_like(data)
|
|
|
|
MOE_MANAGER.reset_loss()
|
|
for layer in layer_list:
|
|
data = layer(data)
|
|
data.backward(grad)
|
|
grad_handler.handle_gradient()
|
|
|
|
assert_equal_in_group(layer_list[0].experts.wi.grad, dist_dict[1].dp_group)
|
|
assert_equal_in_group(layer_list[0].experts.wo.grad, dist_dict[1].dp_group)
|
|
assert_equal_in_group(layer_list[1].experts.wi.grad, dist_dict[2].dp_group)
|
|
assert_equal_in_group(layer_list[1].experts.wo.grad, dist_dict[2].dp_group)
|
|
assert_equal_in_group(layer_list[2].experts.wi.grad, dist_dict[4].dp_group)
|
|
assert_equal_in_group(layer_list[2].experts.wo.grad, dist_dict[4].dp_group)
|
|
# MoE grad handler test passed
|
|
|
|
|
|
@pytest.mark.dist
|
|
@rerun_if_address_is_in_use()
|
|
def test_grad_handler():
|
|
spawn(run_test, 4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_grad_handler()
|