[feat] add zerobubble pp (just a frame now); add POC test for dx_dw; …

…add test for zerobubble;
duanjunwen · Aug 22, 2024 · ee9baed · ee9baed
1 parent 75c9636
commit ee9baed
Show file tree

Hide file tree

Showing 6 changed files with 2,628 additions and 1 deletion.
diff --git a/colossalai/pipeline/__init__.py b/colossalai/pipeline/__init__.py
@@ -1,11 +1,12 @@
 from .p2p import PipelineP2PCommunication
-from .schedule import InterleavedSchedule, OneForwardOneBackwardSchedule, PipelineSchedule
+from .schedule import InterleavedSchedule, OneForwardOneBackwardSchedule, PipelineSchedule, ZeroBubbleVPipeScheduler
 from .stage_manager import PipelineStageManager
 
 __all__ = [
     "PipelineSchedule",
     "OneForwardOneBackwardSchedule",
     "InterleavedSchedule",
+    "ZeroBubbleVPipeScheduler",
     "PipelineP2PCommunication",
     "PipelineStageManager",
 ]
diff --git a/colossalai/pipeline/schedule/__init__.py b/colossalai/pipeline/schedule/__init__.py
@@ -1,9 +1,11 @@
 from .base import PipelineSchedule
 from .interleaved_pp import InterleavedSchedule
 from .one_f_one_b import OneForwardOneBackwardSchedule
+from .zero_bubble_pp import ZeroBubbleVPipeScheduler
 
 __all__ = [
     "PipelineSchedule",
     "OneForwardOneBackwardSchedule",
     "InterleavedSchedule",
+    "ZeroBubbleVPipeScheduler",
 ]