diff --git a/heat/core/communication.py b/heat/core/communication.py index 4df8ac7cab..c6ca3edc13 100644 --- a/heat/core/communication.py +++ b/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 diff --git a/heat/core/linalg/basics.py b/heat/core/linalg/basics.py index bd8240ad76..1fdb40333e 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 @@ -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