From 67cd7c9cdf234b7c425b4a090ac095c567549f48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guti=C3=A9rrez=20Hermosillo=20Muriedas=2C=20Juan=20Pedro?= Date: Wed, 6 May 2026 17:07:11 +0200 Subject: [PATCH 1/9] wip: added incompatibility list --- heat/core/_config.py | 47 ++++++++++++++++++++++++---- heat/core/communication.py | 53 ++++++++++++++++++++++++++------ tests/core/test_communication.py | 7 ++--- 3 files changed, 86 insertions(+), 21 deletions(-) diff --git a/heat/core/_config.py b/heat/core/_config.py index c82063449e..7c20fed5e0 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -27,21 +27,27 @@ class MPILibrary(Enum): class MPILibraryInfo: name: MPILibrary version: str + incompatible_operations: dict[str, list[str]] = dataclasses.field(default_factory=dict) def _get_mpi_library() -> MPILibraryInfo: library = mpi4py.MPI.Get_library_version().split() match library: case ["Open", "MPI", *_]: - return MPILibraryInfo(MPILibrary.OpenMPI, library[2]) + version = library[2] + if version.startswith("v5.0."): + incompatibilities = INCOMPATIBILITIES[MPILibrary.OpenMPI].get("5.0.x", {}) + elif version.startswith("v4.1."): + incompatibilities = INCOMPATIBILITIES[MPILibrary.OpenMPI].get("4.1.x", {}) + return MPILibraryInfo(MPILibrary.OpenMPI, library[2], incompatibilities) case ["Intel(R)", "MPI", *_]: - return MPILibraryInfo(MPILibrary.IntelMPI, library[3]) + return MPILibraryInfo(MPILibrary.IntelMPI, library[3], {}) case ["MPICH", "Version:", *_]: - return MPILibraryInfo(MPILibrary.MPICH, library[2]) + return MPILibraryInfo(MPILibrary.MPICH, library[2], {}) case ["MVAPICH", "Version:", *_]: - return MPILibraryInfo(MPILibrary.MVAPICH, library[2]) + return MPILibraryInfo(MPILibrary.MVAPICH, library[2], {}) case ["===", "ParaStation", "MPI", *_]: - return MPILibraryInfo(MPILibrary.ParaStationMPI, library[3]) + return MPILibraryInfo(MPILibrary.ParaStationMPI, library[3], {}) case _: return MPILibraryInfo(MPILibrary.Other, "unknown") @@ -68,7 +74,7 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: rocm = "rocm" in extensions or "hip" in extensions # Seems to be broken, disabled by default for now # return cuda, rocm - return False, False + return cuda, rocm except Exception as e: # noqa E722 return False, False case MPILibrary.IntelMPI: @@ -93,6 +99,35 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: return False, False +# Library / version / device +INCOMPATIBILITIES: dict[MPILibrary, dict[str, dict[str, list[str]]]] = { + MPILibrary.IntelMPI: {}, + 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", + ] + }, + "4.1.x": {"cuda": []}, + }, +} + + PLATFORM = platform.platform() TORCH_VERSION = torch.__version__ TORCH_CUDA_IS_AVAILABLE = torch.cuda.is_available() diff --git a/heat/core/communication.py b/heat/core/communication.py index 4df8ac7cab..9f19b30c04 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -15,7 +15,7 @@ from .stride_tricks import sanitize_axis -from ._config import GPU_AWARE_MPI +from ._config import GPU_AWARE_MPI, mpi_library class MPIRequest: @@ -55,11 +55,25 @@ def Wait(self, status: MPI.Status = None): Waits for an MPI request to complete """ 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_device = ( + self.tensor.device + if isinstance(self.tensor, torch.Tensor) + else self.tensor.larray.device + ) + recvbuf_device = self.recvbuf.device + + if tensor_device != recvbuf_device: + if isinstance(self.tensor, torch.Tensor): + self.tensor.copy_(self.recvbuf.to(tensor_device)) + else: + self.tensor.larray.copy_(self.recvbuf.to(tensor_device)) def __getattr__(self, name: str) -> Callable: """ @@ -434,9 +448,25 @@ 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. + + 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.get( + "cuda", [] + ): torch.cuda.synchronize(x.device) return x else: @@ -1074,8 +1104,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__ diff --git a/tests/core/test_communication.py b/tests/core/test_communication.py index d64a0ec787..b92273623a 100644 --- a/tests/core/test_communication.py +++ b/tests/core/test_communication.py @@ -2616,11 +2616,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()) From 9ccb3680fc7aaba8ffe4b4f27a190833db1dcae6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guti=C3=A9rrez=20Hermosillo=20Muriedas=2C=20Juan=20Pedro?= Date: Mon, 1 Jun 2026 17:45:18 +0200 Subject: [PATCH 2/9] fix: tests --- heat/core/communication.py | 42 +++++++++++++++++++++----------------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/heat/core/communication.py b/heat/core/communication.py index fcd89d3f1e..a456145b17 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -396,8 +396,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) @@ -877,8 +877,11 @@ def Bcast(self, buf: Union[DNDarray, torch.Tensor, 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__ @@ -1374,7 +1377,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 @@ -1402,8 +1405,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__ @@ -1431,8 +1438,12 @@ 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) + print("Unpermuted") + 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__ @@ -1829,7 +1840,6 @@ def _create_recursive_vectortype( ... datatype, tensor_stride, subarray_sizes ... ) """ - datatype_history = [] current_datatype = datatype i = len(tensor_stride) - 1 @@ -1849,25 +1859,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 @@ -1875,8 +1881,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( From 8b1eb535b710926bfbe7c476dfd551d9c62c2f5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20Pedro=20Guti=C3=A9rrez=20Hermosillo=20Muriedas?= Date: Tue, 16 Jun 2026 09:17:15 +0200 Subject: [PATCH 3/9] Update heat/core/_config.py Co-authored-by: Thomas Saupe <39156931+brownbaerchen@users.noreply.github.com> --- heat/core/_config.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/heat/core/_config.py b/heat/core/_config.py index 7c20fed5e0..c0d114bdf4 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -72,8 +72,6 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: 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 cuda, rocm except Exception as e: # noqa E722 return False, False From dec2138e4dc0fd556061ecdff66767f3559343ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20Pedro=20Guti=C3=A9rrez=20Hermosillo=20Muriedas?= Date: Tue, 16 Jun 2026 09:19:23 +0200 Subject: [PATCH 4/9] Update heat/core/communication.py Co-authored-by: Thomas Saupe <39156931+brownbaerchen@users.noreply.github.com> --- heat/core/communication.py | 1 + 1 file changed, 1 insertion(+) diff --git a/heat/core/communication.py b/heat/core/communication.py index a456145b17..3cb43affe3 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -460,6 +460,7 @@ def as_buffer( 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. + If communication happens on GPU, the stream is synchronized in order to prepare for communication. Parameters ---------- From da9dc6134e8abfd2f4aa296738562ac19eb15f2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20Pedro=20Guti=C3=A9rrez=20Hermosillo=20Muriedas?= Date: Tue, 16 Jun 2026 09:19:46 +0200 Subject: [PATCH 5/9] Update heat/core/communication.py Co-authored-by: Thomas Saupe <39156931+brownbaerchen@users.noreply.github.com> --- heat/core/communication.py | 1 - 1 file changed, 1 deletion(-) diff --git a/heat/core/communication.py b/heat/core/communication.py index 3cb43affe3..0013d30ce2 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -1439,7 +1439,6 @@ def Allgatherv( ) if buf is not None and isinstance(buf, torch.Tensor) and permutation is not None: rbuf = rbuf.permute(permutation) - print("Unpermuted") if buf is not None and not GPU_AWARE_MPI: if isinstance(buf, torch.Tensor) and buf.is_cuda: buf.copy_(rbuf) From 6a11f7a571c5fd5e9e2a0bed7892dac74df170af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guti=C3=A9rrez=20Hermosillo=20Muriedas=2C=20Juan=20Pedro?= Date: Fri, 19 Jun 2026 10:38:19 +0200 Subject: [PATCH 6/9] fix: refactored _conf --- heat/core/_config.py | 192 +++++++++++++++++++++++++++---------- heat/core/communication.py | 6 +- 2 files changed, 144 insertions(+), 54 deletions(-) diff --git a/heat/core/_config.py b/heat/core/_config.py index c0d114bdf4..64137ee083 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -2,6 +2,8 @@ Everything you need to know about the configuration of Heat """ +from mpi4py import MPI +from numpy import isin import torch import platform import mpi4py @@ -27,34 +29,55 @@ class MPILibrary(Enum): class MPILibraryInfo: name: MPILibrary version: str - incompatible_operations: dict[str, list[str]] = dataclasses.field(default_factory=dict) + 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", *_]: - version = library[2] - if version.startswith("v5.0."): - incompatibilities = INCOMPATIBILITIES[MPILibrary.OpenMPI].get("5.0.x", {}) - elif version.startswith("v4.1."): - incompatibilities = INCOMPATIBILITIES[MPILibrary.OpenMPI].get("4.1.x", {}) - return MPILibraryInfo(MPILibrary.OpenMPI, library[2], incompatibilities) - 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"] @@ -68,38 +91,78 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: 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."): + if version.startswith("v4."): rocm = cuda - elif library.version.startswith("v5."): + elif version.startswith("v5."): rocm = "rocm" in extensions or "hip" in extensions - return cuda, rocm - 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: + + finally: + cuda = False + rocm = False + gpu_comp = False + device_incompatibilities = None + + case ["Intel(R)", "MPI", *_]: + library = MPILibrary.IntelMPI + version = library_info[3] + + cuda = False + rocm = False + + case ["MPICH", "Version:", *_]: + library = MPILibrary.MPICH + version = library_info[2] + + cuda = os.environ.get("MV2_USE_CUDA", "0") == "1" + rocm = os.environ.get("MV2_USE_ROCM", "0") == "1" + + case ["MVAPICH", "Version:", *_]: + library = MPILibrary.MVAPICH + version = library_info[2] + cuda = os.environ.get("MPIR_CVAR_ENABLE_HCOLL") == "1" rocm = False - return cuda, rocm - case MPILibrary.CrayMPI: + + case ["CrayMPI", *_]: + library = MPILibrary.CrayMPI + version = library_info[1] + incompatibilities = _match_version_pattern(version, INCOMPATIBILITIES.get(library, {})) + cuda = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" rocm = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" - return cuda, rocm - case MPILibrary.ParaStationMPI: + + case ["===", "ParaStation", "MPI", *_]: + library = MPILibrary.ParaStationMPI + version = library_info[3] cuda = os.environ.get("PSP_CUDA") == "1" rocm = False - return cuda, rocm + case _: - return False, False + library = MPILibrary.Other + version = "unknown" + cuda = False + rocm = False + + incompatibilities = _match_version_pattern(version, INCOMPATIBILITIES.get(library, {})) + device_incompatibilities = ( + incompatibilities[incompatibilities_list_id] + if incompatibilities_list_id in incompatibilities + else None + ) + gpu_comp = (rocm and CUDA_IS_ACTUALLY_ROCM) or (cuda and not CUDA_IS_ACTUALLY_ROCM) + gpu_comp = gpu_comp and isinstance(device_incompatibilities, list) + + return MPILibraryInfo(library, version, cuda, rocm, gpu_comp, device_incompatibilities) # Library / version / device -INCOMPATIBILITIES: dict[MPILibrary, dict[str, dict[str, list[str]]]] = { - MPILibrary.IntelMPI: {}, +# Structure: MPILibrary -> version_pattern -> device -> incompatibilities +# Incompatibilities can be: +# - None: 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": None, "rocm": None}}, MPILibrary.OpenMPI: { "5.0.x": { "cuda": [ @@ -119,10 +182,41 @@ def _check_gpu_aware_mpi(library: MPILibraryInfo) -> tuple[bool, bool]: "Rget", "Rput", "Ireduce", - ] + ], + "rocm": None, }, - "4.1.x": {"cuda": []}, + "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": None, # 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": None, # ROCm not supported + } }, + MPILibrary.Other: { + "*": {"cuda": None, "rocm": None} + }, # Unknown library, assume compatibility unless proven otherwise } @@ -132,11 +226,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.", @@ -148,5 +242,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 0013d30ce2..88865a4cc7 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -15,7 +15,7 @@ from .stride_tricks import sanitize_axis -from ._config import GPU_AWARE_MPI, mpi_library +from ._config import GPU_AWARE_MPI, mpi_library as MPI_LIBRARY class MPIRequest: @@ -475,9 +475,7 @@ def _moveToCompDevice(self, x: torch.Tensor, func: Callable | None) -> torch.Ten The tensor on the relevant device for the MPI function """ if x.is_cuda: - if GPU_AWARE_MPI and func.__name__ not in mpi_library.incompatible_operations.get( - "cuda", [] - ): + if GPU_AWARE_MPI and func.__name__ not in MPI_LIBRARY.incompatible_operations: torch.cuda.synchronize(x.device) return x else: From 3493912e488530884bf14b955644710a675b6186 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guti=C3=A9rrez=20Hermosillo=20Muriedas=2C=20Juan=20Pedro?= Date: Tue, 23 Jun 2026 14:03:51 +0200 Subject: [PATCH 7/9] fix: review comments --- heat/core/_config.py | 71 ++++++++++++++++++++++---------------- heat/core/communication.py | 12 ++----- 2 files changed, 44 insertions(+), 39 deletions(-) diff --git a/heat/core/_config.py b/heat/core/_config.py index 64137ee083..26316b4faa 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -3,7 +3,6 @@ """ from mpi4py import MPI -from numpy import isin import torch import platform import mpi4py @@ -14,6 +13,8 @@ import dataclasses from enum import Enum +from torch._C import _rocm_is_backward_pass + class MPILibrary(Enum): OpenMPI = "ompi" @@ -90,75 +91,85 @@ def _get_mpi_library() -> MPILibraryInfo: # 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 + cuda_is_compatible: bool = cuda_support_flag and "cuda" in extensions if version.startswith("v4."): - rocm = cuda + rocm_is_compatible: bool = cuda_is_compatible elif version.startswith("v5."): - rocm = "rocm" in extensions or "hip" in extensions + rocm_is_compatible: bool = "rocm" in extensions or "hip" in extensions finally: - cuda = False - rocm = False - gpu_comp = False + cuda_is_compatible = False + rocm_is_compatible = False device_incompatibilities = None case ["Intel(R)", "MPI", *_]: library = MPILibrary.IntelMPI version = library_info[3] - cuda = False - rocm = False + cuda_is_compatible = False + rocm_is_compatible = False case ["MPICH", "Version:", *_]: library = MPILibrary.MPICH version = library_info[2] - cuda = os.environ.get("MV2_USE_CUDA", "0") == "1" - rocm = os.environ.get("MV2_USE_ROCM", "0") == "1" + 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 = os.environ.get("MPIR_CVAR_ENABLE_HCOLL") == "1" - rocm = False + 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 = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" - rocm = os.environ.get("MPICH_GPU_SUPPORT_ENABLED") == "1" + 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 = os.environ.get("PSP_CUDA") == "1" - rocm = False + cuda_is_compatible = os.environ.get("PSP_CUDA") == "1" + rocm_is_compatible = False case _: library = MPILibrary.Other version = "unknown" - cuda = False - rocm = False + 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 None + 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, ) - gpu_comp = (rocm and CUDA_IS_ACTUALLY_ROCM) or (cuda and not CUDA_IS_ACTUALLY_ROCM) - gpu_comp = gpu_comp and isinstance(device_incompatibilities, list) - - return MPILibraryInfo(library, version, cuda, rocm, gpu_comp, device_incompatibilities) # Library / version / device # Structure: MPILibrary -> version_pattern -> device -> incompatibilities # Incompatibilities can be: -# - None: All operations are incompatible for this device +# - 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]]] = { @@ -183,7 +194,7 @@ def _get_mpi_library() -> MPILibraryInfo: "Rput", "Ireduce", ], - "rocm": None, + "rocm": False, }, "4.1.x": { "cuda": [], # All operations compatible @@ -199,7 +210,7 @@ def _get_mpi_library() -> MPILibraryInfo: MPILibrary.MPICH: { "*": { "cuda": [], # All operations compatible when MPIR_CVAR_ENABLE_HCOLL=1 - "rocm": None, # ROCm not supported + "rocm": False, # ROCm not supported } }, MPILibrary.CrayMPI: { @@ -211,12 +222,12 @@ def _get_mpi_library() -> MPILibraryInfo: MPILibrary.ParaStationMPI: { "*": { "cuda": [], # All operations compatible when PSP_CUDA=1 - "rocm": None, # ROCm not supported + "rocm": False, } }, MPILibrary.Other: { - "*": {"cuda": None, "rocm": None} - }, # Unknown library, assume compatibility unless proven otherwise + "*": {"cuda": False, "rocm": False} + }, # Unknown library, assume incompatibility unless proven otherwise } diff --git a/heat/core/communication.py b/heat/core/communication.py index 88865a4cc7..021576c7d2 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -62,18 +62,12 @@ def Wait(self, status: MPI.Status = None): # Copy result from CPU back to GPU if needed if self.tensor is not None: - tensor_device = ( - self.tensor.device - if isinstance(self.tensor, torch.Tensor) - else self.tensor.larray.device - ) + 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: - if isinstance(self.tensor, torch.Tensor): - self.tensor.copy_(self.recvbuf.to(tensor_device)) - else: - self.tensor.larray.copy_(self.recvbuf.to(tensor_device)) + tensor.copy_(self.recvbuf.to(tensor_device)) def __getattr__(self, name: str) -> Callable: """ From ab3e4d1fdd6ca4094c629dddc430b7ae27d4d091 Mon Sep 17 00:00:00 2001 From: Thomas Saupe <39156931+brownbaerchen@users.noreply.github.com> Date: Fri, 26 Jun 2026 12:30:32 +0200 Subject: [PATCH 8/9] Apply suggestion from @brownbaerchen --- heat/core/_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heat/core/_config.py b/heat/core/_config.py index 26316b4faa..012e8f12e8 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -173,7 +173,7 @@ def _get_mpi_library() -> MPILibraryInfo: # - [] (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": None, "rocm": None}}, + MPILibrary.IntelMPI: {"*": {"cuda": False, "rocm": False}}, MPILibrary.OpenMPI: { "5.0.x": { "cuda": [ From e1568dbdafc772d66b2e2059d5c681fb4ca7af4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20Pedro=20Guti=C3=A9rrez=20Hermosillo=20Muriedas?= Date: Mon, 29 Jun 2026 09:11:30 +0200 Subject: [PATCH 9/9] Update _config.py --- heat/core/_config.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/heat/core/_config.py b/heat/core/_config.py index 012e8f12e8..017bcc8fb6 100644 --- a/heat/core/_config.py +++ b/heat/core/_config.py @@ -13,8 +13,6 @@ import dataclasses from enum import Enum -from torch._C import _rocm_is_backward_pass - class MPILibrary(Enum): OpenMPI = "ompi"