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

[pipeline]add pipeline policy and bert forward (#4130)

Browse Source 
* add pipeline policy and bert forward to be done

* add bertmodel pipeline forward and make tests

* add Bert_Policy and test for policy

* update formatting

* update formatting

* update the code

* fix bugs

* fix name confilt
This commit is contained in:
Jianghai 2023-07-04 13:46:16 +08:00 committed by Hongxin Liu
parent 5c897ddb94
commit c552cefa93
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_pipeline/test_stage_manager.py
View File

@ -21,7 +21,7 @@ def check_stage_manager():
1: [0, 1], 1: [0, 1],
2: [2, 3], 2: [2, 3],
3: [2, 3], 3: [2, 3],
} }
pg_mesh = ProcessGroupMesh(DP_SIZE, PP_SIZE) pg_mesh = ProcessGroupMesh(DP_SIZE, PP_SIZE)
stage_manager = PipelineStageManager(pg_mesh, PP_DIM) stage_manager = PipelineStageManager(pg_mesh, PP_DIM)
rank = dist.get_rank() rank = dist.get_rank()