github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 19:13:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

[moe] support optimizer checkpoint (#5015)

Browse Source 
* Refactor MoE Manager setup method

* unshard optim ckpt

* optim io

* update transformer version

* update requirements

* update ckpt

* update ckpt

* update ckpt

* fix engine

* fix engine
This commit is contained in:
Xuanlei Zhao
2023-11-08 23:07:03 +08:00
committed by GitHub
parent 67f5331754
commit f71e63b0f3
20 changed files with 738 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_moe/test_grad_handler.py
View File

@@ -6,10 +6,9 @@ import torch.nn as nn
import colossalai
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 tests.test_moe.moe_utils import MoeGradientHandler, assert_not_equal_in_group
from tests.test_moe.moe_utils import MoeGradientHandler
BATCH_SIZE = 4
DIM = 16
@@ -25,7 +24,7 @@ def run_test(rank, world_size, port):
backend="nccl",
)
MOE_MANAGER.setup(42, parallel="EP") # MOE initialization
MOE_MANAGER.setup(parallel="EP") # MOE initialization
num_experts_list = [1, 2, 4]
layer_list = []
for num_experts in num_experts_list:
@@ -41,15 +40,6 @@ def run_test(rank, world_size, port):
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)
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)