mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
108a75b454
Zero bubble can be expressed through `ScheduleFlexibleInterleaved1F1B` by setting `enable_zero_bubble=True`. But instead of having to include this flag in schedule initialization we should create a separate ZeroBubbleSchedule and also transition `Interleaved1F1B` to derive from `ScheduleFlexibleInterleaved1F1B`. Then we dont need to expose `ScheduleFlexibleInterleaved1F1B` since the naming is not obvious Pull Request resolved: https://github.com/pytorch/pytorch/pull/133467 Approved by: https://github.com/wconstab ghstack dependencies: #132691
29 lines
661 B
Python
29 lines
661 B
Python
# Copyright (c) Meta Platforms, Inc. and affiliates
|
|
from ._IR import Pipe, pipe_split, pipeline, SplitPoint
|
|
from .schedules import (
|
|
_ScheduleForwardOnly,
|
|
Schedule1F1B,
|
|
ScheduleFlexibleInterleaved1F1B,
|
|
ScheduleGPipe,
|
|
ScheduleInterleaved1F1B,
|
|
ScheduleInterleavedZeroBubble,
|
|
ScheduleLoopedBFS,
|
|
)
|
|
from .stage import build_stage, PipelineStage
|
|
|
|
|
|
__all__ = [
|
|
"Pipe",
|
|
"pipe_split",
|
|
"SplitPoint",
|
|
"pipeline",
|
|
"PipelineStage",
|
|
"build_stage",
|
|
"Schedule1F1B",
|
|
"ScheduleFlexibleInterleaved1F1B",
|
|
"ScheduleGPipe",
|
|
"ScheduleInterleaved1F1B",
|
|
"ScheduleLoopedBFS",
|
|
"ScheduleInterleavedZeroBubble",
|
|
]
|