github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 11:32:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

[moe] merge moe into main (#4978)

Browse Source 
* update moe module
* support openmoe
This commit is contained in:
Xuanlei Zhao
2023-11-02 10:21:24 +08:00
committed by GitHub
parent 8993c8a817
commit dc003c304c
67 changed files with 7618 additions and 1657 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
tests/test_moe/test_grad_handler.py
View File

@@ -4,40 +4,58 @@ import torch.distributed as dist
import torch.nn as nn
import colossalai
from colossalai.context.moe_context import MOE_CONTEXT
from colossalai.legacy.engine.gradient_handler import MoeGradientHandler
from colossalai.nn.layer.moe import Experts, MoeLayer, Top1Router, UniformNoiseGenerator
from colossalai.moe import SparseMLP
from colossalai.moe.manager import MOE_MANAGER
from colossalai.moe.utils import sync_moe_model_param
from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn
from colossalai.utils import get_current_device
from colossalai.utils.moe import sync_moe_model_param
from tests.test_moe.moe_utils import MoeGradientHandler, assert_not_equal_in_group
BATCH_SIZE = 4
DIM = 16
CONFIG = dict()
def run_test(rank, world_size, port):
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=DIM, out_features=DIM, device=get_current_device())
colossalai.launch(
config=dict(),
rank=rank,
world_size=world_size,
host="localhost",
port=port,
backend="nccl",
)
MOE_CONTEXT.setup(42) # MOE initialization
noisy_func = UniformNoiseGenerator()
router = Top1Router(noisy_func=noisy_func)
MOE_MANAGER.setup(42, parallel="EP") # MOE initialization
num_experts_list = [1, 2, 4]
layer_list = []
for num_experts in num_experts_list:
exp = Experts(expert_module, num_experts, **expert_factor)
moe_layer = MoeLayer(DIM, num_experts, router, exp)
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_current_device())
dist_dict = MOE_MANAGER.parallel_info_dict
assert_not_equal_in_group(layer_list[0].experts.wi.data, dist_dict[1].dp_group)
assert_not_equal_in_group(layer_list[0].experts.wo.data, dist_dict[1].dp_group)
assert_not_equal_in_group(layer_list[1].experts.wi.data, dist_dict[2].dp_group)
assert_not_equal_in_group(layer_list[1].experts.wo.data, dist_dict[2].dp_group)
assert_not_equal_in_group(layer_list[2].experts.wi.data, dist_dict[4].dp_group)
assert_not_equal_in_group(layer_list[2].experts.wo.data, dist_dict[4].dp_group)
sync_moe_model_param(model)
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)
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)
@@ -47,17 +65,18 @@ def run_test(rank, world_size, port):
data = torch.randn(BATCH_SIZE, DIM, device=get_current_device())
grad = torch.randn_like(data)
MOE_CONTEXT.reset_loss()
MOE_MANAGER.reset_loss()
for layer in layer_list:
data, _ = layer(data)
data = layer(data)
data.backward(grad)
grad_handler.handle_gradient()
assert_equal_in_group(layer_list[0].experts.experts[0].weight.grad, dist_dict[1].dp_group)
assert_equal_in_group(layer_list[0].experts.experts[0].bias.grad, dist_dict[1].dp_group)
assert_equal_in_group(layer_list[1].experts.experts[0].weight.grad, dist_dict[2].dp_group)
assert_equal_in_group(layer_list[1].experts.experts[0].bias.grad, dist_dict[2].dp_group)
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