NVIDIA / NVIDIA/Megatron-LM

Support MXFP8 with MFSDP v2 block-atomic sharding

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

Description

## Goal

Add MXFP8 support to MFSDP v2 while preserving the block-atomic sharding format required by MXFP8 block scaling.

## Motivation

MFSDP v2 needs to support real Megatron Core training configurations that use MXFP8. MXFP8 has block-wise scale/data locality requirements, so a naive parameter shard boundary can split a quantization block and create correctness or performance problems.

The design should make MXFP8 sharding explicit instead of silently relying on shard boundaries that happen to be aligned.

## Background

MFSDP v2 currently lives under `megatron_fsdp/experimental` and is being developed around a per-module `fully_shard()` API.

Existing Megatron-FSDP and Megatron Core paths already contain MXFP8-specific handling for quantized parameters, parameter all-gather, transpose buffers, and memory reuse. MFSDP v2 should define the minimal MXFP8-compatible behavior for the new block-atomic sharding path rather than inheriting all existing v1 performance knobs at once.

## Design question

The first implementation needs to choose how to keep MXFP8 blocks atomic across MFSDP v2 shards.

Two candidate approaches:

- Pad more aggressively so each shard boundary is block-aligned.
- Allow uneven all-gather / reduce-scatter sizes so each rank owns only complete MXFP8 blocks without padding everything to uniform shard sizes.

This is a tradeoff between memory usage and latency:

- More aggressive padding keeps collectives simpler and more uniform, but increases memory footprint.
- Uneven all-gather / reduce-scatter can reduce padding overhead, but may increase latency and implementation complexity.

## Intended outcome

- Define the block-atomic sharding format for MXFP8 parameters in MFSDP v2.
- Decide whether the first implementation should prefer padding or uneven collectives.
- Preserve the mapping between MXFP8 data blocks and their scale metadata across sharding, all-gather, reduce-scatter, and parameter update paths.

## Notes

This issue is specifically about MXFP8 support in MFSDP v2. It should be tracked under the `MFSDP v2` project and should not use the old `MFSDPv2` label or broad `module:*` labels.

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.