3 Commits

Author	SHA1	Message	Date
Ke Wen	3090667cf9	[pipelining] pipeline() taking microbatch as example input (#128163) ... Changed the API of `pipeline()` to take microbatch instead of full batch as example args. Main purpose is to: - make this API more atomic; - decouple tracing frontend from runtime info like `num_chunks`. Side effects: - Creates opportunity for varying `num_chunks` of schedules with the same `pipe` object. - User has to create example microbatch input. - Chunk spec stuff are now all moved to runtime side. Pull Request resolved: https://github.com/pytorch/pytorch/pull/128163 Approved by: https://github.com/H-Huang	2024-06-07 15:51:53 +00:00
Ke Wen	07d6ab5aa2	[pipelining] Add pipeline schedules (#125975) ... 1. Add pipeline schedules: - GPipe - 1F1B - Interleaved 1F1B - LoopedBFS 2. Add basic forward and backward tests: test_schedule.py Pull Request resolved: https://github.com/pytorch/pytorch/pull/125975 Approved by: https://github.com/wconstab ghstack dependencies: #125729	2024-05-11 21:17:53 +00:00
Ke Wen	0199ce8d6c	[pipelining] Add microbatch split and merge utils (#125273) ... Pull Request resolved: https://github.com/pytorch/pytorch/pull/125273 Approved by: https://github.com/H-Huang ghstack dependencies: #124776, #124875, #124958	2024-05-02 21:09:47 +00:00