ColossalAI/colossalai/pipeline/rpc
Ziyue Jiang 59e343328d
[Pipeline Middleware ] Fix deadlock when num_microbatch=num_stage (#2156)
* add splitter

* polish code

* remove comment

* fix async nan by moving to cpu first

Co-authored-by: Ziyue Jiang <ziyue.jiang@gmail.com>
2022-12-23 11:38:43 +08:00
..
__init__.py [pipeline/pytree] add pytree to process args and kwargs | provide `data_process_func` to process args and kwargs after forward (#1642) 2022-09-29 10:58:58 +08:00
_pipeline_base.py [Pipeline Middleware ] Fix deadlock when num_microbatch=num_stage (#2156) 2022-12-23 11:38:43 +08:00
_pipeline_schedule.py [Pipeline Middleware ] Fix deadlock when num_microbatch=num_stage (#2156) 2022-12-23 11:38:43 +08:00
utils.py [Pipeline Middleware ] Fix deadlock when num_microbatch=num_stage (#2156) 2022-12-23 11:38:43 +08:00