Skip to content

[C/PyTorch] Userbuffers and comm+GEMM overlap algorithms refactored and moved to TE/common - #1067

Merged
denera merged 36 commits into
NVIDIA:mainfrom
denera:comm-gemm-overlap-refactor
Oct 29, 2024
Merged

[C/PyTorch] Userbuffers and comm+GEMM overlap algorithms refactored and moved to TE/common#1067
denera merged 36 commits into
NVIDIA:mainfrom
denera:comm-gemm-overlap-refactor

Conversation

@denera

@deneradenera commented Jul 31, 2024•
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This PR moves Userbuffers and comm+GEMM overlap algorithms from TE/PyTorch to TE/common with refactored interfaces to remove the PyTorch dependency.

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 refractor

Changes

  • transformer_engine/pytorch/csrc/userbuffers moved to transformer_engine/common/comm_gemm_overlap/userbuffers.
  • transformer_engine/pytorch/csrc/comm_gemm_overlap.h split into transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h and transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp and refactored to remove torch::Tensor dependency.
  • Added new TE/PyTorch wrappers around the refactored comm+GEMM overlap algorithms.
  • Expanded unit tests to cover all overlap algorithms including atomic GEMM overlaps (tested as AG+RS pairs).

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

@denera
deneraforce-pushed the comm-gemm-overlap-refactor branch from 7255ca5 to 7c0cc8dCompareJuly 31, 2024 18:44
@deneradenera self-assigned this Jul 31, 2024
@deneradenera added the enhancement New feature or request label Jul 31, 2024
timmoon10
timmoon10 reviewed Aug 1, 2024
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt Outdated
Comment threadtransformer_engine/pytorch/module/layernorm_linear.py Outdated
Comment threadtransformer_engine/common/util/pybind_helper.h Outdated
@timmoon10
timmoon10 self-requested a review August 1, 2024 00:36
anderson101866
anderson101866 reviewed Aug 1, 2024
View reviewed changes
Comment threadtransformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp Outdated
timmoon10
timmoon10 reviewed Aug 1, 2024
View reviewed changes

@timmoon10timmoon10 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Overall this looks pretty good. My suggestions are quibbles with the API.

