diff --git a/heat/core/_config.py b/heat/core/_config.py index c82063449e..017bcc8fb6 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -2,6 +2,7 @@ Everything you need to know about the configuration of Heat """ +from mpi4py import MPI import torch import platform import mpi4py @@ -27,28 +28,55 @@ class MPILibrary(Enum): class MPILibraryInfo: name: MPILibrary version: str + cuda_compatible: bool = False + rocm_compatible: bool = False + gpu_compatible: bool = False + incompatible_operations: list[str] | None = None + + +# Helper function to match version patterns +def _match_version_pattern( + version: str, patterns: dict[str, dict[str, list[str] | None]] +) -> dict[str, list[str] | None]: + """ + Match a version string against pattern keys (e.g., '5.0.x', '4.1.x', '*'). + Returns the incompatibilities dict for the matching pattern, or {} if no match. + + Parameters + ---------- + version : str + The version string to match (e.g., 'v5.0.1', '4.1.2') + patterns : dict[str, dict[str, list[str] | None]] + Dictionary mapping version patterns to incompatibilities + + Returns + ------- + dict[str, list[str] | None] + The incompatibilities for the matched version pattern, or {} if no match + """ + # First check for wildcard pattern + if "*" in patterns: + return patterns["*"] + + # Then try to match specific version patterns + for pattern, incompatibilities in patterns.items(): + # Convert pattern like '5.0.x' to regex '5\.0\.\d+' + regex_pattern = pattern.replace(".", r"\.").replace("x", r"\d+") + if re.match(f"^{regex_pattern}$", version): + return incompatibilities + + return {} def _get_mpi_library() -> MPILibraryInfo: - library = mpi4py.MPI.Get_library_version().split() - match library: - case ["Open", "MPI", *_]: - return MPILibraryInfo(MPILibrary.OpenMPI, library[2]) - case ["Intel(R)", "MPI", *_]: - return MPILibraryInfo(MPILibrary.IntelMPI, library[3]) - case ["MPICH", "Version:", *_]: - return MPILibraryInfo(MPILibrary.MPICH, library[2]) - case ["MVAPICH", "Version:", *_]: - return MPILibraryInfo(MPILibrary.MVAPICH, library[2]) - case ["===", "ParaStation", "MPI", *_]: - return MPILibraryInfo(MPILibrary.ParaStationMPI, library[3]) - case _: - return MPILibraryInfo(MPILibrary.Other, "unknown") + library_info = mpi4py.MPI.Get_library_version().split() + incompatibilities_list_id = "rocm" if CUDA_IS_ACTUALLY_ROCM else "cuda" + match library_info: + case ["Open", "MPI", *_]: + library = MPILibrary.OpenMPI + version = library_info[2] -def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: - match library.name: - case MPILibrary.OpenMPI: try: parsable_ompi_info = subprocess.check_output( ["ompi_info", "--parsable", "--all"] @@ -61,36 +89,144 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: # Check for extensions match = re.search(r"MPI extensions: (.*)", ompi_info) extensions = [ext.strip() for ext in match.group(0).split(":")[1].split(",")] - cuda = cuda_support_flag and "cuda" in extensions - if library.version.startswith("v4."): - rocm = cuda - elif library.version.startswith("v5."): - rocm = "rocm" in extensions or "hip" in extensions - # Seems to be broken, disabled by default for now - # return cuda, rocm - return False, False - except Exception as e: # noqa E722 - return False, False - case MPILibrary.IntelMPI: - return False, False - case MPILibrary.MVAPICH: - cuda = os.environ.get("MV2_USE_CUDA") == "1" - rocm = os.environ.get("MV2_USE_ROCM") == "1" - return cuda, rocm - case MPILibrary.MPICH: - cuda = os.environ.get("MPIR_CVAR_ENABLE_HCOLL") == "1" - rocm = False - return cuda, rocm - case MPILibrary.CrayMPI: - cuda = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" - rocm = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" - return cuda, rocm - case MPILibrary.ParaStationMPI: - cuda = os.environ.get("PSP_CUDA") == "1" - rocm = False - return cuda, rocm + cuda_is_compatible: bool = cuda_support_flag and "cuda" in extensions + if version.startswith("v4."): + rocm_is_compatible: bool = cuda_is_compatible + elif version.startswith("v5."): + rocm_is_compatible: bool = "rocm" in extensions or "hip" in extensions + + finally: + cuda_is_compatible = False + rocm_is_compatible = False + device_incompatibilities = None + + case ["Intel(R)", "MPI", *_]: + library = MPILibrary.IntelMPI + version = library_info[3] + + cuda_is_compatible = False + rocm_is_compatible = False + + case ["MPICH", "Version:", *_]: + library = MPILibrary.MPICH + version = library_info[2] + + cuda_is_compatible = os.environ.get("MV2_USE_CUDA", "0") == "1" + rocm_is_compatible = os.environ.get("MV2_USE_ROCM", "0") == "1" + + case ["MVAPICH", "Version:", *_]: + library = MPILibrary.MVAPICH + version = library_info[2] + + cuda_is_compatible = os.environ.get("MPIR_CVAR_ENABLE_HCOLL") == "1" + rocm_is_compatible = False + + case ["CrayMPI", *_]: + library = MPILibrary.CrayMPI + version = library_info[1] + incompatibilities = _match_version_pattern(version, INCOMPATIBILITIES.get(library, {})) + + cuda_is_compatible = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" + rocm_is_compatible = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" + + case ["===", "ParaStation", "MPI", *_]: + library = MPILibrary.ParaStationMPI + version = library_info[3] + cuda_is_compatible = os.environ.get("PSP_CUDA") == "1" + rocm_is_compatible = False + case _: - return False, False + library = MPILibrary.Other + version = "unknown" + cuda_is_compatible = False + rocm_is_compatible = False + + incompatibilities = _match_version_pattern(version, INCOMPATIBILITIES.get(library, {})) + + # Passes the incompatibilites of the combination library+device to device_incompatibilities. If the device is not found, it is set to False (non-compatible by default). + device_incompatibilities = ( + incompatibilities[incompatibilities_list_id] + if incompatibilities_list_id in incompatibilities + else False + ) + gpu_is_compatible = (rocm_is_compatible and CUDA_IS_ACTUALLY_ROCM) or ( + cuda_is_compatible and not CUDA_IS_ACTUALLY_ROCM + ) + gpu_is_compatible = gpu_is_compatible and isinstance(device_incompatibilities, list) + + return MPILibraryInfo( + library, + version, + cuda_is_compatible, + rocm_is_compatible, + gpu_is_compatible, + device_incompatibilities, + ) + + +# Library / version / device +# Structure: MPILibrary -> version_pattern -> device -> incompatibilities +# Incompatibilities can be: +# - False: All operations are incompatible for this device +# - [] (empty list): All operations are compatible for this device +# - [list of operation names]: Only the listed operations are incompatible +INCOMPATIBILITIES: dict[MPILibrary, dict[str, dict[str, list[str] | None]]] = { + MPILibrary.IntelMPI: {"*": {"cuda": False, "rocm": False}}, + MPILibrary.OpenMPI: { + "5.0.x": { + "cuda": [ + "Accumulate", + "Compare_and_swap", + "Fetch_and_op", + "Get_Accumulate", + "Iallgather", + "Iallgatherv", + "Iallreduce", + "Ialltoall", + "Ialltoallv", + "Ialltoallw", + "Ibcast", + "Iscan", + "Iexscan", + "Rget", + "Rput", + "Ireduce", + ], + "rocm": False, + }, + "4.1.x": { + "cuda": [], # All operations compatible + "rocm": [], # All operations compatible (ROCm handled same as CUDA in 4.1.x) + }, + }, + MPILibrary.MVAPICH: { + "*": { + "cuda": [], # All operations compatible when MV2_USE_CUDA=1 + "rocm": [], # All operations compatible when MV2_USE_ROCM=1 + } + }, + MPILibrary.MPICH: { + "*": { + "cuda": [], # All operations compatible when MPIR_CVAR_ENABLE_HCOLL=1 + "rocm": False, # ROCm not supported + } + }, + MPILibrary.CrayMPI: { + "*": { + "cuda": [], # All operations compatible when MPICH_GPU_SUPPORT_ENABLED=1 + "rocm": [], # All operations compatible when MPICH_GPU_SUPPORT_ENABLED=1 + } + }, + MPILibrary.ParaStationMPI: { + "*": { + "cuda": [], # All operations compatible when PSP_CUDA=1 + "rocm": False, + } + }, + MPILibrary.Other: { + "*": {"cuda": False, "rocm": False} + }, # Unknown library, assume incompatibility unless proven otherwise +} PLATFORM = platform.platform() @@ -99,11 +235,11 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: CUDA_IS_ACTUALLY_ROCM = "rocm" in TORCH_VERSION mpi_library = _get_mpi_library() -CUDA_AWARE_MPI, ROCM_AWARE_MPI = _check_gpu_aware_mpi(mpi_library) -GPU_AWARE_MPI = False +CUDA_AWARE_MPI, ROCM_AWARE_MPI = mpi_library.cuda_compatible, mpi_library.rocm_compatible +GPU_AWARE_MPI = mpi_library.gpu_compatible # warn the user if CUDA/ROCm-aware MPI is not available, but PyTorch can use GPUs with CUDA/ROCm -if TORCH_CUDA_IS_AVAILABLE: +if TORCH_CUDA_IS_AVAILABLE and not GPU_AWARE_MPI: if not CUDA_IS_ACTUALLY_ROCM and not CUDA_AWARE_MPI: warnings.warn( f"Heat has CUDA GPU-support (PyTorch version {TORCH_VERSION} and `torch.cuda.is_available() = True`), but CUDA-awareness of MPI could not be detected. This may lead to performance degradation as direct MPI-communication between GPUs is not possible.", @@ -115,5 +251,3 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: f"Heat has ROCm GPU-support (PyTorch version {TORCH_VERSION} and `torch.cuda.is_available() = True`), but ROCm-awareness of MPI could not be detected. This may lead to performance degradation as direct MPI-communication between GPUs is not possible.", UserWarning, ) - else: - GPU_AWARE_MPI = True diff --git a/heat/core/communication.py b/heat/core/communication.py index 91e689093c..a7dc0caf66 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -17,7 +17,7 @@ from .stride_tricks import sanitize_axis -from ._config import GPU_AWARE_MPI +from ._config import GPU_AWARE_MPI, mpi_library as MPI_LIBRARY class MPIRequest: @@ -59,11 +59,19 @@ def Wait(self, status: MPI.Status | None = None): if self.handle is None: return self.handle.Wait(status) - if self.tensor is not None and isinstance(self.tensor, torch.Tensor): - if self.permutation is not None: - self.recvbuf = self.recvbuf.permute(self.permutation) - if self.tensor is not None and self.tensor.is_cuda and not GPU_AWARE_MPI: - self.tensor.copy_(self.recvbuf) + + # Apply permutation if needed (for all buffer types) + if self.permutation is not None and self.recvbuf is not None: + self.recvbuf = self.recvbuf.permute(self.permutation) + + # Copy result from CPU back to GPU if needed + if self.tensor is not None: + tensor = self.tensor if isinstance(self.tensor, torch.Tensor) else self.tensor.larray + tensor_device = tensor.device + recvbuf_device = self.recvbuf.device + + if tensor_device != recvbuf_device: + tensor.copy_(self.recvbuf.to(tensor_device)) def __getattr__(self, name: str) -> Callable: """ @@ -400,8 +408,8 @@ def mpi_type_and_elements_of( # chain the types based on the for i in range(len(shape) - 1, -1, -1): mpi_type = mpi_type.Create_vector(shape[i], 1, strides[i]).Create_resized(0, offsets[i]) - mpi_type.Commit() + mpi_type.Commit() if counts is not None: return mpi_type, (counts, displs) @@ -462,9 +470,24 @@ def as_buffer( return [mpi_mem, elements, mpi_type] def _moveToCompDevice(self, x: torch.Tensor, func: Callable | None) -> torch.Tensor: - """Moves the torch tensor to the relevant device, in case the function is not compatible with the MPI+GPU library.""" + """ + Moves the torch tensor to the relevant device, in case the function is not compatible with the MPI+GPU library. + If communication happens on GPU, the stream is synchronized in order to prepare for communication. + + Parameters + ---------- + x: torch.Tensor + The tensor to be moved to the relevant device + func: Callable + The MPI function that is intended to be called with the tensor, used to check for compatibility with the MPI+GPU library + + Returns + ------- + torch.Tensor + The tensor on the relevant device for the MPI function + """ if x.is_cuda: - if GPU_AWARE_MPI: + if GPU_AWARE_MPI and func.__name__ not in MPI_LIBRARY.incompatible_operations: torch.cuda.synchronize(x.device) return x else: @@ -861,8 +884,11 @@ def Bcast(self, buf: Any, root: int = 0) -> None: Rank of the root process, that broadcasts the message """ ret, sbuf, rbuf, buf = self.__broadcast_like(self.handle.Bcast, buf, root) - if buf is not None and isinstance(buf, torch.Tensor) and buf.is_cuda and not GPU_AWARE_MPI: - buf.copy_(rbuf) + if buf is not None and not GPU_AWARE_MPI: + if isinstance(buf, torch.Tensor) and buf.is_cuda: + buf.copy_(rbuf) + elif isinstance(buf, DNDarray) and buf.larray.is_cuda: + buf.larray.copy_(rbuf) return ret Bcast.__doc__ = MPI.Comm.Bcast.__doc__ @@ -1105,8 +1131,11 @@ def Allreduce( The operation to perform upon reduction """ ret, sbuf, rbuf, buf = self.__reduce_like(self.handle.Allreduce, sendbuf, recvbuf, op) - if buf is not None and isinstance(buf, torch.Tensor) and buf.is_cuda and not GPU_AWARE_MPI: - buf.copy_(rbuf) + if buf is not None and not GPU_AWARE_MPI: + if isinstance(buf, torch.Tensor) and buf.is_cuda: + buf.copy_(rbuf) + elif isinstance(buf, DNDarray) and buf.larray.is_cuda: + buf.larray.copy_(rbuf) return ret Allreduce.__doc__ = MPI.Comm.Allreduce.__doc__ @@ -1373,7 +1402,7 @@ def __allgather_like( rbuf = recvbuf mpi_recvbuf = recvbuf - # perform the scatter operation + # perform the allgather operation exit_code = func(mpi_sendbuf, mpi_recvbuf, **kwargs) return exit_code, sbuf, rbuf, original_recvbuf, recv_axis_permutation @@ -1401,8 +1430,12 @@ def Allgather( ) if buf is not None and isinstance(buf, torch.Tensor) and permutation is not None: rbuf = rbuf.permute(permutation) - if isinstance(buf, torch.Tensor) and buf.is_cuda and not GPU_AWARE_MPI: - buf.copy_(rbuf) + + if buf is not None and not GPU_AWARE_MPI: + if isinstance(buf, torch.Tensor) and buf.is_cuda: + buf.copy_(rbuf) + elif isinstance(buf, DNDarray) and buf.larray.is_cuda: + buf.larray.copy_(rbuf) return ret Allgather.__doc__ = MPI.Comm.Allgather.__doc__ @@ -1430,8 +1463,11 @@ def Allgatherv( ) if buf is not None and isinstance(buf, torch.Tensor) and permutation is not None: rbuf = rbuf.permute(permutation) - if isinstance(buf, torch.Tensor) and buf.is_cuda and not GPU_AWARE_MPI: - buf.copy_(rbuf) + if buf is not None and not GPU_AWARE_MPI: + if isinstance(buf, torch.Tensor) and buf.is_cuda: + buf.copy_(rbuf) + elif isinstance(buf, DNDarray) and buf.larray.is_cuda: + buf.larray.copy_(rbuf) return ret Allgatherv.__doc__ = MPI.Comm.Allgatherv.__doc__ @@ -1846,7 +1882,6 @@ def _create_recursive_vectortype( ... datatype, tensor_stride, subarray_sizes ... ) """ - datatype_history = [] current_datatype = datatype i = len(tensor_stride) - 1 @@ -1866,25 +1901,21 @@ def _create_recursive_vectortype( next_size = subarray_sizes[i] new_vector_datatype = current_datatype.Create_vector( next_size, current_size, current_stride - ).Commit() + ) else: if i == len(tensor_stride) - 1: new_vector_datatype = current_datatype.Create_vector( current_size, 1, current_stride - ).Commit() + ) else: - new_vector_datatype = current_datatype.Create_vector( - current_size, 1, 1 - ).Commit() + new_vector_datatype = current_datatype.Create_vector(current_size, 1, 1) - datatype_history.append(new_vector_datatype) # Set extent of the new datatype to the extent of the basic datatype to allow interweaving of data next_stride = tensor_stride[i - 1] new_resized_vector_datatype = new_vector_datatype.Create_resized( 0, datatype.Get_extent()[1] * next_stride - ).Commit() - datatype_history.append(new_resized_vector_datatype) + ) current_datatype = new_resized_vector_datatype i -= 1 @@ -1892,8 +1923,6 @@ def _create_recursive_vectortype( displacement = sum([x * y for x, y in zip(tensor_stride, start)]) * datatype.Get_extent()[1] current_datatype = current_datatype.Create_hindexed_block(1, [displacement]).Commit() - for dt in datatype_history[:-1]: - dt.Free() return current_datatype def Ialltoall( diff --git a/tests/core/test_communication.py b/tests/core/test_communication.py index 716865cb5f..2dc8c874ff 100644 --- a/tests/core/test_communication.py +++ b/tests/core/test_communication.py @@ -2621,11 +2621,8 @@ def test_largecount_workaround_IsendRecv(self): ) def test_largecount_workaround_Allreduce(self): shape = (2**10, 2**11, 2**10) - data = ( - torch.zeros(shape, dtype=torch.bool) - if ht.MPI_WORLD.rank % 2 == 0 - else torch.ones(shape, dtype=torch.bool) - ) + data = torch.zeros(shape, dtype=torch.bool) if ht.MPI_WORLD.rank % 2 == 0 else torch.ones(shape, dtype=torch.bool) + ht.MPI_WORLD.Allreduce(ht.MPI.IN_PLACE, data, op=ht.MPI.SUM) self.assertTrue(data.all())