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

[pipeline] Add Simplified Alpa DP Partition (#2507)

Browse Source 
* add alpa dp split

* add alpa dp split

* use fwd+bwd instead of fwd only

---------

Co-authored-by: Ziyue Jiang <ziyue.jiang@gmail.com>
This commit is contained in:
Ziyue Jiang
2023-03-07 10:34:31 +08:00
committed by GitHub
parent b42d3d28ed
commit 400f63012e
4 changed files with 197 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/pipeline/rpc/_pipeline_base.py
View File

@@ -1115,7 +1115,8 @@ class PipelineEngineBase(ABC, nn.Module):
# let each worker know global worker rref (include itself)
sync_futs = []
for pp_rank in self.pp_rank_to_worker_rref:
fut = self.pp_rank_to_worker_rref[pp_rank].rpc_async().sync_global_worker_rrefs(self.pp_rank_to_worker_rref)
fut = self.pp_rank_to_worker_rref[pp_rank].rpc_async(timeout=0).sync_global_worker_rrefs(
self.pp_rank_to_worker_rref)
sync_futs.append(fut)
for fut in sync_futs: