matmul dtensor experiments

heat/core/communication.py

@@ -2468,5 +2468,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

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
@@ -51,6 +56,38 @@
 ]
 
 
+def _use_dtensor(*DNDarrays) -> bool:
+    if not _DTENSOR_AVAILABLE:
+        return False
+
+    # on evenly distributed DNDarrays and on GPU only
+    for array in DNDarrays:
+        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
+
+    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:
     """
     Computes an upper estimate for the largest singular value of the input 2D DNDarray.
@@ -611,6 +648,31 @@ def matmul(a: DNDarray, b: DNDarray, allow_resplit: bool = False) -> DNDarray:
     sanitation.sanitize_in(a)
     sanitation.sanitize_in(b)
 
+    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_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