mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-05 19:13:01 +00:00
[test] fixed rerun_on_exception and adapted test cases (#487)
This commit is contained in:
@@ -11,6 +11,7 @@ from colossalai.context.moe_context import MOE_CONTEXT
|
||||
from colossalai.utils.moe import sync_moe_model_param
|
||||
from colossalai.engine.gradient_handler import MoeGradientHandler
|
||||
from colossalai.testing import assert_equal_in_group
|
||||
from colossalai.testing import rerun_on_exception
|
||||
|
||||
BATCH_SIZE = 4
|
||||
DIM = 16
|
||||
@@ -62,6 +63,7 @@ def run_test(rank, world_size, port):
|
||||
|
||||
|
||||
@pytest.mark.dist
|
||||
@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*")
|
||||
def test_grad_handler():
|
||||
world_size = 4
|
||||
run_func = partial(run_test, world_size=world_size, port=free_port())
|
||||
|
Reference in New Issue
Block a user