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

[zero] solve hang

Browse Source
This commit is contained in:
botbw
2024-07-09 08:14:00 +00:00
committed by hxwang
parent 2431694564
commit b303ffe9f3
8 changed files with 218 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -113,65 +113,43 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
@parameterize(
"test_config",
[
{
"tp_size": 2,
"pp_size": 1,
"ep_size": 1,
"zero_stage": 2,
"precision": "fp32",
}, # [dp(2) + tp(2)] + [moe_dp(4)]
{
"tp_size": 2,
"pp_size": 1,
"ep_size": 2,
"zero_stage": 2,
"precision": "fp32",
}, # [dp(2) + tp(2)] + [ep(2) + moe_dp(2)]
{
"tp_size": 1,
"pp_size": 2,
"num_microbatches": 2,
"pp_size": 1,
"ep_size": 1,
"zero_stage": 2,
"precision": "fp32",
}, # [dp(2) + pp(2)] + [moe_dp(4)]
{
"tp_size": 1,
"pp_size": 2,
"num_microbatches": 2,
"ep_size": 1,
"zero_stage": 2,
"precision": "fp32",
}, # [dp(2) + pp(2)] + [moe_dp(4)]
{
"tp_size": 1,
"pp_size": 2,
"num_microbatches": 2,
"ep_size": 4,
"zero_stage": 2,
"precision": "fp32",
}, # [dp(2) + pp(2)] + [ep(4))]
{
"tp_size": 1,
"pp_size": 1,
"ep_size": 2,
"zero_stage": 2,
"precision": "fp32",
}, # [dp(4)] + [ep(2) + moe_tp(2)]
{
"tp_size": 1,
"pp_size": 1,
"ep_size": 4,
"zero_stage": 2,
"precision": "fp32"
}, # full dp for non-moe and full ep for moe
{
"tp_size": 1,
"pp_size": 1,
"ep_size": 1,
"zero_stage": 2,
"precision": "fp32"
}, # full dp for moe and non-moe
# {
# "tp_size": 1,
# "pp_size": 2,
# "num_microbatches": 2,
# "ep_size": 1,
# "zero_stage": 1,
# "precision": "fp32",
# }, # [dp(2) + pp(2)] + [moe_dp(4)]
# {
# "tp_size": 1,
# "pp_size": 2,
# "num_microbatches": 2,
# "ep_size": 4,
# "zero_stage": 1,
# "precision": "fp32",
# }, # [dp(2) + pp(2)] + [ep(4))]
# {
# "tp_size": 1,
# "pp_size": 1,
# "ep_size": 2,
# "zero_stage": 0,
# "precision": "fp32",
# }, # [dp(4)] + [ep(2) + moe_tp(2)]
# {
# "tp_size": 1,
# "pp_size": 1,
# "ep_size": 4,
# "zero_stage": 0,
# "precision": "fp32"
# }, # full dp for non-moe and full ep for moe
],
)
def run_mixtral_test(test_config):