NVIDIA / NVIDIA/Megatron-LM

[Design Doc] HybridModel 1F1B Overlap Support

Open
#5,854 1 comment 1 reaction 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.