mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
198b5fd2d4
Added the DualPipeV schedule according to http://github.com/deepseek-ai/DualPipe/blob/main/dualpipe/dualpipev.py#L11 <img width="3633" height="486" alt="image" src="https://github.com/user-attachments/assets/4e843bb9-87cd-4d11-936c-7dfe8ee12f16" /> This schedule doesn't perform the actual "overlap" during execution, but provides the scaffolding and schedule definition we need to run it E2E in torchtitan. Supporting the overlapped operation will be worked on in following PRs. Tests: ```sh python test/distributed/pipelining/test_schedule_multiproc.py -k test_v_shape_schedules python test/distributed/pipelining/test_schedule.py -k test_pipeline_order_for_v_schedules ``` Also tested in TorchTitan and is running. Pull Request resolved: https://github.com/pytorch/pytorch/pull/159591 Approved by: https://github.com/wconstab
31 lines
689 B
Python
31 lines
689 B
Python
# Copyright (c) Meta Platforms, Inc. and affiliates
|
|
from ._IR import Pipe, pipe_split, pipeline, SplitPoint
|
|
from .schedules import (
|
|
_ScheduleForwardOnly,
|
|
Schedule1F1B,
|
|
ScheduleDualPipeV,
|
|
ScheduleGPipe,
|
|
ScheduleInterleaved1F1B,
|
|
ScheduleInterleavedZeroBubble,
|
|
ScheduleLoopedBFS,
|
|
ScheduleZBVZeroBubble,
|
|
)
|
|
from .stage import build_stage, PipelineStage
|
|
|
|
|
|
__all__ = [
|
|
"Pipe",
|
|
"pipe_split",
|
|
"SplitPoint",
|
|
"pipeline",
|
|
"PipelineStage",
|
|
"build_stage",
|
|
"Schedule1F1B",
|
|
"ScheduleGPipe",
|
|
"ScheduleInterleaved1F1B",
|
|
"ScheduleLoopedBFS",
|
|
"ScheduleInterleavedZeroBubble",
|
|
"ScheduleZBVZeroBubble",
|
|
"ScheduleDualPipeV",
|
|
]
|