NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Increasing throughput with vpp > 1

Open
#1,535 3 comments 1 reaction 0 assignees View on GitHub
bug
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**
When we turn vpp > 1, I see a throughput drop as compared to not using vpp. The primary reason for this is the large "all_gather" latency incurred at the start of each vpp stage (maroon bars in image-1). Also, we see that the 1st “all_gather” operation for each vpp stage takes a long time, with subsequent “all_gather” happens instantly (2nd picture). So, for each model chunk corresponding to the vpp stages, the first time a batch flows through it, it takes time and all subsequent passes through the model chunk for all microbatches happens instantly.

The reason I believe that the 1st “all_gather” takes a long time, is because the first vpp stage is responsible for setting up the communication buffers and memory allocations, which might be taking that longer time. All subsequent batches reuse this established stream.

**Describe the solution you'd like**
I haven’t depp dived much into this. But, I would believe that by default Megatron/TE dont reuse the ‘all_gather’ buffers across vpp stages (hence the above phenemenon), such that each virtual stage is treated independently in terms of buffer lifecycle and microbatch processing. Instead, can we ensure that we pay this cost once per vpp stage, such that we dont pay this penalty for the next vpp model chunk on the same physical device? Basically, can we amortize the cost of the first expensive all_gather across all VPP stages.

**Describe alternatives you've considered**
We see the throughput drop if the number of microbatches aren't enough to cover for the time it takes for the all_Gather operations to take place across all the vpp stages, and this phenomenon only increases with vpp increasing (which is counter-intuitive).

![Image](https://github.com/user-attachments/assets/e6256e8c-efe5-420e-bca7-428ee016c093)

![Image](https://github.com/user-attachments/assets/c5f3353b-d956-41c2-a483-b13c89e75ce0)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.