From ee9de89bb6aab3824efd1d47c286682629fd2afa Mon Sep 17 00:00:00 2001 From: Claudia Comito <39374113+ClaudiaComito@users.noreply.github.com> Date: Tue, 21 Apr 2026 18:14:01 +0200 Subject: [PATCH 1/4] test: nccl backend for matmul --- heat/core/linalg/basics.py | 97 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/heat/core/linalg/basics.py b/heat/core/linalg/basics.py index c52e64c40e..16f1e8edc9 100644 --- a/heat/core/linalg/basics.py +++ b/heat/core/linalg/basics.py @@ -50,6 +50,80 @@ "vector_norm", ] +import torch.distributed as dist + +try: + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Shard, Replicate + from torch.distributed.tensor.placement_types import Partial + + _DTENSOR_AVAILABLE = True +except ImportError: + _DTENSOR_AVAILABLE = False + +_DTENSOR_MESH = None + + +def _get_or_create_mesh(): + import os + + global _DTENSOR_MESH + if _DTENSOR_MESH is None: + if not dist.is_initialized(): + rank = int(os.environ.get("OMPI_COMM_WORLD_RANK", "0")) + world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", "1")) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "127.0.0.1") + os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500") + + backend = "nccl" if torch.cuda.is_available() else "gloo" + dist.init_process_group(backend=backend) + + device_type = "cuda" if torch.cuda.is_available() else "cpu" + _DTENSOR_MESH = init_device_mesh(device_type, (dist.get_world_size(),)) + return _DTENSOR_MESH + + +def _use_dtensor(a: DNDarray, b: DNDarray) -> bool: + if not _DTENSOR_AVAILABLE: + return False + + # enforce nccl backend to guarantee performance over mpi ring topologies + if not torch.cuda.is_available(): + return False + + # ensure 2d matrices as dtensor matmul behavior differs on higher dims + if a.ndim != 2 or b.ndim != 2: + return False + + # avoid faketensor propagation crashes by enforcing strict divisibility + comm_size = a.comm.size + if a.split is not None and a.gshape[a.split] % comm_size != 0: + return False + if b.split is not None and b.gshape[b.split] % comm_size != 0: + return False + + return True + + +def _to_dtensor(dndarray: DNDarray, mesh) -> "DTensor": + placements = [Replicate()] if dndarray.split is None else [Shard(dndarray.split)] + return DTensor.from_local(dndarray.larray, mesh, placements) + + +def _from_dtensor(dtensor: "DTensor", target_split: int) -> torch.Tensor: + target_placement = Replicate() if target_split is None else Shard(target_split) + + # force network redistribution if tensor is incomplete or layout mismatches + if ( + any(isinstance(p, Partial) for p in dtensor.placements) + or dtensor.placements[0] != target_placement + ): + dtensor = dtensor.redistribute(dtensor.device_mesh, [target_placement]) + + return dtensor.to_local() + def _estimate_largest_singularvalue(A: DNDarray, algorithm: str = "fro") -> DNDarray: """ @@ -611,6 +685,29 @@ def matmul(a: DNDarray, b: DNDarray, allow_resplit: bool = False) -> DNDarray: sanitation.sanitize_in(a) sanitation.sanitize_in(b) + if _use_dtensor(a, b): + try: + mesh = _get_or_create_mesh() + + dt_a = _to_dtensor(a, mesh) + dt_b = _to_dtensor(b, mesh) + + dt_c = torch.matmul(dt_a, dt_b) + + # heat infers the final split directly from the inputs + expected_split = a.split if a.split is not None else b.split + local_c = _from_dtensor(dt_c, expected_split) + + gshape_c = (a.gshape[0], b.gshape[1]) + + return DNDarray( + local_c, gshape_c, a.dtype, expected_split, a.device, a.comm, balanced=True + ) + except Exception as e: + import logging + + logging.warning(f"dtensor routing failed, falling back to mpi: {e}") + batch_dim = max(a.ndim, b.ndim) - 2 # -1 for vector vector multiplication batched = batch_dim > 0 From cb1159e906e481281e78691dc811af5303bc5919 Mon Sep 17 00:00:00 2001 From: Claudia Comito <39374113+ClaudiaComito@users.noreply.github.com> Date: Wed, 22 Apr 2026 09:25:25 +0200 Subject: [PATCH 2/4] introduce torch.distr mesh creation --- heat/core/communication.py | 68 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/heat/core/communication.py b/heat/core/communication.py index 5ee875b5ad..985ff92d15 100644 --- a/heat/core/communication.py +++ b/heat/core/communication.py @@ -2467,5 +2467,73 @@ def use_comm(comm: Communication = None): __default_comm = sanitize_comm(comm) +# DTensor mesh initialization + +import torch.distributed as dist + +try: + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Shard, Replicate + from torch.distributed.tensor.placement_types import Partial + + _DTENSOR_AVAILABLE = True +except ImportError: + _DTENSOR_AVAILABLE = False + +_DTENSOR_MESHES = {} + + +def _get_or_create_mesh(device): + """ + Initializes a PyTorch Distributed ProcessGroup and DeviceMesh that + mirrors the underlying MPI communicator for the given device. + """ + import os + + global _DTENSOR_MESHES + mesh_device_type = "cuda" if str(device)[:3] == "gpu" else "cpu" + + if mesh_device_type not in _DTENSOR_MESHES: + if not dist.is_initialized(): + # Map MPI ranks to PyTorch Distributed environment variables + rank = int(os.environ.get("OMPI_COMM_WORLD_RANK", "0")) + world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", "1")) + + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + + # Map MPI local rank to physical GPU ID (only if using GPUs) + if mesh_device_type == "cuda" and torch.cuda.is_available(): + local_rank = int(os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + + # Network Configuration + if "MASTER_ADDR" not in os.environ: + import logging + + logging.info("MASTER_ADDR not found in environment. Defaulting to 127.0.0.1") + os.environ["MASTER_ADDR"] = "127.0.0.1" + + if "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = "6000" + + # Initialize Process Group + if mesh_device_type == "cuda": + backend = "nccl" + elif dist.is_mpi_available(): + backend = "mpi" + else: + backend = "gloo" + + dist.init_process_group(backend=backend) + + # Create Device Mesh for the given device type + _DTENSOR_MESHES[mesh_device_type] = init_device_mesh( + mesh_device_type, (dist.get_world_size(),) + ) + + return _DTENSOR_MESHES[mesh_device_type] + + # import at the end of file to break circular dependencies from .dndarray import DNDarray From f6c498adb4972f3c3053778aab3a1bb1488cfd51 Mon Sep 17 00:00:00 2001 From: Claudia Comito <39374113+ClaudiaComito@users.noreply.github.com> Date: Wed, 22 Apr 2026 09:27:37 +0200 Subject: [PATCH 3/4] torch.distr backend for even GPU matmul --- heat/core/linalg/basics.py | 63 ++++++++------------------------------ 1 file changed, 13 insertions(+), 50 deletions(-) diff --git a/heat/core/linalg/basics.py b/heat/core/linalg/basics.py index 16f1e8edc9..6bc0ea4de0 100644 --- a/heat/core/linalg/basics.py +++ b/heat/core/linalg/basics.py @@ -12,6 +12,11 @@ from torch._C import Value from ..communication import MPI +from ..communication import _get_or_create_mesh, _DTENSOR_AVAILABLE + +if _DTENSOR_AVAILABLE: + from torch.distributed.tensor import DTensor, Shard, Replicate + from torch.distributed.tensor.placement_types import Partial from .. import arithmetics from .. import complex_math from .. import constants @@ -50,59 +55,17 @@ "vector_norm", ] -import torch.distributed as dist - -try: - from torch.distributed.device_mesh import init_device_mesh - from torch.distributed.tensor import DTensor, Shard, Replicate - from torch.distributed.tensor.placement_types import Partial - - _DTENSOR_AVAILABLE = True -except ImportError: - _DTENSOR_AVAILABLE = False - -_DTENSOR_MESH = None - - -def _get_or_create_mesh(): - import os - global _DTENSOR_MESH - if _DTENSOR_MESH is None: - if not dist.is_initialized(): - rank = int(os.environ.get("OMPI_COMM_WORLD_RANK", "0")) - world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", "1")) - os.environ["RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "127.0.0.1") - os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500") - - backend = "nccl" if torch.cuda.is_available() else "gloo" - dist.init_process_group(backend=backend) - - device_type = "cuda" if torch.cuda.is_available() else "cpu" - _DTENSOR_MESH = init_device_mesh(device_type, (dist.get_world_size(),)) - return _DTENSOR_MESH - - -def _use_dtensor(a: DNDarray, b: DNDarray) -> bool: +def _use_dtensor(*DNDarrays) -> bool: if not _DTENSOR_AVAILABLE: return False - # enforce nccl backend to guarantee performance over mpi ring topologies - if not torch.cuda.is_available(): - return False - - # ensure 2d matrices as dtensor matmul behavior differs on higher dims - if a.ndim != 2 or b.ndim != 2: - return False - - # avoid faketensor propagation crashes by enforcing strict divisibility - comm_size = a.comm.size - if a.split is not None and a.gshape[a.split] % comm_size != 0: - return False - if b.split is not None and b.gshape[b.split] % comm_size != 0: - return False + # on evenly distributed DNDarrays and on GPU only + for array in DNDarrays: + if not array.is_distributed() or str(array.device)[:3] != "gpu": + return False + if array.split is not None and array.gshape[array.split] % array.comm.size != 0: + return False return True @@ -687,7 +650,7 @@ def matmul(a: DNDarray, b: DNDarray, allow_resplit: bool = False) -> DNDarray: if _use_dtensor(a, b): try: - mesh = _get_or_create_mesh() + mesh = _get_or_create_mesh(a.device) dt_a = _to_dtensor(a, mesh) dt_b = _to_dtensor(b, mesh) From 1f35d8342de3d7272a4bcad4f18fbac17223fec8 Mon Sep 17 00:00:00 2001 From: Claudia Comito <39374113+ClaudiaComito@users.noreply.github.com> Date: Wed, 22 Apr 2026 11:45:46 +0200 Subject: [PATCH 4/4] at least 1 distributed array for DTensor backend --- heat/core/linalg/basics.py | 36 +++++++++++++++++++----------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/heat/core/linalg/basics.py b/heat/core/linalg/basics.py index 6bc0ea4de0..8e27b1b105 100644 --- a/heat/core/linalg/basics.py +++ b/heat/core/linalg/basics.py @@ -62,7 +62,7 @@ def _use_dtensor(*DNDarrays) -> bool: # on evenly distributed DNDarrays and on GPU only for array in DNDarrays: - if not array.is_distributed() or str(array.device)[:3] != "gpu": + if not str(array.device)[:3] == "gpu": return False if array.split is not None and array.gshape[array.split] % array.comm.size != 0: return False @@ -648,28 +648,30 @@ def matmul(a: DNDarray, b: DNDarray, allow_resplit: bool = False) -> DNDarray: sanitation.sanitize_in(a) sanitation.sanitize_in(b) - if _use_dtensor(a, b): - try: - mesh = _get_or_create_mesh(a.device) + if a.is_distributed() or b.is_distributed(): + # route through DTensor if it makes sense + if _use_dtensor(a, b): + try: + mesh = _get_or_create_mesh(a.device) - dt_a = _to_dtensor(a, mesh) - dt_b = _to_dtensor(b, mesh) + dt_a = _to_dtensor(a, mesh) + dt_b = _to_dtensor(b, mesh) - dt_c = torch.matmul(dt_a, dt_b) + dt_c = torch.matmul(dt_a, dt_b) - # heat infers the final split directly from the inputs - expected_split = a.split if a.split is not None else b.split - local_c = _from_dtensor(dt_c, expected_split) + # heat infers the final split directly from the inputs + expected_split = a.split if a.split is not None else b.split + local_c = _from_dtensor(dt_c, expected_split) - gshape_c = (a.gshape[0], b.gshape[1]) + gshape_c = (a.gshape[0], b.gshape[1]) - return DNDarray( - local_c, gshape_c, a.dtype, expected_split, a.device, a.comm, balanced=True - ) - except Exception as e: - import logging + return DNDarray( + local_c, gshape_c, a.dtype, expected_split, a.device, a.comm, balanced=True + ) + except Exception as e: + import logging - logging.warning(f"dtensor routing failed, falling back to mpi: {e}") + logging.warning(f"dtensor routing failed, falling back to mpi: {e}") batch_dim = max(a.ndim, b.ndim) - 2 # -1 for vector vector multiplication batched = batch_dim > 0