github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 01:28:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

polish pp middleware (#2476)

Browse Source 
Co-authored-by: Ziyue Jiang <ziyue.jiang@gmail.com>
This commit is contained in:
Ziyue Jiang
2023-01-13 16:56:01 +08:00
committed by GitHub
parent a5dc4253c6
commit fef5c949c3
3 changed files with 3 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -211,7 +211,7 @@ class WorkerBase(ABC):
refcount = 0
with self.output_list_condition_lock:
if refcount < lifecycle:
if refcount <= lifecycle:
self.output_list[key] = output_work_item
self.output_list_condition_lock.notify_all()
@@ -390,7 +390,7 @@ class WorkerBase(ABC):
subscribe_forward_futures[target_index] = []
else:
subscribe_forward_futures[target_index] = producer_worker_rref.rpc_async().get_output_by_key(
producer_output_key, rank=self.pp_rank)
producer_output_key, rank=self.pp_rank, offsets=offsets)
else:
for i in range(producer_num):

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

@@ -29,9 +29,6 @@ class FillDrainWorker(WorkerBase):
target_key = UniqueKey(target_microbatch_id, target_phase)
with self.work_list_condition_lock:
self.work_list_condition_lock.wait_for(lambda: target_key in self.work_list)
return target_key