Comment threadtransformer_engine/common/include/transformer_engine/comm_gemm_overlap.h Outdated
Comment threadtransformer_engine/common/include/transformer_engine/activation.h Outdated
Comment threadtransformer_engine/pytorch/cpp_extensions/gemm.py Outdated
Comment threadtransformer_engine/pytorch/cpp_extensions/gemm.py Outdated
@denera
deneraforce-pushed the comm-gemm-overlap-refactor branch from 2e55bb2 to dd8cc21CompareAugust 1, 2024 20:56
ksivaman
ksivaman reviewed Aug 2, 2024
View reviewed changes
Comment threadtransformer_engine/pytorch/cpp_extensions/gemm.py Outdated
@denera
deneraforce-pushed the comm-gemm-overlap-refactor branch 6 times, most recently from b3fb503 to 51519ebCompareAugust 7, 2024 08:25
This was referenced Aug 8, 2024
Merged
Merged
vasunvidia
vasunvidia reviewed Aug 12, 2024
View reviewed changes
Comment threadtransformer_engine/pytorch/module/base.py Outdated
anderson101866
anderson101866 reviewed Aug 20, 2024
View reviewed changes
Comment threadtransformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp Outdated
@denera
deneraforce-pushed the comm-gemm-overlap-refactor branch 8 times, most recently from 4847133 to 10feff5CompareAugust 28, 2024 23:14
pre-commit-ciBotand others added 25 commits October 29, 2024 10:05
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
be46007
for more information, see https://pre-commit.ci
@denera
fixed compile errors when building with NVTE_UB_WITH_MPI=1
af0d4b4
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
69cc14f
for more information, see https://pre-commit.ci
@denera
fixed default bootstrap backend
3ff6a45
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
switched default bootstrap backend priority to MPI > Gloo > NCCL
87ccb55
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
e271f61
for more information, see https://pre-commit.ci
@denera
updated bootstrap backend documentation
a1dbad4
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
close UB bootstrap socket to avoid interfering with CUDA Multicast sh…
815ac11
…areable file handle send/recv
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
added torch::Tensor wrappers for communication buffer and atomic coun…
0c4edad
…ters so PyTorch can factor externally allocated memory into its garbage collection threshold
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
482452c
for more information, see https://pre-commit.ci
@denera
automated handling of world, local and node ranks/sizes within C++ Co…
9cd8c6d
…mmOverlapHelper to simplify Python function signatures
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
34503ce
for more information, see https://pre-commit.ci
@denera
fixed incorrect read of environment variables
ce7f036
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
corrected priority for _SOCKET_IFNAME environment variables in UB boo…
e4a32fb
…tstrapping
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
moved multicast support check to cuda_runtime.h and replaced cudaDevi…
2cb003c
…ceGetProp call with cached sm_count()
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
d6be5f1
for more information, see https://pre-commit.ci
@denera
removed commented out old code and replaced external collective funct…
342b898
…ion type defines with aliases
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
compile-time CUDA version guard for CUDA Driver Multicast attribute
d7b72d7
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
added compile-time CUDA version guards to Multicast code in Userbuffers
1afe8c5
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
1d9d276
for more information, see https://pre-commit.ci
@denera
condensed UB docs, corrected const violations
959dce8
Signed-off-by: Alp Dener <adener@nvidia.com>
@pre-commit-ci@denera
[pre-commit.ci] auto fixes from pre-commit.com hooks
2af2893
for more information, see https://pre-commit.ci
@denera
fixed autodoc rst for UB calls, added CUDA version guard on Multicast…
3b378fc
… UB kernels
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
fixed incorrect UB type reporting for P2P overlaps, comment reformatting
b087b85
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
add docstring to tex.ubuf_built_with_mpi()
d84c462
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
deneraforce-pushed the comm-gemm-overlap-refactor branch from 7b3cf0a to d84c462CompareOctober 29, 2024 15:05
@denera
denera merged commit 933294d into NVIDIA:mainOct 29, 2024
timmoon10 pushed a commit to timmoon10/TransformerEngine that referenced this pull request Nov 7, 2024
@denera
[C/PyTorch] Userbuffers and comm+GEMM overlap algorithms refactored a…
258a3d0
…nd moved to TE/common (NVIDIA#1067)
* moved userbuffers code to TE/common
Signed-off-by: Alp Dener <adener@nvidia.com>
* moved comm+GEMM overlap code to TE/common
Signed-off-by: Alp Dener <adener@nvidia.com>
* removed PyTorch depdency from comm+GEMM overlap in TE/common
Signed-off-by: Alp Dener <adener@nvidia.com>
* added TE/PyTorch wrappers for refactored comm+GEMM overlap code in TE/common
Signed-off-by: Alp Dener <adener@nvidia.com>
* updated TE/PyTorch Python API to match the refactored comm+GEMM overlap code
Signed-off-by: Alp Dener <adener@nvidia.com>
* updated unit tests to work with refactored comm+GEMM overlap code
Signed-off-by: Alp Dener <adener@nvidia.com>
* added a pylint exception to comm+GEMM overlap test runner
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fixing linting errors
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* added documentation for te.initialize_ub
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fixed compile errors when building with NVTE_UB_WITH_MPI=1
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fixed default bootstrap backend
Signed-off-by: Alp Dener <adener@nvidia.com>
* switched default bootstrap backend priority to MPI > Gloo > NCCL
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* updated bootstrap backend documentation
Signed-off-by: Alp Dener <adener@nvidia.com>
* close UB bootstrap socket to avoid interfering with CUDA Multicast shareable file handle send/recv
Signed-off-by: Alp Dener <adener@nvidia.com>
* added torch::Tensor wrappers for communication buffer and atomic counters so PyTorch can factor externally allocated memory into its garbage collection threshold
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* automated handling of world, local and node ranks/sizes within C++ CommOverlapHelper to simplify Python function signatures
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fixed incorrect read of environment variables
Signed-off-by: Alp Dener <adener@nvidia.com>
* corrected priority for _SOCKET_IFNAME environment variables in UB bootstrapping
Signed-off-by: Alp Dener <adener@nvidia.com>
* moved multicast support check to cuda_runtime.h and replaced cudaDeviceGetProp call with cached sm_count()
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* removed commented out old code and replaced external collective function type defines with aliases
Signed-off-by: Alp Dener <adener@nvidia.com>
* compile-time CUDA version guard for CUDA Driver Multicast attribute
Signed-off-by: Alp Dener <adener@nvidia.com>
* added compile-time CUDA version guards to Multicast code in Userbuffers
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* condensed UB docs, corrected const violations
Signed-off-by: Alp Dener <adener@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fixed autodoc rst for UB calls, added CUDA version guard on Multicast UB kernels
Signed-off-by: Alp Dener <adener@nvidia.com>
* fixed incorrect UB type reporting for P2P overlaps, comment reformatting
Signed-off-by: Alp Dener <adener@nvidia.com>
* add docstring to tex.ubuf_built_with_mpi()
Signed-off-by: Alp Dener <adener@nvidia.com>
timmoon10 pushed a commit to timmoon10/TransformerEngine that referenced this pull request Nov 7, 2024
@denera
Merge branch 'blackwell-devel' into 'blackwell-devel'
de230d0
[C/PyTorch] Userbuffers and comm+GEMM overlap algorithms refactored and moved to TE/common (NVIDIA#1067)
See merge request dl/transformerengine/transformerengine!205
@deneradenera mentioned this pull request Mar 9, 2026
Closed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ksivamanksivamanksivaman approved these changes
@ptrendxptrendxAwaiting requested review from ptrendxptrendx is a code owner
@vasunvidiavasunvidiaAwaiting requested review from vasunvidia
+1 more reviewer
@anderson101866anderson101866anderson101866 left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

@deneradenera

Labels

enhancementNew feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@denera@timmoon10@anderson101866@ksivaman@vasunvidia