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

[chore] arg pass & remove drop token

Browse Source
This commit is contained in:
hxwang
2024-07-12 09:08:16 +00:00
committed by Hongxin Liu
parent 8dbb86899d
commit 102b784a10
5 changed files with 41 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_moe/test_moe_checkpoint.py
View File

@@ -13,7 +13,6 @@ from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM
import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
from colossalai.checkpoint_io import MoECheckpointIO
from colossalai.tensor.moe_tensor.api import is_moe_tensor
from colossalai.testing import parameterize, spawn
from colossalai.testing.utils import spawn
@@ -115,7 +114,6 @@ def check_moe_checkpoint(test_config):
pp_size=2,
ep_size=2,
tp_size=1,
checkpoint_io=MoECheckpointIO,
microbatch_size=1,
zero_stage=1,
)