github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-17 23:18:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

[moe] init moe plugin comm setting with sp

Browse Source
This commit is contained in:
hxwang
2024-07-18 08:37:06 +00:00
committed by Hongxin Liu
parent 09d6280d3e
commit 877d94bb8c
7 changed files with 101 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/test_shardformer/test_model/test_shard_mixtral.py
View File

@@ -141,12 +141,12 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
[
# {
# "tp_size": 1,
# "pp_size": 2,
# "pp_size": 1,
# "num_microbatches": 2,
# "ep_size": 2,
# "zero_stage": 1,
# "zero_stage": 0,
# "overlap_communication": False,
# "precision": "fp32",
# "precision": "fp16",
# }, # [dp(4)] + [moe_dp(4)]
# {
# "tp_size": 1,
@@ -169,7 +169,7 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
{ # Ulysess + Flash attention
"tp_size": 1,
"pp_size": 1,
"sp_size": 4,
"sp_size": 2,
"ep_size": 1,
"enable_sequence_parallelism": True,
"sequence_parallelism_mode": "all_to_all",