[Design Doc] HybridModel 1F1B Overlap Support
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
`HybridModel` (see #4620) builds hybrid models out of separate decoder layers (e.g., SSM, attention, MoE) with each its own module in the decoder stack. Because of that, the fine-grained 1F1B overlap schedule that `GPTModel` uses to hide MoE all-to-all communication behind another microbatch's compute cannot be applied to it: there is no single "layer" that contains both a compute part and a communication part.
This design refactors hybrid layers into combined units (via `HybridStack`) and adds the schedule callables and pipeline support needed to enable 1F1B expert-parallel overlap for hybrid models, covering PP/VPP and Megatron-FSDP. The implementation is under review in #4798.
Please feel free to read our [design doc](https://github.com/user-attachments/files/30103872/HybridModel.1F1B.Overlap.Support.pdf) and ask any follow-up questions in this issue!
Authors: @Wohox (Pingtian Li), @Connor-XY (Yan Xu), @guihong-nv (Guihong Li)
Contributor guide
Research direction
Start by reading the HybridModel background in #4620 and the linked HybridModel 1F1B Overlap Support design document. Then review the implementation under review in #4798, focusing on HybridStack, schedule callables, pipeline support, PP/VPP, and Megatron-FSDP. Done means the proposed hybrid-model support enables 1F1B expert-parallel overlap across the stated configurations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100