Skip to content

Integrate MXFP8 weights into experimental MFSDP - #7114

Open
wujingyue wants to merge 3 commits into
NVIDIA:pull-request/7265from
wujingyue:agent/mfsdp-grouped-dbuffer
Open

Integrate MXFP8 weights into experimental MFSDP#7114
wujingyue wants to merge 3 commits into
NVIDIA:pull-request/7265from
wujingyue:agent/mfsdp-grouped-dbuffer

Conversation

@wujingyue

@wujingyue wujingyue commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

Integrate QuantizedDBuffer into experimental MFSDP for MXFP8 training.

Depends on #7265 (QuantizedDBuffer and buffer tests), which depends on #7257 (DBuffer layout allocation) and #6978 (BlockAtomic placement).

Remaining scope:

  • Group parameters by effective storage dtype and support MXFP8 parameter materialization.
  • Initialize main weights from preserved high-precision values when available; add DBuffer.copy_from for initialization.
  • Quantize optimizer-layout weights and redistribute to compute layouts.
  • Add numerical training comparisons against an unsharded TE model with FusedAdam, including ZeRO-1/3, multiple weights and BF16 biases.

The split preserves the complete source tree of the previously tested #7114 head; only commit ancestry and PR boundaries change. Buffer implementation and its standalone tests are reviewed in #7265.

Limitations: 2D MXFP8 weights with dimensions divisible by 32 and aligned row shards. Compact scales require padded TE allocations when necessary. Two previously identified CPU/meta initialization test failures remain outside this change.

@copy-pr-bot

copy-pr-bot Bot commented Sep 7, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@wujingyue wujingyue changed the title Add MXFP8 grouped DBuffer support Add MXFP8 GroupedDBuffer support on BlockAtomic DBuffer Sep 8, 2026
@wujingyue
wujingyue force-pushed the agent/mfsdp-grouped-dbuffer branch from 56ace1a to b00d660 Compare September 8, 2026 19:26
@wujingyue
wujingyue changed the base branch from main to pull-request/6978 September 8, 2026 19:26
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread tests/unit_tests/distributed/mfsdp_v2/test_grouped_dbuffer.py Outdated
@wujingyue wujingyue changed the title Add MXFP8 GroupedDBuffer support on BlockAtomic DBuffer Add MXFP8 support to experimental MFSDP Sep 11, 2026
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
@wujingyue wujingyue changed the title Add MXFP8 support to experimental MFSDP Add MXFP8 support to MFSDP v2 Sep 11, 2026
@wujingyue
wujingyue force-pushed the agent/mfsdp-grouped-dbuffer branch from f506fb8 to 333855f Compare September 11, 2026 22:02
@wujingyue
wujingyue changed the base branch from pull-request/6978 to pull-request/7257 September 11, 2026 22:02
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/grouped_dbuffer.py Outdated

def redistribute(
self, new_placements: Iterable[Placement], *, out: "QuantizedDBuffer | None" = None
) -> "QuantizedDBuffer":

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this return a new QuantizedDBuffer even when out is specified?

@wujingyue wujingyue changed the title Add MXFP8 support to MFSDP v2 Integrate MXFP8 weights into experimental MFSDP Sep 11, 2026
@wujingyue
wujingyue changed the base branch from pull-request/7257 to pull-request/7265 September 11, 2026 23:13
@wujingyue
wujingyue marked this pull request as ready for review September 11, 2026 23:15
@wujingyue
wujingyue requested a review from a team as a code owner September 11, 2026 23:15
@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test 8689e05

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@wujingyue
wujingyue force-pushed the agent/mfsdp-grouped-dbuffer branch from 8689e05 to 8ee05d3 Compare September 11, 2026 23:18
@wujingyue
wujingyue requested review from a team as code owners September 11, 2026 23:18
@copy-pr-bot

copy-pr-bot Bot commented Sep 11, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test 8ee05d3

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test 38fa0f6

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant