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

[moe] add parallel strategy for shared_expert && fix test for deepseek (#6063)

Browse Source
This commit is contained in:
botbw
2024-09-18 10:09:01 +08:00
committed by GitHub
parent 63314ce4e4
commit 4fa6b9509c
2 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_shardformer/test_model/test_shard_deepseek.py
View File

@@ -20,14 +20,15 @@ from tests.test_moe.moe_utils import assert_loose_close, check_model_equal
NUM_BATCH = 8
NUM_TOK_PER_BATCH, NUM_EXPERTS = 64, 4
NUM_LAYERS = 4
HIDDEN_SIZE_PER_HEAD = 4
HIDDEN_SIZE_PER_HEAD = 8
NUM_HEADS = 8
TOP_K = 2
def run_deepseek_commom(config: Tuple[int, ...]):
def run_deepseek_commom(parallel_config: Tuple[int, ...]):
Randomizer.reset_index()
stage, ep_size, pp_size, tp_size, sp_size = config
print(f"rank {dist.get_rank()} testing {parallel_config}")
stage, ep_size, pp_size, tp_size, sp_size = parallel_config
world_size = dist.get_world_size()
rank = dist.get_rank()
dtype, precision = torch.bfloat16, "bf16"
@@ -65,6 +66,7 @@ def run_deepseek_commom(config: Tuple[int, ...]):
attn_implementation="flash_attention_2",
torch_dtype="float16",
n_routed_experts=NUM_EXPERTS,
n_shared_experts=2,
num_experts_per_tok=TOP_K,
trust_remote_code=True,
)
@@ -159,7 +161,7 @@ def run_deepseek_commom(config: Tuple[int, ...]):
if rank == world_size - 1:
shutil.rmtree(model_dir)
print(f"rank {dist.get_rank()} test passed")
print(f"rank {dist.get_rank()} passed {parallel_config}")
@parameterize(