Skip to content

Wire pipes into torchcomms conda build and update toml configs #966

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
snarayankh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Wire pipes into torchcomms conda build and update toml configs #966

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ option(USE_RCCLX "Whether to build RCCLX or not" OFF)
option(USE_XCCL "Whether to build XCCL or not" OFF)
option(USE_TRANSPORT "Whether to build TRANSPORT or not" ON)
option(USE_TRITON "Whether to build Triton device bitcode or not" OFF)
option(USE_PIPES "Whether to link against libpipes.so for dispatch/combine" OFF)
option(BUILD_TESTS "Whether to build tests or not" OFF)
message(STATUS " USE_NCCL : ${USE_NCCL}")
message(STATUS " USE_NCCLX : ${USE_NCCLX}")
Expand All @@ -23,6 +24,7 @@ message(STATUS " USE_RCCLX : ${USE_RCCLX}")
message(STATUS " USE_XCCL : ${USE_XCCL}")
message(STATUS " USE_TRANSPORT : ${USE_TRANSPORT}")
message(STATUS " USE_TRITON : ${USE_TRITON}")
message(STATUS " USE_PIPES : ${USE_PIPES}")
message(STATUS " BUILD_TESTS : ${BUILD_TESTS}")

if(NOT DEFINED ENV{TORCH_CUDA_ARCH_LIST})
Expand Down Expand Up @@ -224,6 +226,23 @@ if (USE_TRITON)
endif()
include(comms/torchcomms/hooks/fr/CMakeLists.txt)

# --- Pipes (optional: device API headers + libpipes.so) ---
if (USE_PIPES)
# Existence check only -- pipes is header-only for
# torchcomms (no linking needed). The ncclx backend
# includes pipes device headers (CopyUtils.cuh etc.)
# but does not call into libpipes.so at link time;
# all pipes functionality is invoked at runtime by
# MCCL which links libpipes.so directly.
find_library(PIPES_LIBRARY pipes PATHS ${CONDA_LIB} REQUIRED)
message(STATUS "Found libpipes.so: ${PIPES_LIBRARY}")

# Make pipes headers available to the ncclx target (for CopyUtils.cuh etc.)
if(TARGET torchcomms_comms_ncclx)
target_compile_definitions(torchcomms_comms_ncclx PRIVATE USE_PIPES=1)
endif()
endif()

# Install targets to Python package structure
install(TARGETS torchcomms
LIBRARY DESTINATION .
Expand Down
2 changes: 2 additions & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ def flag_str(val: bool):
# Transport is CUDA-only; disable by default on ROCm but allow explicit opt-in.
USE_TRANSPORT = flag_enabled("USE_TRANSPORT", not IS_ROCM)
USE_TRITON = flag_enabled("USE_TRITON", False)
USE_PIPES = flag_enabled("USE_PIPES", False)

requirement_path = os.path.join(ROOT, "requirements.txt")
try:
Expand Down Expand Up @@ -132,6 +133,7 @@ def build_cmake(self, ext):
f"-DUSE_XCCL={flag_str(USE_XCCL)}",
f"-DUSE_TRANSPORT={flag_str(USE_TRANSPORT)}",
f"-DUSE_TRITON={flag_str(USE_TRITON)}",
f"-DUSE_PIPES={flag_str(USE_PIPES)}",
]
build_args = ["--", "-j"]

Expand Down
Loading