[VPP] Defer embedding initialization sync until all local model chunks are built
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
### Describe the bug
With pipeline parallelism and virtual pipeline parallelism enabled, each pipeline rank constructs its local virtual-pipeline model chunks sequentially.
However, tied-embedding initialization synchronization is performed while an individual chunk is still being constructed.
During construction, `LanguageModule.setup_embeddings_and_output_layer()` prepares duplicated embedding weights and immediately performs the initial tied-embedding `all_reduce`。The first and last pipeline ranks reach the same embedding-group collective while constructing different virtual stages。
For example, with PP=6 and VPP=3, construction proceeds approximately as:
```text
PP0: build VP0 -> embedding all_reduce (wait) -> build VP1 -> build VP2
PP5: build VP0 -> build VP1 -> build VP2 -> embedding all_reduce
```
PP0 cannot continue constructing VP1 and VP2 until PP5 finishes constructing its earlier chunks and reaches the matching collective.
Both supported VPP construction paths iterate over local VP chunks sequentially:
- [`training.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/training.py#L2374-L2399)
- [`dist_utils.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/models/dist_utils.py#L417-L464)
### Steps/Code to reproduce bug
Use a model configuration with:
- `pipeline_model_parallel_size > 1`
- `virtual_pipeline_model_parallel_size > 1`
- Multi-Token Prediction (MTP) enabled, for example `--mtp-num-layers 1`
### Expected behavior
Initial embedding synchronization should not block construction of the remaining local VP chunks. Each pipeline rank should be able to continue building its local virtual stages without waiting for other ranks reach the matching collective.
```
PP0: build VP0 -> build VP1 -> build VP2 -> embedding all_reduce
PP5: build VP0 -> build VP1 -> build VP2 -> embedding all_reduce
```
### Proposed change
Defer tied-embedding initialization synchronization until all local VPP model chunks have been constructed, then perform the synchronization before continuing model setup.
Contributor guide
Research direction
Start with the VPP chunk-construction loops in training.py and megatron/training/models/dist_utils.py, then trace LanguageModule.setup_embeddings_and_output_layer() where the initial tied-embedding all_reduce occurs. Reproduce with pipeline and virtual pipeline parallelism greater than one and MTP enabled. Done means all local VP chunks are built before the initialization synchronization, without the ranks blocking during intermediate chunk construction.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100