NVIDIA / NVIDIA/Megatron-LM

[VPP] Defer embedding initialization sync until all local model chunks are built

Open
#6,747 0 comments 0 reactions 0 assignees View on GitHub
bug community-request
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.