NVIDIA / NVIDIA/Megatron-LM

Support uneven communication for MFSDP v2

Open
#6,587 3 comments 0 reactions 1 assignee Claimed by @XiaodaNV View on GitHub
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Goal

Support communication where data-parallel ranks contribute and receive unequal
amounts of data, without requiring padding to a uniform size.

## Motivation

TensorAtomic placement can assign different amounts of data to each rank, so
its materialization and gradient-reduction paths need uneven communication.

This is a general communication need, not specific to TensorAtomic. MXFP8
block-atomic sharding may also benefit from uneven communication to reduce
padding overhead.

Separating the work lets the solution be designed around the broader set of
callers and their performance requirements.

## Related issues

- #6368 — Add TensorAtomic DBuffer placement
- #5615 — Support MXFP8 with MFSDP v2 block-atomic sharding

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.