From 5c897ddb9433d034db937113839a805ec74e243e Mon Sep 17 00:00:00 2001 From: Hongxin Liu Date: Tue, 27 Jun 2023 16:17:01 +0800 Subject: [PATCH] [pipeline] add stage manager (#4093) * [pipeline] add stage manager * [test] add pipeline stage manager test * [pipeline] add docstring for stage manager --- tests/test_pipeline/test_stage_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pipeline/test_stage_manager.py b/tests/test_pipeline/test_stage_manager.py index 67a2e9053..be4591d58 100644 --- a/tests/test_pipeline/test_stage_manager.py +++ b/tests/test_pipeline/test_stage_manager.py @@ -21,7 +21,7 @@ def check_stage_manager(): 1: [0, 1], 2: [2, 3], 3: [2, 3], - } + } pg_mesh = ProcessGroupMesh(DP_SIZE, PP_SIZE) stage_manager = PipelineStageManager(pg_mesh, PP_DIM) rank = dist.get_rank()