NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Low contention nvidia symmetric_memory all2all for up to 2x improvement on H100s

Open
#1,787 2 comments 0 reactions 0 assignees View on GitHub
community-request enhancement module: moe
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.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

* All2All Symmetric Memory implementation reduces contention and can provide a 2x speedup to all2all performance which is a critical bottleneck in many MoE implementations.

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

Implement the low contention all2all kernel with nvshmem. You can build it with the PyPi

More info and benchmarks on this thread: https://x.com/cloneofsimo/status/1962795533933912158

Image

Image

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Proposed implementation**
If you have a proposed implementation for the feature state it here or link to a PR.

**Additional context**
Add any other context or screenshots about the feature request here.

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.