[pipeline] add stage manager (#4093)

* [pipeline] add stage manager * [test] add pipeline stage manager test * [pipeline] add docstring for stage manager
2023-06-27 16:17:01 +08:00 · 2023-06-27 16:17:01 +08:00 · 5c897ddb94
parent e8e7e49243
commit 5c897ddb94
1 changed files with 1 additions and 1 deletions
--- 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()