Skip to content

Fix PyTorch CUDACachingAllocator::snapshot() API compatibility for conda builds #65

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
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix PyTorch CUDACachingAllocator::snapshot() API compatibility for conda builds #65

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
1 change: 1 addition & 0 deletions comms/torchcomms/nccl/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ add_library(torchcomms_comms_nccl MODULE
${TORCHCOMMS_NCCL_SOURCES}
${TORCHCOMMS_CUDA_API_SOURCE}
)
target_compile_definitions(torchcomms_comms_nccl PRIVATE TORCHCOMMS_CONDA_BUILD)
set_target_properties(torchcomms_comms_nccl PROPERTIES
PREFIX ""
OUTPUT_NAME "_comms_nccl"
Expand Down
4 changes: 4 additions & 0 deletions comms/torchcomms/nccl/TorchCommNCCLCCA.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,11 @@ void CachingAllocatorHookImpl::registerMemPreHook() {
int device = c10::cuda::current_device();
// We assume no mem pool and no comm has been created yet, we just loop up the
// snapshot of the default pool for the current device.
#ifdef TORCHCOMMS_CONDA_BUILD
auto snapshot = c10::cuda::CUDACachingAllocator::snapshot();
#else
auto snapshot = c10::cuda::CUDACachingAllocator::snapshot({device, 0});
#endif
for (const auto& segmentInfo : snapshot.segments) {
// NOLINTNEXTLINE(performance-no-int-to-ptr)
void* addr = reinterpret_cast<void*>(segmentInfo.address);
Expand Down
6 changes: 5 additions & 1 deletion comms/torchcomms/ncclx/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,11 @@ add_library(torchcomms_comms_ncclx MODULE
${TORCHCOMMS_NCCLX_SOURCES}
${TORCHCOMMS_CUDA_API_SOURCE}
)
target_compile_definitions(torchcomms_comms_ncclx PRIVATE MOCK_SCUBA_DATA CTRAN_DISABLE_TCPDM)
target_compile_definitions(torchcomms_comms_ncclx PRIVATE
MOCK_SCUBA_DATA
CTRAN_DISABLE_TCPDM
TORCHCOMMS_CONDA_BUILD
)
set_target_properties(torchcomms_comms_ncclx PROPERTIES
PREFIX ""
OUTPUT_NAME "_comms_ncclx"
Expand Down
4 changes: 4 additions & 0 deletions comms/torchcomms/ncclx/TorchCommNCCLXCCA.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,11 @@ void CachingAllocatorHookImpl::registerMemPreHook() {
int device = c10::cuda::current_device();
// We assume no mem pool and no comm has been created yet, we just loop up the
// snapshot of the default pool for the current device.
#ifdef TORCHCOMMS_CONDA_BUILD
auto snapshot = c10::cuda::CUDACachingAllocator::snapshot();
#else
auto snapshot = c10::cuda::CUDACachingAllocator::snapshot({device, 0});
#endif
for (const auto& segmentInfo : snapshot.segments) {
// NOLINTNEXTLINE(performance-no-int-to-ptr)
void* addr = reinterpret_cast<void*>(segmentInfo.address);
Expand Down
Loading