Skip to content

Comments

MXFP8 support in Userbuffers#1711

Merged
ptrendx merged 55 commits intoNVIDIA:mainfrom
timmoon10:mxfp8-ub-debug
May 7, 2025
Merged

MXFP8 support in Userbuffers#1711
ptrendx merged 55 commits intoNVIDIA:mainfrom
timmoon10:mxfp8-ub-debug

Conversation

@timmoon10
Copy link
Collaborator

@timmoon10 timmoon10 commented Apr 22, 2025 •
edited
Loading

Description

This PR adds basic support for overlapped tensor-parallel communication with MXFP8 data, with the following limitations:

  • Tensor dimensions must be multiples of 128
  • MXFP8 scaling factors are communicated with blocking NCCL collectives
  • The wgrad GEMM + grad output AG case is not supported.

In the process, I've done some significant refactoring:

  • Removed quantized tensor logic from the Userbuffers communicators, so they only take responsibility for the Userbuffers buffer itself.
  • Added helper functions to expose the Userbuffers buffer as a quantized tensor and handle scaling factors
  • Restore support for Userbuffers in te.Sequential
  • Reorganize module implementation for clarity

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Add MXFP8 support in Userbuffers
  • Move Userbuffers logic for quantized tensors out of UB communicator and into helper functions
  • Restore Userbuffers support in te.Sequential

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@timmoon10
Initial work toward restoring UB support in te.Sequential
25bcbda 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Forward UB linear runs, but has numerical error
49c6a02 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Debug UB forward tests
4bacf2a 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Minor tweaks
bd1d50a 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Remove Python checks for MXFP8 UB linear forward
b8b325b 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Add dim check for MXFP8 full tiles
c8b2c51 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Move QuantizedTensor logic out of UB comm and into Python helper func…
9f562b6 
…tion

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Support MXFP8 AGs
c7a5e65 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Coalesce NCCL all-gathers for MXFP8 all-gather
0c1a98f 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Merge branch 'main' into mxfp8-ub-debug
4304ddf
@timmoon10
Merge branch 'main' into optimize-wgrad-allgather
15c34ec
@timmoon10
Initial impl of backward UB linear in te.Sequential
0917b20 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Merge branch 'optimize-wgrad-allgather' into mxfp8-ub-debug
33a3dbb
@timmoon10
Debug UB linear backward with no quantization
a86cbb9 
dgrad GEMM + dx RS is still broken.

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Fix chunk dims for dgrad GEMM + dx RS
3ff2955 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Debugging MXFP8 UB cases
0942ee3 
Still failing with dy AG + wgrad GEMM

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Use NCCL to overlap dy AG with dgrad GEMM
df59ba0 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Merge branch 'main' into mxfp8-ub-debug
1047a09 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Debug UB GEMM tests
ffa65bf 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Initial refactoring of linear module forward
8531147 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10 timmoon10 added the enhancement New feature or request label Apr 22, 2025
@timmoon10
Refactor linear module backward
0767cea 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Debug linear module UB tests
1a03c8f 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Tweak test tensor dims
0639166 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Merge branch 'main' into mxfp8-ub-debug
6f855aa 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Do not store autograd context within wgrad GEMM closure
d4ac2ea 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Fix linter warnings
7803cb0 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Copy link
Collaborator Author

timmoon10 commented Apr 28, 2025

/te-ci pytorch L1

denera
denera previously approved these changes Apr 29, 2025
View reviewed changes
Copy link
Collaborator

@denera denera left a comment

Choose a reason for hiding this comment

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

LGTM

timmoon10 added 2 commits May 1, 2025 18:59
@timmoon10
Merge branch 'main' into mxfp8-ub-debug
dfb53ca 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Fix RS out dims
546e02a 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Copy link
Collaborator Author

timmoon10 commented May 1, 2025

/te-ci L1

@timmoon10
Copy link
Collaborator Author

timmoon10 commented May 2, 2025

/te-ci pytorch L1

ptrendx
ptrendx reviewed May 6, 2025
View reviewed changes
timmoon10 and others added 5 commits May 6, 2025 21:54
@timmoon10
Merge branch 'main' into mxfp8-ub-debug
fe02a7e 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Restore support for internal quantized tensors
19459ec 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Add tests for MXFP8 GEMM with UB
0efa1a9 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
fb5ca6e 
for more information, see https://pre-commit.ci
@timmoon10
Fix linter warnings
70057c6 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Copy link
Collaborator Author

timmoon10 commented May 7, 2025

/te-ci pytorch L1

@timmoon10
Debug test failures
7153804 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Copy link
Collaborator Author

timmoon10 commented May 7, 2025

/te-ci pytorch L1

timmoon10 and others added 2 commits May 7, 2025 07:35
@timmoon10
Debug test failures
d1fc045 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Merge branch 'main' into mxfp8-ub-debug
6994f29
@timmoon10
Copy link
Collaborator Author

timmoon10 commented May 7, 2025

/te-ci pytorch L1

ptrendx
ptrendx approved these changes May 7, 2025
View reviewed changes
Copy link
Member

@ptrendx ptrendx left a comment

Choose a reason for hiding this comment

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

LGTM

@ptrendx ptrendx merged commit ce0b46c into NVIDIA:main May 7, 2025
17 of 24 checks passed
@timmoon10 timmoon10 mentioned this pull request Jun 11, 2025
Merged
13 tasks
@timmoon10 timmoon10 mentioned this pull request Jul 15, 2025
Merged
13 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ptrendx ptrendx ptrendx approved these changes

@denera denera denera left review comments

@ksivaman ksivaman Awaiting requested review from ksivaman

Assignees

No one assigned

Labels

enhancement New feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@timmoon10 @zhongbozhu @denera @ptrendx