Skip to content

Add minimal DBuffer implementation#4835

Merged
wujingyue merged 27 commits into
NVIDIA:mainfrom
wujingyue:dbuffer
Jun 17, 2026
Merged

Add minimal DBuffer implementation#4835
wujingyue merged 27 commits into
NVIDIA:mainfrom
wujingyue:dbuffer

Conversation

@wujingyue

@wujingyue wujingyue commented May 17, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add a minimal Megatron-FSDP DBuffer implementation following http://nv/mfsdp-design
  • Add global flat-buffer layout computation with row-aligned offsets and padded global size.
  • Add unit coverage for construction, tensor views, DTensor conversion, scatter/allgather, allreduce/reduce-scatter, and 2D mesh cases.

For broader context on how DBuffer will be used, refer to #4976.

@copy-pr-bot

copy-pr-bot Bot commented May 17, 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 force-pushed the dbuffer branch 2 times, most recently from 3973505 to 4e9f580 Compare May 18, 2026 05:48
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py Outdated
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py Outdated
@wujingyue wujingyue marked this pull request as ready for review May 18, 2026 06:35
@wujingyue wujingyue requested review from a team as code owners May 18, 2026 06:35
@svcnvidia-nemo-ci svcnvidia-nemo-ci requested a review from a team May 18, 2026 06:35
@wujingyue wujingyue changed the title Add minimal DBuffer prototype Add minimal DBuffer implementation May 18, 2026
Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py Outdated
@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test

wujingyue added 24 commits June 16, 2026 22:12
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test

@svcnvidia-nemo-ci

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27655533140

@svcnvidia-nemo-ci

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27656783364

@svcnvidia-nemo-ci

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27657812515

@svcnvidia-nemo-ci

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27660486306

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

Labels

Approved All necessary approvals have been made complexity: medium MFSDPv2 Run tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants