feat(parallel): add topology-aware DLO planning

astrai/parallel/__init__.py

@@ -15,9 +15,19 @@
     get_rank,
     get_world_size,
     only_on_rank,
+    resolve_local_device_index,
     setup_parallel,
     spawn_parallel_fn,
 )
+from astrai.parallel.topology import (
+    DLOMeasurement,
+    DLOTopologyPlan,
+    GPUTopology,
+    build_parallel_groups,
+    parse_device_order,
+    parse_nvidia_topology,
+    select_dlo_plan,
+)
 
 __all__ = [
     "get_world_size",
@@ -36,4 +46,12 @@
     "FSDPExecutor",
     "create_ref_model",
     "broadcast_state_dict",
+    "DLOMeasurement",
+    "DLOTopologyPlan",
+    "GPUTopology",
+    "build_parallel_groups",
+    "parse_device_order",
+    "parse_nvidia_topology",
+    "resolve_local_device_index",
+    "select_dlo_plan",
 ]

astrai/parallel/executor.py

@@ -291,7 +291,15 @@ def _prepare_model(self, model: nn.Module) -> nn.Module:
         if not self.use_distributed:
             logger.warning("DDP backend selected but world_size=1, model not wrapped")
             return model
-        local_rank = int(os.environ.get("LOCAL_RANK", get_rank()))
+        local_device = os.environ.get("LOCAL_DEVICE")
+        device_index = (
+            torch.device(local_device).index if local_device is not None else None
+        )
+        local_rank = (
+            device_index
+            if device_index is not None
+            else int(os.environ.get("LOCAL_RANK", get_rank()))
+        )
         model = DDP(
             model,
             device_ids=[local_rank],

astrai/parallel/setup.py

@@ -12,11 +12,29 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
+from astrai.parallel.topology import parse_device_order
 from astrai.signal_handler import install_early_signal_handlers
 
 logger = logging.getLogger(__name__)
 
 
+def resolve_local_device_index(
+    local_rank: int,
+    local_world_size: int,
+    device_type: str,
+) -> int:
+    """Map a logical local rank to an accelerator selected by the planner."""
+
+    if not 0 <= local_rank < local_world_size:
+        raise ValueError(
+            f"local rank {local_rank} is outside local world size {local_world_size}"
+        )
+    value = os.environ.get("ASTRAI_DEVICE_ORDER")
+    if value is None or device_type == "cpu":
+        return local_rank
+    return parse_device_order(value, local_world_size)[local_rank]
+
+
 def find_free_port() -> str:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         s.bind(("", 0))
@@ -54,15 +72,18 @@ def setup_parallel(
         yield dist.group.WORLD
         return
 
+    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", str(world_size)))
+    device_index = resolve_local_device_index(local_rank, local_world_size, device_type)
+
     if world_size <= 1:
-        device_id = torch.device(device_type, local_rank)
+        device_id = torch.device(device_type, device_index)
         os.environ["LOCAL_RANK"] = str(local_rank)
         os.environ["WORLD_SIZE"] = "1"
         os.environ["LOCAL_DEVICE"] = str(device_id)
         yield None
         return
 
-    device_id = torch.device(device_type, local_rank)
+    device_id = torch.device(device_type, device_index)
 
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port