diff --git a/docs/env_vars.rst b/docs/env_vars.rst index 9dd86e8..e4fae79 100644 --- a/docs/env_vars.rst +++ b/docs/env_vars.rst @@ -12,15 +12,18 @@ CUDECOMP_ENABLE_NCCL_UBR :code:`CUDECOMP_ENABLE_NCCL_UBR` controls whether cuDecomp registers its communication buffers with the NCCL library using :code:`ncclCommRegister`/:code:`ncclCommDeregister` (i.e., user buffer registration). Registration can improve NCCL send/receive performance in some scenarios. See the `User Buffer Registration `_ section of the NCCL documentation for more details. +This option requires CUDA VMM workspace allocations. Setting :code:`CUDECOMP_ENABLE_NCCL_UBR` also requests the cuMem allocation path used by :code:`cudecompMalloc`; if cuMem support is unavailable, NCCL user buffer registration is disabled. See :code:`CUDECOMP_ENABLE_CUMEM` for more details on cuMem workspace allocations. -Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature. +Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature when cuMem support is available. CUDECOMP_ENABLE_CUMEM ------------------------ -(since v0.5.0, requires CUDA 12.3 driver/toolkit or newer) +(since v0.5.0, requires CUDA 11.3 driver/toolkit or newer) -:code:`CUDECOMP_ENABLE_CUMEM` controls whether cuDecomp uses :code:`cuMem*` APIs to allocate fabric-registered workspace buffers via :code:`cudecompMalloc`. This option can improve the performance of -some MPI distributions on multi-node NVLink (MNNVL) capable systems. +:code:`CUDECOMP_ENABLE_CUMEM` controls whether cuDecomp uses :code:`cuMem*` APIs to allocate VMM workspace buffers via :code:`cudecompMalloc`. +The allocations always request POSIX file-descriptor export support. +When built with CUDA 12.3 or newer, running with a CUDA 12.3 or newer driver, and running on a device that supports fabric handles, cuDecomp also requests fabric export support. +This option can improve the performance of some MPI distributions on multi-node NVLink (MNNVL) capable systems. Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature. diff --git a/include/internal/halo.h b/include/internal/halo.h index 105d032..f604657 100644 --- a/include/internal/halo.h +++ b/include/internal/halo.h @@ -148,6 +148,7 @@ void cudecompUpdateHalos_(int ax, const cudecompHandle_t handle, const cudecompG bool input_has_padding = anyNonzeros(padding); if (c == 2 && (input_has_padding || haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend) || + (haloBackendRequiresNccl(grid_desc->config.halo_comm_backend) && handle->nccl_enable_ubr) || (managed && haloBackendRequiresMpi(grid_desc->config.halo_comm_backend)) || (handle->cuda_cumem_enable && comm_info.mnnvl_active && haloBackendRequiresMpi(grid_desc->config.halo_comm_backend)))) { @@ -155,6 +156,7 @@ void cudecompUpdateHalos_(int ax, const cudecompHandle_t handle, const cudecompG // For managed memory, always stage to work space if using MPI. // If using MPI and communicator has MNNVL connections, stage to work space if fabric-allocated. // For any memory, always stage to workspace if using NVSHMEM. + // For NCCL user buffer registration, stage to the registered workspace instead of application input. // Can revisit for NVSHMEM if input is NVSHMEM allocated. c = 1; } diff --git a/include/internal/transpose.h b/include/internal/transpose.h index d68e149..ff4b66c 100644 --- a/include/internal/transpose.h +++ b/include/internal/transpose.h @@ -374,6 +374,11 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c // in to workspace (which should be nvshmem allocated). Can revisit support for input/output // arrays allocated with nvshmem. enable = false; + } else if (transposeBackendRequiresNccl(grid_desc->config.transpose_comm_backend) && handle->nccl_enable_ubr) { + // Note: NCCL user buffer registration requires both the source and destination buffers to be registered. + // cuDecomp only registers the workspace, so keep NCCL communication staged through workspace instead of + // directly using input/output buffers that may not be registered. + enable = false; } else if (transposeBackendRequiresMpi(grid_desc->config.transpose_comm_backend)) { // Note: For MPI, disable special cases if input or output pointers are to managed memory // since MPI performance directly from managed memory is not great diff --git a/src/autotune.cc b/src/autotune.cc index 858660e..6a3d495 100644 --- a/src/autotune.cc +++ b/src/autotune.cc @@ -292,16 +292,32 @@ void autotuneTransposeBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_d transposeWorkspaceGuard(work_nvshmem, {handle, grid_desc, CUDECOMP_TRANSPOSE_COMM_NVSHMEM}); // Check if there is enough memory for separate non-NVSHMEM allocated work buffer + const bool allow_nvshmem_workspace_fallback = !handle->cuda_cumem_enable && !handle->nccl_enable_ubr; auto ret = cudaMalloc(&work, work_sz); - if (ret == cudaErrorMemoryAllocation) { + int any_oom = (ret == cudaErrorMemoryAllocation); + CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &any_oom, 1, MPI_INT, MPI_LOR, handle->mpi_comm)); + if (any_oom) { + if (ret == cudaSuccess) { + CHECK_CUDA(cudaFree(work)); + work = nullptr; + } else if (ret == cudaErrorMemoryAllocation) { + cudaGetLastError(); // Reset CUDA error state + } else { + CHECK_CUDA(ret); + } + if (!allow_nvshmem_workspace_fallback) { + THROW_CUDA_ERROR( + "Cannot allocate separate non-NVSHMEM workspace during autotuning while cuMem or NCCL user buffer " + "registration is enabled."); + } if (handle->rank == 0) { printf("CUDECOMP:WARN: Cannot allocate separate workspace for non-NVSHMEM backends during " "autotuning. Using NVSHMEM allocated workspace for all backends, which may cause issues " "for some MPI implementations. See documentation for more details and suggested workarounds.\n"); } work = work_nvshmem; - cudaGetLastError(); // Reset CUDA error state } else { + CHECK_CUDA(ret); CHECK_CUDA(cudaFree(work)); auto backend = (need_nccl) ? CUDECOMP_TRANSPOSE_COMM_NCCL : CUDECOMP_TRANSPOSE_COMM_MPI_P2P; tmp = grid_desc->config.transpose_comm_backend; @@ -746,16 +762,32 @@ void autotuneHaloBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, work_nvshmem_guard = haloWorkspaceGuard(work_nvshmem, {handle, grid_desc, CUDECOMP_HALO_COMM_NVSHMEM}); // Check if there is enough memory for separate non-NVSHMEM allocated work buffer + const bool allow_nvshmem_workspace_fallback = !handle->cuda_cumem_enable && !handle->nccl_enable_ubr; auto ret = cudaMalloc(&work, work_sz); - if (ret == cudaErrorMemoryAllocation) { + int any_oom = (ret == cudaErrorMemoryAllocation); + CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &any_oom, 1, MPI_INT, MPI_LOR, handle->mpi_comm)); + if (any_oom) { + if (ret == cudaSuccess) { + CHECK_CUDA(cudaFree(work)); + work = nullptr; + } else if (ret == cudaErrorMemoryAllocation) { + cudaGetLastError(); // Reset CUDA error state + } else { + CHECK_CUDA(ret); + } + if (!allow_nvshmem_workspace_fallback) { + THROW_CUDA_ERROR( + "Cannot allocate separate non-NVSHMEM workspace during autotuning while cuMem or NCCL user buffer " + "registration is enabled."); + } if (handle->rank == 0) { printf("CUDECOMP:WARN: Cannot allocate separate workspace for non-NVSHMEM backends during " "autotuning. Using NVSHMEM allocated workspace for all backends, which may cause issues " "for some MPI implementations. See documentation for more details and suggested workarounds.\n"); } work = work_nvshmem; - cudaGetLastError(); // Reset CUDA error state } else { + CHECK_CUDA(ret); CHECK_CUDA(cudaFree(work)); auto backend = (need_nccl) ? CUDECOMP_HALO_COMM_NCCL : CUDECOMP_HALO_COMM_MPI; tmp = grid_desc->config.halo_comm_backend; diff --git a/src/cudecomp.cc b/src/cudecomp.cc index f3981be..1f3be86 100644 --- a/src/cudecomp.cc +++ b/src/cudecomp.cc @@ -389,44 +389,66 @@ static void gatherGlobalMPIInfo(cudecompHandle_t& handle) { static void getCudecompEnvVars(cudecompHandle_t& handle) { // Check CUDECOMP_ENABLE_NCCL_UBR (NCCL user buffer registration) - handle->nccl_enable_ubr = checkEnvVar("CUDECOMP_ENABLE_NCCL_UBR"); + bool nccl_ubr_requested = checkEnvVar("CUDECOMP_ENABLE_NCCL_UBR"); // Check CUDECOMP_ENABLE_CUMEM (CUDA VMM allocations for work buffers) - handle->cuda_cumem_enable = checkEnvVar("CUDECOMP_ENABLE_CUMEM"); + bool cuda_cumem_requested = checkEnvVar("CUDECOMP_ENABLE_CUMEM"); + handle->cuda_cumem_enable = cuda_cumem_requested || nccl_ubr_requested; if (handle->cuda_cumem_enable) { -#if CUDART_VERSION < 12030 +#if CUDART_VERSION < 11030 if (handle->rank == 0) { - printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but CUDA version used for compilation does not " - "support fabric allocations. Disabling this feature.\n"); + if (cuda_cumem_requested) { + printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but CUDA version used for compilation does not " + "support CUDA VMM allocations. Disabling this feature.\n"); + } + if (nccl_ubr_requested) { + printf("CUDECOMP:WARN: CUDECOMP_ENABLE_NCCL_UBR is set but CUDA version used for compilation does not " + "support CUDA VMM allocations. Disabling this feature.\n"); + } } handle->cuda_cumem_enable = false; #else int driverVersion; CHECK_CUDA(cudaDriverGetVersion(&driverVersion)); - if (driverVersion < 12030) { + if (driverVersion < 11030) { if (handle->rank == 0) { - printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but installed driver does not " - "support fabric allocations. Disabling this feature.\n"); + if (cuda_cumem_requested) { + printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but installed driver does not " + "support CUDA VMM allocations. Disabling this feature.\n"); + } + if (nccl_ubr_requested) { + printf("CUDECOMP:WARN: CUDECOMP_ENABLE_NCCL_UBR is set but installed driver does not " + "support CUDA VMM allocations. Disabling this feature.\n"); + } } handle->cuda_cumem_enable = false; } else { - // Check if fabric allocation type is supported int dev; CUdevice cu_dev; CHECK_CUDA(cudaGetDevice(&dev)); CHECK_CUDA_DRV(cuDeviceGet(&cu_dev, dev)); - int flag = 0; - CHECK_CUDA_DRV(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_dev)); - if (!flag) { - if (handle->rank == 0) { - printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but device does not " - "support fabric allocations. Disabling this feature.\n"); + + int vmm_supported = 0; + int posix_fd_supported = 0; + CHECK_CUDA_DRV( + cuDeviceGetAttribute(&vmm_supported, CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, cu_dev)); + CHECK_CUDA_DRV(cuDeviceGetAttribute(&posix_fd_supported, + CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR_SUPPORTED, cu_dev)); + handle->cuda_cumem_enable = vmm_supported && posix_fd_supported; + if (!handle->cuda_cumem_enable && handle->rank == 0) { + if (cuda_cumem_requested) { + printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but the current device does not support CUDA VMM " + "allocations with POSIX file-descriptor handles. Disabling this feature.\n"); + } + if (nccl_ubr_requested) { + printf("CUDECOMP:WARN: CUDECOMP_ENABLE_NCCL_UBR is set but the current device does not support CUDA VMM " + "allocations with POSIX file-descriptor handles. Disabling this feature.\n"); } - handle->cuda_cumem_enable = false; } } #endif } + handle->nccl_enable_ubr = nccl_ubr_requested && handle->cuda_cumem_enable; // Check CUDECOMP_ENABLE_CUDA_GRAPHS (CUDA Graphs usage in pipelined backends) handle->cuda_graphs_enable = checkEnvVar("CUDECOMP_ENABLE_CUDA_GRAPHS"); @@ -601,7 +623,7 @@ static void cleanupFailedGridDescCreate(cudecompHandle_t handle, cudecompGridDes releaseUnusedHandleResources(handle, release_streams); } -#if CUDART_VERSION >= 12030 +#if CUDART_VERSION >= 11030 struct cuMemAllocationGuard { ~cuMemAllocationGuard() noexcept { if (!active) return; @@ -1268,16 +1290,28 @@ cudecompResult_t cudecompMalloc(cudecompHandle_t handle, cudecompGridDesc_t grid #endif } else { if (handle->cuda_cumem_enable) { -#if CUDART_VERSION >= 12030 +#if CUDART_VERSION >= 11030 int dev; CUdevice cu_dev; CHECK_CUDA(cudaGetDevice(&dev)); CHECK_CUDA_DRV(cuDeviceGet(&cu_dev, dev)); + int requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; +#if CUDART_VERSION >= 12030 + int driverVersion; + CHECK_CUDA(cudaDriverGetVersion(&driverVersion)); + if (driverVersion >= 12030) { + int fabric_supported = 0; + CHECK_CUDA_DRV( + cuDeviceGetAttribute(&fabric_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_dev)); + if (fabric_supported) requestedHandleTypes |= CU_MEM_HANDLE_TYPE_FABRIC; + } +#endif + CUmemAllocationProp prop = {}; prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + prop.requestedHandleTypes = static_cast(requestedHandleTypes); prop.location.id = cu_dev; // Check for RDMA support @@ -1286,15 +1320,35 @@ cudecompResult_t cudecompMalloc(cudecompHandle_t handle, cudecompGridDesc_t grid cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, cu_dev)); if (flag) prop.allocFlags.gpuDirectRDMACapable = 1; - // Align allocation size to required granularity + // Keep the caller-requested size so any retry can realign from the original value. + size_t original_buffer_size_bytes = buffer_size_bytes; size_t granularity; - CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM)); - buffer_size_bytes = (buffer_size_bytes + granularity - 1) / granularity * granularity; + CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity; // Allocate memory cuMemAllocationGuard cumem_guard; cumem_guard.size = buffer_size_bytes; - CHECK_CUDA_DRV(cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0)); + CUresult err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0); +#if CUDART_VERSION >= 12030 + if ((requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) && + (err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED)) { + // Fabric handles are useful when the platform supports them, but regular NCCL user buffer registration only + // requires POSIX FD export support. If Fabric creation is unavailable at runtime, keep VMM enabled and fall + // back to a POSIX-FD-only allocation. + requestedHandleTypes &= ~CU_MEM_HANDLE_TYPE_FABRIC; + prop.requestedHandleTypes = static_cast(requestedHandleTypes); + CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity; + cumem_guard.size = buffer_size_bytes; + err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0); + } +#endif + if (CUDA_SUCCESS != err) { + const char* error_str; + cuFnTable.pfn_cuGetErrorString(err, &error_str); + throw cudecomp::CudaError(__FILE__, __LINE__, error_str); + } cumem_guard.handle_created = true; CHECK_CUDA_DRV(cuMemAddressReserve(&cumem_guard.ptr, buffer_size_bytes, granularity, 0, 0)); cumem_guard.address_reserved = true; @@ -1382,7 +1436,7 @@ cudecompResult_t cudecompFree(cudecompHandle_t handle, cudecompGridDesc_t grid_d } else { if (handle->cuda_cumem_enable) { -#if CUDART_VERSION >= 12030 +#if CUDART_VERSION >= 11030 if (buffer) { CUmemGenericAllocationHandle cumem_handle; CHECK_CUDA_DRV(cuMemRetainAllocationHandle(&cumem_handle, buffer));