From 7e876d69b599c647bb8c4d375e85dbb8cd3b8b38 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 17:00:10 +0800 Subject: [PATCH] feat(parallel): add topology-aware DLO planning --- astrai/parallel/__init__.py | 18 + astrai/parallel/executor.py | 10 +- astrai/parallel/setup.py | 25 +- astrai/parallel/topology.py | 359 ++++++++++++++++ benchmarks/infraswe/README.md | 45 ++ .../infraswe/astrai-topology-dlo-draft.json | 105 +++++ benchmarks/results/dlo_topology_8x5060ti.json | 157 +++++++ .../dlo_topology_8x5060ti_infraswe_score.json | 114 +++++ .../benchmarks/topology_aware_dlo_8x5060ti.md | 87 ++++ docs/guides/distributed.md | 62 +++ scripts/tools/benchmark_dlo_topology.py | 395 ++++++++++++++++++ scripts/tools/plan_dlo_topology.py | 114 +++++ tests/parallel/test_infraswe_draft.py | 89 ++++ tests/parallel/test_parallel.py | 85 +++- tests/parallel/test_topology.py | 111 +++++ tests/parallel/test_topology_cli.py | 36 ++ 16 files changed, 1808 insertions(+), 4 deletions(-) create mode 100644 astrai/parallel/topology.py create mode 100644 benchmarks/infraswe/README.md create mode 100644 benchmarks/infraswe/astrai-topology-dlo-draft.json create mode 100644 benchmarks/results/dlo_topology_8x5060ti.json create mode 100644 benchmarks/results/dlo_topology_8x5060ti_infraswe_score.json create mode 100644 docs/benchmarks/topology_aware_dlo_8x5060ti.md create mode 100644 scripts/tools/benchmark_dlo_topology.py create mode 100644 scripts/tools/plan_dlo_topology.py create mode 100644 tests/parallel/test_infraswe_draft.py create mode 100644 tests/parallel/test_topology.py create mode 100644 tests/parallel/test_topology_cli.py diff --git a/astrai/parallel/__init__.py b/astrai/parallel/__init__.py index b565f8d4..d3329a92 100644 --- a/astrai/parallel/__init__.py +++ b/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", ] diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index d0c12c86..8342d4f1 100644 --- a/astrai/parallel/executor.py +++ b/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], diff --git a/astrai/parallel/setup.py b/astrai/parallel/setup.py index 4709d75a..df705da3 100644 --- a/astrai/parallel/setup.py +++ b/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 diff --git a/astrai/parallel/topology.py b/astrai/parallel/topology.py new file mode 100644 index 00000000..5f82dc9f --- /dev/null +++ b/astrai/parallel/topology.py @@ -0,0 +1,359 @@ +"""Topology-aware device placement for DP/SP layerwise offload. + +The planner models the rank order used by vLLM-Omni diffusion workers: sequence +parallel ranks are contiguous and data parallel ranks are the outer dimension. +Distributed layerwise offload (DLO) therefore communicates over DP when DP is +larger than one, otherwise it communicates over SP. Topology labels are only a +fallback: measured collective timings take precedence when supplied. +""" + +from __future__ import annotations + +import itertools +import re +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass + +_ANSI_ESCAPE = re.compile(r"\x1b(?:\[[0-?]*[ -/]*[@-~]|\][^\x07]*(?:\x07|\x1b\\))") +_GPU_TOKEN = re.compile(r"^GPU(\d+)$") +_KNOWN_LINKS = {"X", "PIX", "PXB", "PHB", "NODE", "SYS"} +_LINK_WEIGHT = { + "SYS": 1.0, + "NODE": 2.0, + "PHB": 3.0, + "PXB": 4.0, + "PIX": 5.0, + "X": 7.0, +} + + +def _link_weight(link: str) -> float: + if link.startswith("NV") and link[2:].isdigit(): + return 6.0 + int(link[2:]) / 100.0 + try: + return _LINK_WEIGHT[link] + except KeyError as exc: + raise ValueError(f"unknown GPU topology link: {link}") from exc + + +@dataclass(frozen=True) +class GPUTopology: + """A symmetric GPU connectivity matrix parsed from ``nvidia-smi topo -m``.""" + + devices: tuple[int, ...] + links: Mapping[tuple[int, int], str] + + def link(self, left: int, right: int) -> str: + if left == right: + return "X" + try: + return self.links[(left, right)] + except KeyError as exc: + raise ValueError(f"missing topology link GPU{left} -> GPU{right}") from exc + + def affinity(self, left: int, right: int) -> float: + return _link_weight(self.link(left, right)) + + +@dataclass(frozen=True) +class DLOMeasurement: + """Measured steady-state collective cost for one physical rank mapping.""" + + dp_size: int + sp_size: int + device_order: tuple[int, ...] + dlo_all_gather_ms: float + sp_collective_ms: float + + @property + def combined_ms(self) -> float: + return self.dlo_all_gather_ms + self.sp_collective_ms + + +@dataclass(frozen=True) +class DLOTopologyPlan: + """A DP/SP shape and logical-rank-to-physical-device mapping.""" + + dp_size: int + sp_size: int + device_order: tuple[int, ...] + dlo_group_kind: str + dlo_groups: tuple[tuple[int, ...], ...] + sp_groups: tuple[tuple[int, ...], ...] + selection_source: str + score: float + + def as_dict(self) -> dict[str, object]: + return { + "dp_size": self.dp_size, + "sp_size": self.sp_size, + "device_order": list(self.device_order), + "dlo_group_kind": self.dlo_group_kind, + "dlo_groups": [list(group) for group in self.dlo_groups], + "sp_groups": [list(group) for group in self.sp_groups], + "selection_source": self.selection_source, + "score": self.score, + } + + +def parse_nvidia_topology(text: str) -> GPUTopology: + """Parse and validate the GPU matrix emitted by ``nvidia-smi topo -m``. + + The parser intentionally fails closed on partial or asymmetric matrices. + CPU-affinity and NUMA columns following the GPU matrix are ignored. + """ + + clean = _ANSI_ESCAPE.sub("", text) + lines = [line.strip() for line in clean.splitlines() if line.strip()] + header_index = next( + ( + index + for index, line in enumerate(lines) + if ( + len([token for token in line.split() if _GPU_TOKEN.match(token)]) >= 2 + or ( + len([token for token in line.split() if _GPU_TOKEN.match(token)]) + == 1 + and len(line.split()) > 1 + and line.split()[1] not in _KNOWN_LINKS + ) + ) + ), + None, + ) + if header_index is None: + raise ValueError("nvidia-smi topology output has no GPU header") + + header = lines[header_index].split() + devices = tuple( + int(match.group(1)) + for token in header + if (match := _GPU_TOKEN.match(token)) is not None + ) + if not devices or len(set(devices)) != len(devices): + raise ValueError("GPU topology header is empty or contains duplicate devices") + + rows: dict[int, tuple[str, ...]] = {} + for line in lines[header_index + 1 :]: + tokens = line.split() + match = _GPU_TOKEN.match(tokens[0]) if tokens else None + if match is None: + continue + device = int(match.group(1)) + if device in rows: + raise ValueError(f"duplicate GPU{device} topology row") + if len(tokens) < len(devices) + 1: + raise ValueError(f"incomplete GPU{device} topology row") + rows[device] = tuple(tokens[1 : len(devices) + 1]) + + if set(rows) != set(devices): + missing = sorted(set(devices) - set(rows)) + extra = sorted(set(rows) - set(devices)) + raise ValueError(f"GPU topology row mismatch: missing={missing}, extra={extra}") + + links: dict[tuple[int, int], str] = {} + for row_device in devices: + for column, column_device in enumerate(devices): + link = rows[row_device][column] + if link not in _KNOWN_LINKS and not ( + link.startswith("NV") and link[2:].isdigit() + ): + raise ValueError( + f"unknown topology token {link!r} for GPU{row_device}/GPU{column_device}" + ) + if row_device == column_device and link != "X": + raise ValueError(f"GPU{row_device} diagonal must be X, got {link}") + links[(row_device, column_device)] = link + + for left in devices: + for right in devices: + if links[(left, right)] != links[(right, left)]: + raise ValueError( + f"asymmetric GPU topology: GPU{left}/GPU{right} is " + f"{links[(left, right)]}/{links[(right, left)]}" + ) + return GPUTopology(devices=devices, links=links) + + +def build_parallel_groups( + device_order: Sequence[int], dp_size: int, sp_size: int +) -> tuple[tuple[tuple[int, ...], ...], tuple[tuple[int, ...], ...]]: + """Return physical DP and SP groups for an SP-fastest logical rank order.""" + + order = tuple(device_order) + if dp_size <= 0 or sp_size <= 0: + raise ValueError("DP and SP sizes must be positive") + if dp_size * sp_size != len(order): + raise ValueError("DP * SP must equal the device count") + if len(set(order)) != len(order): + raise ValueError("device_order must not contain duplicates") + + sp_groups = tuple( + tuple(order[dp_rank * sp_size : (dp_rank + 1) * sp_size]) + for dp_rank in range(dp_size) + ) + dp_groups = tuple( + tuple(order[dp_rank * sp_size + sp_rank] for dp_rank in range(dp_size)) + for sp_rank in range(sp_size) + ) + return dp_groups, sp_groups + + +def dlo_groups_for_plan( + device_order: Sequence[int], dp_size: int, sp_size: int +) -> tuple[str, tuple[tuple[int, ...], ...], tuple[tuple[int, ...], ...]]: + """Select the DLO group using DP-first, SP-second backend semantics.""" + + dp_groups, sp_groups = build_parallel_groups(device_order, dp_size, sp_size) + if dp_size > 1: + return "dp", dp_groups, sp_groups + if sp_size > 1: + return "sp", sp_groups, sp_groups + return "rank-local", tuple((device,) for device in device_order), sp_groups + + +def _mean_group_affinity( + topology: GPUTopology, groups: Sequence[Sequence[int]] +) -> float: + values = [ + topology.affinity(left, right) + for group in groups + for left_index, left in enumerate(group) + for right in group[left_index + 1 :] + ] + return sum(values) / len(values) if values else _link_weight("X") + + +def topology_score( + topology: GPUTopology, + device_order: Sequence[int], + dp_size: int, + sp_size: int, + *, + dlo_weight: float = 4.0, + sp_weight: float = 1.0, +) -> float: + """Score a mapping using link labels when measurements are unavailable.""" + + if set(device_order) != set(topology.devices): + raise ValueError("device_order must contain every topology device exactly once") + _, dlo_groups, sp_groups = dlo_groups_for_plan(device_order, dp_size, sp_size) + return dlo_weight * _mean_group_affinity( + topology, dlo_groups + ) + sp_weight * _mean_group_affinity(topology, sp_groups) + + +def optimize_device_order( + topology: GPUTopology, + dp_size: int, + sp_size: int, + *, + exhaustive_limit: int = 8, +) -> tuple[tuple[int, ...], float]: + """Find a deterministic topology-label optimum for a fixed DP/SP shape.""" + + if dp_size * sp_size != len(topology.devices): + raise ValueError("DP * SP must equal the topology device count") + natural = tuple(sorted(topology.devices)) + if len(natural) > exhaustive_limit: + return natural, topology_score(topology, natural, dp_size, sp_size) + + best_order = natural + best_score = float("-inf") + for order in itertools.permutations(natural): + score = topology_score(topology, order, dp_size, sp_size) + if score > best_score or (score == best_score and order < best_order): + best_order = order + best_score = score + return best_order, best_score + + +def select_dlo_plan( + topology: GPUTopology, + *, + concurrent_requests: int = 1, + dp_size: int | None = None, + sp_size: int | None = None, + measurements: Iterable[DLOMeasurement] = (), +) -> DLOTopologyPlan: + """Select a DP/SP plan, preferring valid measured collective results. + + When no shape is requested explicitly, the largest DP divisor that does + not exceed ``concurrent_requests`` is selected. This prevents a single + request from being silently assigned to multiple data-parallel replicas. + """ + + world_size = len(topology.devices) + if concurrent_requests <= 0: + raise ValueError("concurrent_requests must be positive") + if (dp_size is None) != (sp_size is None): + raise ValueError("dp_size and sp_size must be provided together") + if dp_size is None: + eligible = [ + candidate + for candidate in range(1, world_size + 1) + if world_size % candidate == 0 and candidate <= concurrent_requests + ] + dp_size = max(eligible) + sp_size = world_size // dp_size + assert sp_size is not None + if dp_size * sp_size != world_size: + raise ValueError("DP * SP must equal the topology device count") + if dp_size > concurrent_requests: + raise ValueError("DP size cannot exceed concurrent request capacity") + + valid_measurements = [] + for measurement in measurements: + if measurement.dp_size != dp_size or measurement.sp_size != sp_size: + continue + if set(measurement.device_order) != set(topology.devices): + raise ValueError("measured device_order does not match topology devices") + if measurement.dlo_all_gather_ms <= 0 or measurement.sp_collective_ms <= 0: + raise ValueError("measured collective latency must be positive") + valid_measurements.append(measurement) + + if valid_measurements: + best_measurement = min( + valid_measurements, + key=lambda item: (item.combined_ms, item.device_order), + ) + order = best_measurement.device_order + score = best_measurement.combined_ms + source = "measured-collectives" + else: + order, score = optimize_device_order(topology, dp_size, sp_size) + source = "topology-label-fallback" + + kind, dlo_groups, sp_groups = dlo_groups_for_plan(order, dp_size, sp_size) + return DLOTopologyPlan( + dp_size=dp_size, + sp_size=sp_size, + device_order=order, + dlo_group_kind=kind, + dlo_groups=dlo_groups, + sp_groups=sp_groups, + selection_source=source, + score=score, + ) + + +def parse_device_order(value: str, local_world_size: int) -> tuple[int, ...]: + """Parse ``ASTRAI_DEVICE_ORDER`` as a complete logical-to-visible map.""" + + try: + order = tuple(int(item.strip()) for item in value.split(",")) + except ValueError as exc: + raise ValueError( + "ASTRAI_DEVICE_ORDER must be comma-separated integers" + ) from exc + if len(order) != local_world_size: + raise ValueError( + f"ASTRAI_DEVICE_ORDER has {len(order)} devices, expected {local_world_size}" + ) + expected = set(range(local_world_size)) + if set(order) != expected: + raise ValueError( + "ASTRAI_DEVICE_ORDER must be a permutation of local visible device " + f"indices 0..{local_world_size - 1}" + ) + return order diff --git a/benchmarks/infraswe/README.md b/benchmarks/infraswe/README.md new file mode 100644 index 00000000..1ef5117a --- /dev/null +++ b/benchmarks/infraswe/README.md @@ -0,0 +1,45 @@ +# InfraSWE Draft: AstrAI topology-aware DLO + +This directory binds the benchmark to AstrAI as an explicit repository target. AstrAI +is not one of InfraSWE v0.5's ten built-in projects, so using a built-in default would +silently score the change against the wrong host project. The Draft therefore uses +`target.mode = repository` and remains `D3-contract-proposed`; it does not claim human +review, sealing, or official evaluation. + +The Draft was validated and resolved with InfraSWE commit +`811bc775ed5b3a6ec853219245f3469f78818020`: + +```bash +PYTHONPATH=src .venv/bin/python -m infraswe.cli draft validate \ + /path/to/AstrAI/benchmarks/infraswe/astrai-topology-dlo-draft.json + +PYTHONPATH=src .venv/bin/python -m infraswe.cli draft resolve \ + --local-draft /path/to/AstrAI/benchmarks/infraswe/astrai-topology-dlo-draft.json \ + --output /tmp/astrai-draft-resolution.json +``` + +Digest construction is deterministic: + +- target/baseline: SHA-256 of the target commit ID `1fad50d8476abe7d4b4cb527eb6e174c511fd409`; +- candidate: SHA-256 of the ordered per-file SHA-256 list for the planner, launch + integration, benchmark scripts, tests, and distributed guide; +- project profile: the same construction over `pyproject.toml`, `README.md`, the + parallel public API, and distributed guide; +- workload: SHA-256 of `benchmarks/results/dlo_topology_8x5060ti.json`; +- acceptance/probes: ordered per-file digests of the topology tests and captured + benchmark result; +- precedents: ordered per-file digests of vLLM-Omni's DLO design and backend at the + retrieval cutoff. + +The 7-replay worst-rank collective result is the fast-loop evidence. The same workload +artifact also records successful 2-step and 10-step MiniMax-H3 T2VA requests with +online FP8 and DLO AllGather across eight GPUs. Both remain provisional until an +AstrAI maintainer reviews the contract and advances the Draft lifecycle. + +Running InfraSWE's frozen `project-fit-kernel-v0.5` formula over the visible +acceptance evidence produces a diagnostic ProjectFit of **86.73/100** and a +BenchmarkTrust score of **93.06/100**. The machine-readable score card is +`benchmarks/results/dlo_topology_8x5060ti_infraswe_score.json`; it records every +subcomponent input and its rationale. The score is deliberately marked non-official. +InfraSWE leaves the official score unresolved until the Draft is sealed, hidden probes +are complete, and the evidence manifest is verified. diff --git a/benchmarks/infraswe/astrai-topology-dlo-draft.json b/benchmarks/infraswe/astrai-topology-dlo-draft.json new file mode 100644 index 00000000..4b600183 --- /dev/null +++ b/benchmarks/infraswe/astrai-topology-dlo-draft.json @@ -0,0 +1,105 @@ +{ + "schema_version": "0.5", + "draft": { + "id": "astrai-topology-aware-dlo-dpsp-v1", + "revision": 1, + "state": "D3-contract-proposed", + "created_by": "0z5a" + }, + "target": { + "mode": "repository", + "repository": "https://github.com/ViperEkura/AstrAI", + "revision": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74", + "project_profile_sha256": "sha256:1a013d117f58397189ba0f07b8bf417abf4905a165c3098a1957d2fcb228cbf6" + }, + "candidate": { + "kind": "git-diff", + "revision": "sha256:9ef929c8f646ffebb319c0c5e559d269371261c8a789d20d7409112844f582a9", + "intent": "integrate", + "implementation_kind": "framework", + "entrypoints": [ + "astrai.parallel.topology.select_dlo_plan", + "astrai.parallel.setup.resolve_local_device_index", + "astrai.parallel.executor.DDPExecutor._prepare_model", + "scripts/tools/benchmark_dlo_topology.py" + ], + "operator_family": "communication-collective", + "phase": "communication", + "backend": "cuda", + "primary_host_candidate": "astrai" + }, + "baseline": { + "mode": "target-head", + "revision": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74" + }, + "deployment": { + "workload_portfolio": { + "id": "astrai-topology-dlo-collectives-v1", + "sha256": "sha256:fa215b71961e452f1642dfaf2f53062715a6350f12ea011b36b404abaa461130", + "path": "benchmarks/results/dlo_topology_8x5060ti.json" + }, + "required_cells": [ + "cuda-sm120-8x-rtx5060ti" + ], + "optional_cells": [ + "cuda-sm90-8x", + "cuda-sm100-8x" + ], + "request_or_step_protocol": { + "id": "astrai-topology-dlo-launch-v1", + "sha256": "sha256:e47fd486ddd2fe59959b5a0e585cd3a38c4c816e9fec1912fb71fed752765e05", + "path": "docs/guides/distributed.md" + } + }, + "retrieval": { + "enabled": true, + "corpus_cutoff": "2026-09-02T00:00:00Z", + "sources": [ + "target-code", + "merged-prs", + "rejected-prs", + "review-comments", + "ci-failures", + "release-notes" + ], + "precedent_set_sha256": "sha256:49d9686b940054dcf0b4414baf1c2d2cabb37382d129fbdecb9269d7d5121fd1" + }, + "acceptance_contract": { + "status": "proposed", + "path": "tests/parallel/test_topology.py", + "sha256": "sha256:7ce14b99f74ca9fdb04c06a6c5225792535d412272996e3f89b7e486ac6bf93a", + "probe_set_sha256": "sha256:a43bdafee50ab450264c33ad8c655fb10c7f150a26a38cd21bba71791009e58f", + "hidden_probe_policy_sha256": "sha256:3a9abef329e13ba9da1d678abefcbc0953e19f128223a6797330b4a7d4e10ae8" + }, + "project_objectives": { + "edge_ecosystem_policy": "experimental", + "profile_set_sha256": "sha256:f98050cb02f598b88b6819b288a0768af9c65d996e7684b1bbfb920198dd2005" + }, + "benchmark_loop": { + "fast_stage_max_official_fraction": 0.05, + "affected_stage_max_official_fraction": 0.2, + "official_replays": 7, + "early_exit_on_hard_gate": true, + "affected_case_selection": "required", + "benchmark_budget_policy_id": "draft-staged-budget-v0.5", + "evidence_policy_id": "v0.4-evidence-ladder-plus-seal-v0.5", + "precompile": { + "mode": "auto", + "trigger": "when-compilation-required", + "cache_policy": "content-addressed-evidence-identity", + "cache_miss_action": "precompile-before-timed-cases", + "timing_phases": [ + "precompile", + "cold-start", + "steady-state" + ], + "steady_state_compile_allowed": false + } + }, + "scoring": { + "formula_template_id": "project-fit-kernel-v0.5", + "provisional_scoring_allowed": true, + "official_scoring_requires_seal": true, + "project_season": "astrai-2026q3" + } +} diff --git a/benchmarks/results/dlo_topology_8x5060ti.json b/benchmarks/results/dlo_topology_8x5060ti.json new file mode 100644 index 00000000..49ed05cd --- /dev/null +++ b/benchmarks/results/dlo_topology_8x5060ti.json @@ -0,0 +1,157 @@ +{ + "schema_version": "1", + "source_raw_sha256": "98fbb8966b7aa0a032ca1fa269ee1cd240d34b0bc4a5d918e3fdad31b040212f", + "environment": { + "timestamp_utc": "2026-09-02T08:11:38.370183+00:00", + "gpu": "8x NVIDIA GeForce RTX 5060 Ti 16 GiB", + "compute_capability": "12.0", + "torch": "2.11.0+cu128", + "cuda": "12.8", + "nccl": "2.28.9", + "topology": "PHB pairs (0,1), (2,3), (4,5), (6,7); all other pairs NODE; NUMA 0", + "dlo_full_payload_bytes": 268435456, + "sp_per_rank_payload_bytes": 67108864, + "warmup": 5, + "iterations_per_trial": 20, + "trials": 7, + "timing_scope": "worst-rank steady-state collective" + }, + "results": [ + { + "name": "dp1-sp8-natural", + "dp_size": 1, + "sp_size": 8, + "device_order": [0, 1, 2, 3, 4, 5, 6, 7], + "dlo_group_kind": "sp", + "dlo_all_gather_median_ms": 41.1263, + "dlo_all_gather_p99_ms": 41.1770, + "sp_all_to_all_median_ms": 20.7118, + "sp_all_to_all_p99_ms": 20.8178 + }, + { + "name": "dp2-sp4-natural", + "dp_size": 2, + "sp_size": 4, + "device_order": [0, 1, 2, 3, 4, 5, 6, 7], + "dlo_group_kind": "dp", + "dlo_all_gather_median_ms": 57.3858, + "dlo_all_gather_p99_ms": 67.7326, + "sp_all_to_all_median_ms": 18.4796, + "sp_all_to_all_p99_ms": 18.5259 + }, + { + "name": "dp2-sp4-topology", + "dp_size": 2, + "sp_size": 4, + "device_order": [0, 2, 4, 6, 1, 3, 5, 7], + "dlo_group_kind": "dp", + "dlo_all_gather_median_ms": 52.6750, + "dlo_all_gather_p99_ms": 52.8270, + "sp_all_to_all_median_ms": 18.8081, + "sp_all_to_all_p99_ms": 18.8933 + }, + { + "name": "dp4-sp2-natural", + "dp_size": 4, + "sp_size": 2, + "device_order": [0, 1, 2, 3, 4, 5, 6, 7], + "dlo_group_kind": "dp", + "dlo_all_gather_median_ms": 43.9203, + "dlo_all_gather_p99_ms": 43.9913, + "sp_all_to_all_median_ms": 11.3262, + "sp_all_to_all_p99_ms": 11.5051 + }, + { + "name": "dp4-sp2-topology", + "dp_size": 4, + "sp_size": 2, + "device_order": [0, 2, 1, 3, 4, 6, 5, 7], + "dlo_group_kind": "dp", + "dlo_all_gather_median_ms": 44.0285, + "dlo_all_gather_p99_ms": 44.0651, + "sp_all_to_all_median_ms": 16.5605, + "sp_all_to_all_p99_ms": 16.6675 + }, + { + "name": "dp8-sp1-natural", + "dp_size": 8, + "sp_size": 1, + "device_order": [0, 1, 2, 3, 4, 5, 6, 7], + "dlo_group_kind": "dp", + "dlo_all_gather_median_ms": 41.0465, + "dlo_all_gather_p99_ms": 41.0798, + "sp_all_to_all_median_ms": 0.3494, + "sp_all_to_all_p99_ms": 0.3497 + } + ], + "selected": { + "single_request": "dp1-sp8-natural", + "two_requests": "dp2-sp4-topology", + "four_requests": "dp4-sp2-natural", + "eight_requests": "dp8-sp1-natural" + }, + "end_to_end": { + "model": "MiniMax-H3", + "runtime": { + "vllm_omni_commit": "e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6", + "torch": "2.13.0+cu132", + "cuda": "13.2", + "nccl": "2.29.7", + "quantization": "online-fp8", + "parallelism": "dp1-sp8-text-tp8-vae-tile8", + "distributed_layerwise_offload": "all-gather", + "engine_initialization_s": 306.46, + "idle_process_gpu_memory_gib": 1.86, + "maximum_full_fp8_block_mib": 615.6, + "maximum_shard_mib_per_rank": 76.9 + }, + "request": { + "task": "t2va", + "width": 832, + "height": 480, + "fps": 24, + "requested_duration_s": 4.0, + "generated_frames": 107, + "generated_audio_duration_s": 4.45, + "audio_sample_rate_hz": 32000, + "seed": 1101 + }, + "samples": [ + { + "name": "first-valid-request", + "steps": 2, + "client_e2e_s": 17.317, + "server_e2e_s": 16.877134, + "denoise_step_latency_ms": 8437.763, + "peak_gpu_memory_mib": 13998, + "all_gpus_peak_utilization_percent": 100 + }, + { + "name": "warm-request", + "steps": 2, + "client_e2e_s": 13.714, + "server_e2e_s": 13.291982, + "denoise_step_latency_ms": 6645.439, + "peak_gpu_memory_mib": 14174, + "all_gpus_peak_utilization_percent": 100 + }, + { + "name": "stability-request", + "steps": 10, + "client_e2e_s": 71.719, + "server_e2e_s": 71.320681, + "denoise_step_latency_ms": 7131.966, + "peak_gpu_memory_mib": 14206, + "all_gpus_peak_utilization_percent": 100, + "output_bytes": 1022884, + "output_sha256": "1b3e2d9a13ed292c2bb16ab09fe795b9ec27e4038755ba70ca8c5d72dbc0fdfc", + "video_stream": "h264 832x480 24fps", + "audio_stream": "aac stereo 32000hz" + } + ], + "excluded_probe": { + "steps": 1, + "reason": "MiniMax-H3 sigma schedules require at least two entries" + } + } +} diff --git a/benchmarks/results/dlo_topology_8x5060ti_infraswe_score.json b/benchmarks/results/dlo_topology_8x5060ti_infraswe_score.json new file mode 100644 index 00000000..d6635f31 --- /dev/null +++ b/benchmarks/results/dlo_topology_8x5060ti_infraswe_score.json @@ -0,0 +1,114 @@ +{ + "schema_version": "0.5", + "score_kind": "diagnostic-project-fit", + "score_is_official": false, + "draft_id": "astrai-topology-aware-dlo-dpsp-v1", + "draft_state": "D3-contract-proposed", + "formula_template_id": "project-fit-kernel-v0.5", + "diagnostic_project_fit_100": 86.73007069241775, + "component_values": { + "evolutionary_maintainability": 0.7420138271431783, + "project_contract_fit": 1.0, + "performance_reuse_utilization": 0.9306048591020996, + "operational_fit": 0.9173147546424018 + }, + "component_floors": { + "evolutionary_maintainability": 0.6, + "project_contract_fit": 0.6, + "performance_reuse_utilization": 0.4, + "operational_fit": 0.6 + }, + "subcomponent_inputs": { + "evolutionary_maintainability": { + "evolution": 0.5, + "locality": 0.8, + "tests": 1.0, + "failure": 1.0, + "contract": 1.0 + }, + "project_contract_fit": { + "integration": 1.0, + "interface": 1.0, + "lifecycle": 1.0, + "buildtest": 1.0, + "policy": 1.0 + }, + "performance_reuse_utilization": { + "attainment": 1.0, + "coverage": 0.75, + "retention": 1.0, + "family": 1.0, + "compile": 1.0 + }, + "operational_fit": { + "replay": 1.0, + "load": 0.75, + "resource": 1.0, + "coldsteady": 1.0 + } + }, + "input_rationale": { + "evolution": "First AstrAI integration; no upstream maintenance or release history exists yet.", + "locality": "The runtime change is isolated to parallel planning, setup, and DDP device selection, with separate tools, tests, and documentation.", + "tests": "The complete AstrAI suite passed: 631 passed and 103 environment-dependent tests skipped.", + "failure": "Malformed, partial, asymmetric, and unknown topology inputs fail closed; invalid device permutations are rejected.", + "contract": "Digest-bound planner, CLI, launch, and DDP mapping acceptance tests all passed.", + "integration": "The planner is exported through astrai.parallel and consumed by the existing setup and executor paths.", + "interface": "Environment variables and planner CLI output are documented and tested.", + "lifecycle": "Logical LOCAL_RANK is preserved while LOCAL_DEVICE selects the mapped accelerator.", + "buildtest": "The full test suite, Ruff formatting, and import-order checks passed.", + "policy": "The candidate adds no dependency and leaves natural rank order as the safe fallback.", + "attainment": "Measured selection chooses the best tested mapping for each request-concurrency shape.", + "coverage": "The required 8x RTX 5060 Ti cell was tested; optional SM90 and SM100 cells were not.", + "retention": "The planner retains natural order for DP4xSP2, avoiding the measured 9.67% topology-label regression.", + "family": "DP1/2/4/8 shapes and MiniMax-H3 online-FP8 end-to-end inference were covered.", + "compile": "The Python planner adds no timed compilation path.", + "replay": "Collectives use seven trials with worst-rank median and p99 reporting.", + "load": "Single-request MiniMax-H3 was exercised end to end; DP2/4/8 deployed multi-request service load was not.", + "resource": "Peak GPU memory, retained headroom, all-GPU utilization, and clean service shutdown were checked.", + "coldsteady": "Initialization, first valid request, warm request, and 10-step stability request were recorded." + }, + "benchmark_trust": { + "formula_version": "benchmark-trust-v0.5", + "status": "scored", + "score_100": 93.06048591020996, + "components": { + "reproducibility": 1.0, + "evidence": 0.75, + "statistics": 1.0, + "environment": 1.0 + }, + "failure_codes": [ + "RAW_CAPTURE_NOT_CHECKED_IN", + "DRAFT_UNSEALED" + ] + }, + "official_project_fit": { + "status": "unresolved", + "score_100": null, + "failure_codes": [ + "DRAFT_SEAL_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED" + ] + }, + "comparison_cell": { + "target_project_profile_sha256": "sha256:1a013d117f58397189ba0f07b8bf417abf4905a165c3098a1957d2fcb228cbf6", + "target_repository_or_baseline_sha256": "sha256:671ae3f0331ace4f60b4331e963a61987662a7e3234596d0f057834c307acb74", + "semantic_contract_sha256": "sha256:e47fd486ddd2fe59959b5a0e585cd3a38c4c816e9fec1912fb71fed752765e05", + "acceptance_contract_sha256": "sha256:7ce14b99f74ca9fdb04c06a6c5225792535d412272996e3f89b7e486ac6bf93a", + "probe_set_sha256": "sha256:a43bdafee50ab450264c33ad8c655fb10c7f150a26a38cd21bba71791009e58f", + "workload_portfolio_sha256": "sha256:fa215b71961e452f1642dfaf2f53062715a6350f12ea011b36b404abaa461130", + "performance_target_sha256": "sha256:698990c49a23964b0f03b3f77a92e81e5417c27885c70c185695ed2b33eb1f61", + "required_deployment_cell_set_sha256": "sha256:41bdf6bc1d59535dadd4f2c70bce7765fff1824637d44741da815ebf9fedd38a", + "project_season": "astrai-2026q3", + "cross_project_ranking_allowed": false + }, + "execution": { + "infraswe_commit": "811bc775ed5b3a6ec853219245f3469f78818020", + "draft_resolution_sha256": "ce886cb16bd5f0f00faa62449b29014aff1b0bb14e92df818221535f82eba7a4", + "infraswe_draft_engine_tests": "53 passed", + "astrai_tests": "631 passed, 103 skipped", + "astrai_lint": "passed" + } +} diff --git a/docs/benchmarks/topology_aware_dlo_8x5060ti.md b/docs/benchmarks/topology_aware_dlo_8x5060ti.md new file mode 100644 index 00000000..12efa6dc --- /dev/null +++ b/docs/benchmarks/topology_aware_dlo_8x5060ti.md @@ -0,0 +1,87 @@ +# Topology-aware DLO on 8x RTX 5060 Ti + +This benchmark validates AstrAI's DP/SP rank planner with actual NCCL collectives on +the target host. The collective section models finalized one-byte INT8/FP8 DLO +weights. A separate end-to-end section validates the same DLO group semantics with +an online-FP8 MiniMax-H3 service. + +## Environment + +- 8x NVIDIA GeForce RTX 5060 Ti 16 GiB, compute capability 12.0 +- PyTorch 2.11.0+cu128, CUDA 12.8, NCCL 2.28.9 +- One NUMA node; PHB pairs `(0,1)`, `(2,3)`, `(4,5)`, `(6,7)`; other pairs `NODE` +- DLO AllGather reconstructs a 256 MiB full `uint8` payload +- SP all-to-all sends 64 MiB per rank +- 5 warmups, 20 operations per trial, 7 trials +- Each sample is the slowest rank, so one straggler cannot be hidden by averaging + +Reproduce with: + +```bash +python scripts/tools/benchmark_dlo_topology.py \ + --output dlo-topology.json \ + --markdown-output dlo-topology.md \ + --dp-sizes 1,2,4,8 \ + --dlo-payload-mib 256 \ + --sp-payload-mib 64 \ + --warmup 5 \ + --iterations 20 \ + --trials 7 +``` + +## Results + +| Candidate | Device order | DLO group | DLO median / p99 (ms) | SP median / p99 (ms) | Combined median (ms) | +|---|---|---|---:|---:|---:| +| DP1 x SP8 natural | `0,1,2,3,4,5,6,7` | SP8 | 41.1263 / 41.1770 | 20.7118 / 20.8178 | 61.8381 | +| DP2 x SP4 natural | `0,1,2,3,4,5,6,7` | 4x DP2 `NODE` | 57.3858 / 67.7326 | 18.4796 / 18.5259 | 75.8654 | +| DP2 x SP4 topology | `0,2,4,6,1,3,5,7` | 4x DP2 `PHB` | 52.6750 / 52.8270 | 18.8081 / 18.8933 | **71.4831** | +| DP4 x SP2 natural | `0,1,2,3,4,5,6,7` | 2x DP4 `NODE` | 43.9203 / 43.9913 | 11.3262 / 11.5051 | **55.2465** | +| DP4 x SP2 label optimum | `0,2,1,3,4,6,5,7` | 2x DP4 mixed | 44.0285 / 44.0651 | 16.5605 / 16.6675 | 60.5890 | +| DP8 x SP1 natural | `0,1,2,3,4,5,6,7` | DP8 | 41.0465 / 41.0798 | 0.3494 / 0.3497 | 41.3959 | + +The raw 7-trial capture has SHA-256 +`98fbb8966b7aa0a032ca1fa269ee1cd240d34b0bc4a5d918e3fdad31b040212f`; +the checked-in machine-readable summary is +[`benchmarks/results/dlo_topology_8x5060ti.json`](../../benchmarks/results/dlo_topology_8x5060ti.json). + +## MiniMax-H3 online-FP8 end-to-end validation + +The same host also ran current vLLM-Omni main (`e51fe6ec1`) with PyTorch +2.13.0+cu132, CUDA 13.2, and NCCL 2.29.7. MiniMax-H3 was loaded with online FP8, +DP1 x SP8, text-encoder TP8, VAE tile parallelism across eight GPUs, and distributed +layerwise offload with AllGather. Startup selected the SP8 group for DLO, allocated a +615.6 MiB maximum full FP8 block plus a 76.9 MiB shard per rank, and completed engine +initialization in 306.46 seconds. Idle process-scoped GPU memory was 1.86 GiB per +worker. + +All requests used T2VA, 832x480, 24 FPS, a nominal four-second duration, and seed +1101. The model produced 107 frames and 4.45 seconds of 32 kHz audio. The one-step +probe was rejected by H3's expected minimum-two-entry sigma schedule and is excluded +from the benchmark. + +| Sample | Steps | Client E2E (s) | Server E2E (s) | Denoise latency (ms/step) | Peak GPU memory (MiB) | +|---|---:|---:|---:|---:|---:| +| First valid request | 2 | 17.317 | 16.877 | 8,437.763 | 13,998 | +| Warm request | 2 | 13.714 | 13.292 | 6,645.439 | 14,174 | +| Stability request | 10 | 71.719 | 71.320 | 7,131.966 | 14,206 | + +All eight GPUs reached 100% utilization during every valid sample. The 10-step +request retained 2,105 MiB of headroom on a 16,311 MiB GPU and returned a 1,022,884 +byte MP4 (`1b3e2d9a13ed292c2bb16ab09fe795b9ec27e4038755ba70ca8c5d72dbc0fdfc`). +`ffprobe` verified an 832x480 H.264 stream at 24 FPS and stereo AAC at 32 kHz. + +## Interpretation + +For two concurrent requests, mapping each DLO DP2 group to a PHB pair reduces the +combined median by 5.78% and removes the natural mapping's 67.7 ms AllGather p99 +outlier. For four requests, the label-optimal DLO grouping is the wrong end-to-end +choice: it makes SP traffic cross the slower links and raises combined median by +9.67%. The measured planner correctly selects the natural mapping instead. + +DP shapes represent different request concurrency and should not be ranked solely by +the sum in the table. A single MiniMax-H3 request must use DP1 x SP8; DP2/4/8 are +eligible only when at least 2/4/8 independent requests are available. + +This result supports the planner's central policy: topology labels provide a safe +initial candidate, but measured DLO and SP collectives decide the physical mapping. diff --git a/docs/guides/distributed.md b/docs/guides/distributed.md index 063567d9..60d4a10d 100644 --- a/docs/guides/distributed.md +++ b/docs/guides/distributed.md @@ -8,6 +8,7 @@ AstrAI supports three parallel modes: **single GPU** (`none`), **Data Parallel** - [Parallel Modes](#parallel-modes) - [Gradient Accumulation](#gradient-accumulation) - [Process Launching](#process-launching) +- [Topology-aware DP/SP Placement](#topology-aware-dpsp-placement) - [NCCL Troubleshooting](#nccl-troubleshooting) - [Checkpoint Saving](#checkpoint-saving) - [Total Steps Calculation](#total-steps-calculation) @@ -144,6 +145,67 @@ The current training CLI still uses `--nprocs` when calculating scheduler `total Raw Slurm variables such as `SLURM_PROCID`, `SLURM_NTASKS`, and `SLURM_LOCALID` are not recognized automatically. Launch through `torchrun`, or map the scheduler's variables to `RANK`, `WORLD_SIZE`, `LOCAL_RANK`, `MASTER_ADDR`, and `MASTER_PORT` before starting AstrAI. The same requirement applies to launchers that expose only OpenMPI-specific variables. +## Topology-aware DP/SP Placement + +AstrAI can plan and apply a logical-rank-to-device mapping for runtimes that combine +data parallelism (DP), sequence parallelism (SP), and distributed layerwise offload +(DLO). The planner follows the SP-fastest rank layout used by vLLM-Omni diffusion +workers. DLO uses the DP group when `DP > 1`, falls back to the SP group when +`DP = 1` and `SP > 1`, and otherwise remains rank-local. Tensor-parallel groups are +deliberately not treated as DLO weight-sharding groups. + +Capture the live topology and inspect a single-request plan: + +```bash +nvidia-smi topo -m > topology.txt +python scripts/tools/plan_dlo_topology.py \ + --topology-file topology.txt \ + --concurrent-requests 1 \ + --output dlo-plan.json +``` + +A single request cannot select `DP > 1`; on eight visible GPUs the automatic plan is +therefore DP1 x SP8. Multiple independent requests permit a larger DP dimension: + +```bash +python scripts/tools/plan_dlo_topology.py \ + --concurrent-requests 4 \ + --measurements dlo-topology.json +``` + +Topology labels are only a fallback. PCIe labels do not include every relevant +transport detail, so benchmark candidate mappings on the target host and feed the +result back into the planner: + +```bash +python scripts/tools/benchmark_dlo_topology.py \ + --output dlo-topology.json \ + --markdown-output dlo-topology.md \ + --dp-sizes 1,2,4,8 \ + --dlo-payload-mib 256 \ + --sp-payload-mib 64 +``` + +The DLO payload uses `uint8`, matching the byte width of finalized online INT8/FP8 +weights. The benchmark reports the slowest rank in each trial for both DLO AllGather +and SP all-to-all. Measurements are compared only within the same DP/SP shape; the +request concurrency determines the shape, while timings determine its device order. + +To apply the selected mapping to AstrAI's launcher, export the generated +`ASTRAI_DEVICE_ORDER`. Values refer to indices within the current +`CUDA_VISIBLE_DEVICES` set. AstrAI preserves the launcher's logical `LOCAL_RANK` and +records the mapped accelerator separately in `LOCAL_DEVICE`: + +```bash +export ASTRAI_DEVICE_ORDER=0,2,4,6,1,3,5,7 +python scripts/tools/train.py --nprocs=8 --parallel_mode=fsdp ... +``` + +The plan JSON also emits the equivalent `CUDA_VISIBLE_DEVICES` ordering and DP/SP +flags for an external vLLM-Omni MiniMax-H3 launch. AstrAI does not implement the +diffusion offloader itself; it owns the reusable placement, validation, and benchmark +contract used to configure that runtime. + ## NCCL Troubleshooting The following variables are troubleshooting options for hardware or network configurations where NCCL hangs or fails. They are not general requirements and can reduce performance by disabling peer-to-peer or GPUDirect RDMA paths: diff --git a/scripts/tools/benchmark_dlo_topology.py b/scripts/tools/benchmark_dlo_topology.py new file mode 100644 index 00000000..47a1e7b8 --- /dev/null +++ b/scripts/tools/benchmark_dlo_topology.py @@ -0,0 +1,395 @@ +"""Benchmark topology candidates for INT8 DLO AllGather and SP all-to-all.""" + +from __future__ import annotations + +import json +import math +import os +import socket +import statistics +import subprocess +import time +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from functools import partial +from pathlib import Path + +import click +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from astrai.parallel.topology import ( + dlo_groups_for_plan, + optimize_device_order, + parse_nvidia_topology, + topology_score, +) + + +@dataclass(frozen=True) +class BenchmarkCase: + name: str + dp_size: int + sp_size: int + device_order: tuple[int, ...] + topology_label_score: float + + +def percentile(values: Iterable[float], quantile: float) -> float: + ordered = sorted(values) + if not ordered: + raise ValueError("percentile requires samples") + rank = (len(ordered) - 1) * quantile + lower = math.floor(rank) + upper = math.ceil(rank) + if lower == upper: + return ordered[lower] + fraction = rank - lower + return ordered[lower] * (1 - fraction) + ordered[upper] * fraction + + +def summarize(samples_ms: Sequence[float]) -> dict[str, object]: + return { + "median_ms": statistics.median(samples_ms), + "p90_ms": percentile(samples_ms, 0.90), + "p99_ms": percentile(samples_ms, 0.99), + "min_ms": min(samples_ms), + "max_ms": max(samples_ms), + "samples_ms": list(samples_ms), + } + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind(("127.0.0.1", 0)) + return int(listener.getsockname()[1]) + + +def _make_groups( + dlo_groups: Sequence[tuple[int, ...]], + sp_groups: Sequence[tuple[int, ...]], +) -> dict[tuple[int, ...], dist.ProcessGroup]: + unique_groups = tuple(dict.fromkeys((*dlo_groups, *sp_groups))) + return { + ranks: dist.new_group(ranks=list(ranks), backend="nccl") + for ranks in unique_groups + } + + +def _rank_group( + rank: int, + groups: Sequence[tuple[int, ...]], + handles: dict[tuple[int, ...], dist.ProcessGroup], +) -> tuple[tuple[int, ...], dist.ProcessGroup]: + group = next(group for group in groups if rank in group) + return group, handles[group] + + +def _measure_collective( + operation: Callable[[], None], + *, + warmup: int, + iterations: int, + trials: int, + device: torch.device, +) -> list[float]: + for _ in range(warmup): + operation() + torch.cuda.synchronize(device) + dist.barrier() + + samples = [] + for _ in range(trials): + dist.barrier() + torch.cuda.synchronize(device) + started = time.perf_counter() + for _ in range(iterations): + operation() + torch.cuda.synchronize(device) + elapsed_ms = (time.perf_counter() - started) * 1000 / iterations + worst_rank = torch.tensor(elapsed_ms, dtype=torch.float64, device=device) + dist.all_reduce(worst_rank, op=dist.ReduceOp.MAX) + if dist.get_rank() == 0: + samples.append(float(worst_rank.item())) + return samples + + +def _benchmark_worker( + rank: int, + world_size: int, + port: int, + cases: tuple[BenchmarkCase, ...], + dlo_payload_bytes: int, + sp_payload_bytes: int, + warmup: int, + iterations: int, + trials: int, + topology_text: str, + output: str, + markdown_output: str | None, +) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + "nccl", + rank=rank, + world_size=world_size, + device_id=device, + ) + results: list[dict[str, object]] = [] + try: + for case in cases: + kind, dlo_groups, sp_groups = dlo_groups_for_plan( + case.device_order, case.dp_size, case.sp_size + ) + handles = _make_groups(dlo_groups, sp_groups) + dlo_ranks, dlo_handle = _rank_group(rank, dlo_groups, handles) + sp_ranks, sp_handle = _rank_group(rank, sp_groups, handles) + + local_dlo_bytes = dlo_payload_bytes // len(dlo_ranks) + if local_dlo_bytes * len(dlo_ranks) != dlo_payload_bytes: + raise ValueError( + "DLO payload bytes must be divisible by DLO group size" + ) + if sp_payload_bytes % len(sp_ranks): + raise ValueError("SP payload bytes must be divisible by SP group size") + + dlo_input = torch.full( + (local_dlo_bytes,), rank, dtype=torch.uint8, device=device + ) + dlo_output = torch.empty( + (dlo_payload_bytes,), dtype=torch.uint8, device=device + ) + sp_input = torch.full( + (sp_payload_bytes,), rank, dtype=torch.uint8, device=device + ) + sp_output = torch.empty_like(sp_input) + + dlo_all_gather = partial( + dist.all_gather_into_tensor, + dlo_output, + dlo_input, + group=dlo_handle, + ) + sp_all_to_all = partial( + dist.all_to_all_single, + sp_output, + sp_input, + group=sp_handle, + ) + + dlo_samples = _measure_collective( + dlo_all_gather, + warmup=warmup, + iterations=iterations, + trials=trials, + device=device, + ) + sp_samples = _measure_collective( + sp_all_to_all, + warmup=warmup, + iterations=iterations, + trials=trials, + device=device, + ) + if rank == 0: + results.append( + { + "name": case.name, + "dp_size": case.dp_size, + "sp_size": case.sp_size, + "device_order": list(case.device_order), + "dlo_group_kind": kind, + "dlo_groups": [list(group) for group in dlo_groups], + "sp_groups": [list(group) for group in sp_groups], + "topology_label_score": case.topology_label_score, + "dlo_all_gather": summarize(dlo_samples), + "sp_all_to_all": summarize(sp_samples), + } + ) + dist.barrier() + + if rank == 0: + properties = torch.cuda.get_device_properties(device) + payload = { + "metadata": { + "timestamp_utc": datetime.now(UTC).isoformat(), + "gpu_name": properties.name, + "gpu_count": world_size, + "compute_capability": f"{properties.major}.{properties.minor}", + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "nccl_version": ".".join( + str(value) for value in torch.cuda.nccl.version() + ), + "dtype": "uint8-int8-transfer-representative", + "dlo_full_payload_bytes": dlo_payload_bytes, + "sp_per_rank_payload_bytes": sp_payload_bytes, + "warmup": warmup, + "iterations": iterations, + "trials": trials, + "timing_scope": "worst-rank steady-state collective", + "topology": topology_text, + }, + "results": results, + } + output_path = Path(output) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + if markdown_output is not None: + markdown_path = Path(markdown_output) + markdown_path.parent.mkdir(parents=True, exist_ok=True) + markdown_path.write_text(render_markdown(payload), encoding="utf-8") + finally: + dist.destroy_process_group() + + +def render_markdown(payload: dict[str, object]) -> str: + metadata = payload["metadata"] + results = payload["results"] + assert isinstance(metadata, dict) + assert isinstance(results, list) + lines = [ + "# Topology-aware DLO DP/SP benchmark", + "", + f"- GPU: {metadata['gpu_count']}x {metadata['gpu_name']}", + f"- Compute capability: {metadata['compute_capability']}", + ( + f"- PyTorch / CUDA / NCCL: {metadata['torch_version']} / " + f"{metadata['cuda_version']} / {metadata['nccl_version']}" + ), + f"- DLO full INT8 payload: {metadata['dlo_full_payload_bytes']} bytes", + f"- SP payload per rank: {metadata['sp_per_rank_payload_bytes']} bytes", + f"- Timing: {metadata['trials']} trials, worst rank per trial", + "", + ( + "| Candidate | DP x SP | Device order | DLO group | DLO median (ms) | " + "DLO p99 (ms) | SP median (ms) | SP p99 (ms) |" + ), + "|---|---:|---|---|---:|---:|---:|---:|", + ] + for result in results: + assert isinstance(result, dict) + dlo = result["dlo_all_gather"] + sp = result["sp_all_to_all"] + assert isinstance(dlo, dict) + assert isinstance(sp, dict) + lines.append( + f"| {result['name']} | {result['dp_size']}x{result['sp_size']} | " + f"{','.join(str(value) for value in result['device_order'])} | " + f"{result['dlo_group_kind']} | {dlo['median_ms']:.4f} | " + f"{dlo['p99_ms']:.4f} | {sp['median_ms']:.4f} | {sp['p99_ms']:.4f} |" + ) + lines.append("") + return "\n".join(lines) + + +def build_cases( + topology_text: str, dp_sizes: Sequence[int] +) -> tuple[BenchmarkCase, ...]: + topology = parse_nvidia_topology(topology_text) + natural = tuple(sorted(topology.devices)) + cases = [] + for dp_size in dp_sizes: + if len(natural) % dp_size: + raise click.BadParameter( + f"DP size {dp_size} does not divide {len(natural)} GPUs" + ) + sp_size = len(natural) // dp_size + optimized, optimized_score = optimize_device_order(topology, dp_size, sp_size) + candidates = (("natural", natural), ("topology", optimized)) + seen = set() + for name, order in candidates: + if order in seen: + continue + seen.add(order) + cases.append( + BenchmarkCase( + name=f"dp{dp_size}-sp{sp_size}-{name}", + dp_size=dp_size, + sp_size=sp_size, + device_order=order, + topology_label_score=( + optimized_score + if order == optimized + else topology_score(topology, order, dp_size, sp_size) + ), + ) + ) + return tuple(cases) + + +def parse_dp_sizes(value: str) -> tuple[int, ...]: + try: + sizes = tuple(dict.fromkeys(int(item.strip()) for item in value.split(","))) + except ValueError as exc: + raise click.BadParameter("DP sizes must be comma-separated integers") from exc + if not sizes or any(size <= 0 for size in sizes): + raise click.BadParameter("DP sizes must be positive") + return sizes + + +@click.command(help=__doc__) +@click.option("--output", type=click.Path(path_type=Path), required=True) +@click.option("--markdown-output", type=click.Path(path_type=Path)) +@click.option("--dp-sizes", default="1,2,4", show_default=True) +@click.option( + "--dlo-payload-mib", type=click.IntRange(min=1), default=256, show_default=True +) +@click.option( + "--sp-payload-mib", type=click.IntRange(min=1), default=64, show_default=True +) +@click.option("--warmup", type=click.IntRange(min=1), default=5, show_default=True) +@click.option("--iterations", type=click.IntRange(min=1), default=20, show_default=True) +@click.option("--trials", type=click.IntRange(min=3), default=7, show_default=True) +def main( + output: Path, + markdown_output: Path | None, + dp_sizes: str, + dlo_payload_mib: int, + sp_payload_mib: int, + warmup: int, + iterations: int, + trials: int, +) -> None: + if not torch.cuda.is_available(): + raise click.ClickException("CUDA is required") + topology_text = subprocess.run( + ["nvidia-smi", "topo", "-m"], + check=True, + capture_output=True, + text=True, + ).stdout + cases = build_cases(topology_text, parse_dp_sizes(dp_sizes)) + world_size = torch.cuda.device_count() + if world_size != len(parse_nvidia_topology(topology_text).devices): + raise click.ClickException("visible CUDA devices and topology matrix disagree") + mp.spawn( + _benchmark_worker, + args=( + world_size, + _find_free_port(), + cases, + dlo_payload_mib * 1024 * 1024, + sp_payload_mib * 1024 * 1024, + warmup, + iterations, + trials, + topology_text, + str(output), + str(markdown_output) if markdown_output is not None else None, + ), + nprocs=world_size, + join=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/plan_dlo_topology.py b/scripts/tools/plan_dlo_topology.py new file mode 100644 index 00000000..de2854a6 --- /dev/null +++ b/scripts/tools/plan_dlo_topology.py @@ -0,0 +1,114 @@ +"""Plan topology-aware DP/SP placement for distributed layerwise offload.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from typing import Any + +import click + +from astrai.parallel.topology import ( + DLOMeasurement, + parse_nvidia_topology, + select_dlo_plan, +) + + +def read_topology(path: Path | None) -> str: + if path is not None: + return path.read_text(encoding="utf-8") + result = subprocess.run( + ["nvidia-smi", "topo", "-m"], + check=True, + capture_output=True, + text=True, + ) + return result.stdout + + +def read_measurements(path: Path | None) -> tuple[DLOMeasurement, ...]: + if path is None: + return () + payload: Any = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict) or not isinstance(payload.get("results"), list): + raise click.ClickException("measurement JSON must contain a results list") + measurements = [] + for item in payload["results"]: + try: + measurements.append( + DLOMeasurement( + dp_size=int(item["dp_size"]), + sp_size=int(item["sp_size"]), + device_order=tuple(int(value) for value in item["device_order"]), + dlo_all_gather_ms=float(item["dlo_all_gather"]["median_ms"]), + sp_collective_ms=float(item["sp_all_to_all"]["median_ms"]), + ) + ) + except (KeyError, TypeError, ValueError) as exc: + raise click.ClickException(f"invalid measurement result: {item!r}") from exc + return tuple(measurements) + + +@click.command(help=__doc__) +@click.option( + "--topology-file", + type=click.Path(path_type=Path, exists=True, readable=True), + help="Read captured nvidia-smi topo -m output instead of probing locally.", +) +@click.option( + "--measurements", + type=click.Path(path_type=Path, exists=True, readable=True), + help="Prefer collective timings from benchmark_dlo_topology.py.", +) +@click.option( + "--concurrent-requests", type=click.IntRange(min=1), default=1, show_default=True +) +@click.option("--dp-size", type=click.IntRange(min=1)) +@click.option("--sp-size", type=click.IntRange(min=1)) +@click.option("--output", type=click.Path(path_type=Path)) +def main( + topology_file: Path | None, + measurements: Path | None, + concurrent_requests: int, + dp_size: int | None, + sp_size: int | None, + output: Path | None, +) -> None: + if (dp_size is None) != (sp_size is None): + raise click.UsageError("--dp-size and --sp-size must be provided together") + topology = parse_nvidia_topology(read_topology(topology_file)) + plan = select_dlo_plan( + topology, + concurrent_requests=concurrent_requests, + dp_size=dp_size, + sp_size=sp_size, + measurements=read_measurements(measurements), + ) + payload = plan.as_dict() + order = ",".join(str(device) for device in plan.device_order) + payload["launch"] = { + "astrai_env": {"ASTRAI_DEVICE_ORDER": order}, + "vllm_omni_env": {"CUDA_VISIBLE_DEVICES": order}, + "vllm_omni_args": [ + "--num-gpus", + str(len(plan.device_order)), + "--data-parallel-size", + str(plan.dp_size), + "--usp", + str(plan.sp_size), + "--ring", + "1", + "--enable-distributed-layerwise-offload", + ], + } + rendered = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if output is not None: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(rendered, encoding="utf-8") + click.echo(rendered, nl=False) + + +if __name__ == "__main__": + main() diff --git a/tests/parallel/test_infraswe_draft.py b/tests/parallel/test_infraswe_draft.py new file mode 100644 index 00000000..b75a4a19 --- /dev/null +++ b/tests/parallel/test_infraswe_draft.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +PROJECT_ROOT = Path(__file__).parents[2] +DRAFT_PATH = PROJECT_ROOT / "benchmarks/infraswe/astrai-topology-dlo-draft.json" + + +def _digest_bytes(payload: bytes) -> str: + return "sha256:" + hashlib.sha256(payload).hexdigest() + + +def _digest_file(path: str) -> str: + return _digest_bytes((PROJECT_ROOT / path).read_bytes()) + + +def _digest_file_list(paths: tuple[str, ...]) -> str: + manifest = "".join( + f"{hashlib.sha256((PROJECT_ROOT / path).read_bytes()).hexdigest()} {path}\n" + for path in paths + ) + return _digest_bytes(manifest.encode()) + + +def test_infraswe_draft_local_material_digests_are_current() -> None: + draft = json.loads(DRAFT_PATH.read_text(encoding="utf-8")) + assert draft["target"]["repository"] == "https://github.com/ViperEkura/AstrAI" + assert draft["target"]["revision"] == _digest_bytes( + b"1fad50d8476abe7d4b4cb527eb6e174c511fd409" + ) + assert draft["baseline"]["revision"] == draft["target"]["revision"] + + candidate_files = ( + "astrai/parallel/topology.py", + "astrai/parallel/setup.py", + "astrai/parallel/executor.py", + "astrai/parallel/__init__.py", + "scripts/tools/plan_dlo_topology.py", + "scripts/tools/benchmark_dlo_topology.py", + "tests/parallel/test_topology.py", + "tests/parallel/test_topology_cli.py", + "tests/parallel/test_parallel.py", + "docs/guides/distributed.md", + ) + assert draft["candidate"]["revision"] == _digest_file_list(candidate_files) + + project_profile_files = ( + "pyproject.toml", + "README.md", + "astrai/parallel/__init__.py", + "docs/guides/distributed.md", + ) + assert draft["target"]["project_profile_sha256"] == _digest_file_list( + project_profile_files + ) + + workload_path = "benchmarks/results/dlo_topology_8x5060ti.json" + assert draft["deployment"]["workload_portfolio"]["sha256"] == _digest_file( + workload_path + ) + assert draft["deployment"]["request_or_step_protocol"]["sha256"] == ( + _digest_file("docs/guides/distributed.md") + ) + + acceptance_files = ( + "tests/parallel/test_topology.py", + "tests/parallel/test_topology_cli.py", + "tests/parallel/test_parallel.py", + workload_path, + ) + assert draft["acceptance_contract"]["sha256"] == _digest_file_list(acceptance_files) + probe_files = ( + "tests/parallel/test_topology.py", + "tests/parallel/test_topology_cli.py", + ) + assert draft["acceptance_contract"]["probe_set_sha256"] == _digest_file_list( + probe_files + ) + + profile_files = ( + "astrai/parallel/topology.py", + "docs/guides/distributed.md", + "docs/benchmarks/topology_aware_dlo_8x5060ti.md", + ) + assert draft["project_objectives"]["profile_set_sha256"] == _digest_file_list( + profile_files + ) diff --git a/tests/parallel/test_parallel.py b/tests/parallel/test_parallel.py index c05b0848..6c0a3e34 100644 --- a/tests/parallel/test_parallel.py +++ b/tests/parallel/test_parallel.py @@ -1,7 +1,17 @@ +import os + +import pytest import torch import torch.distributed as dist -from astrai.parallel import get_rank, only_on_rank, spawn_parallel_fn +from astrai.parallel import ( + DDPExecutor, + get_rank, + only_on_rank, + resolve_local_device_index, + setup_parallel, + spawn_parallel_fn, +) @only_on_rank(0) @@ -30,3 +40,76 @@ def test_spawn_only_on_rank(): def test_spawn_all_reduce(): spawn_parallel_fn(all_reduce, world_size=2, backend="gloo") + + +def test_device_order_maps_accelerators_but_not_cpu(monkeypatch): + monkeypatch.setenv("ASTRAI_DEVICE_ORDER", "2,0,3,1") + assert resolve_local_device_index(1, 4, "cuda") == 0 + assert resolve_local_device_index(1, 4, "xpu") == 0 + assert resolve_local_device_index(1, 4, "cpu") == 1 + + +def test_device_order_fails_closed(monkeypatch): + monkeypatch.setenv("ASTRAI_DEVICE_ORDER", "0,0") + with pytest.raises(ValueError, match="permutation"): + resolve_local_device_index(0, 2, "cuda") + + +def test_device_order_uses_local_world_size_and_validates_rank(monkeypatch): + monkeypatch.setenv("ASTRAI_DEVICE_ORDER", "2,0,3,1") + assert resolve_local_device_index(3, 4, "cuda") == 1 + + with pytest.raises(ValueError, match="outside local world size"): + resolve_local_device_index(4, 4, "cuda") + + +def test_setup_parallel_uses_torchrun_local_world_size(monkeypatch): + state = {"initialized": False} + captured = {} + + def fake_init_process_group(**kwargs): + captured.update(kwargs) + state["initialized"] = True + + def fake_destroy_process_group(): + state["initialized"] = False + + monkeypatch.setenv("LOCAL_WORLD_SIZE", "4") + monkeypatch.setenv("ASTRAI_DEVICE_ORDER", "2,0,3,1") + monkeypatch.setenv("MASTER_ADDR", "previous") + monkeypatch.setenv("MASTER_PORT", "previous") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_DEVICE", "cpu") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", fake_init_process_group) + monkeypatch.setattr(dist, "destroy_process_group", fake_destroy_process_group) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + with setup_parallel(rank=4, world_size=8, local_rank=0, device_type="cuda"): + assert captured["device_id"] == torch.device("cuda", 2) + assert captured["rank"] == 4 + assert captured["world_size"] == 8 + assert os.environ["LOCAL_RANK"] == "0" + assert os.environ["LOCAL_DEVICE"] == "cuda:2" + + +def test_ddp_uses_mapped_local_device_without_changing_logical_rank(monkeypatch): + captured = {} + sentinel = object() + + def fake_ddp(model, **kwargs): + captured.update(kwargs) + return sentinel + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setenv("LOCAL_DEVICE", "cuda:2") + monkeypatch.setattr("astrai.parallel.executor.DDP", fake_ddp) + + wrapped = DDPExecutor()._prepare_model(torch.nn.Linear(1, 1)) + + assert wrapped is sentinel + assert captured["device_ids"] == [2] + assert captured["output_device"] == 2 + assert os.environ["LOCAL_RANK"] == "0" diff --git a/tests/parallel/test_topology.py b/tests/parallel/test_topology.py new file mode 100644 index 00000000..49a98941 --- /dev/null +++ b/tests/parallel/test_topology.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import pytest + +from astrai.parallel.topology import ( + DLOMeasurement, + build_parallel_groups, + parse_device_order, + parse_nvidia_topology, + select_dlo_plan, +) + +TOPOLOGY_8X_5060_TI = """ + GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7 CPU Affinity NUMA Affinity +GPU0 X PHB NODE NODE NODE NODE NODE NODE 0-63 0 +GPU1 PHB X NODE NODE NODE NODE NODE NODE 0-63 0 +GPU2 NODE NODE X PHB NODE NODE NODE NODE 0-63 0 +GPU3 NODE NODE PHB X NODE NODE NODE NODE 0-63 0 +GPU4 NODE NODE NODE NODE X PHB NODE NODE 0-63 0 +GPU5 NODE NODE NODE NODE PHB X NODE NODE 0-63 0 +GPU6 NODE NODE NODE NODE NODE NODE X PHB 0-63 0 +GPU7 NODE NODE NODE NODE NODE NODE PHB X 0-63 0 +Legend: + X = Self + PHB = Connection traversing PCIe as well as a PCIe Host Bridge +""" + + +def test_parse_nvidia_topology_and_reject_partial_matrix() -> None: + topology = parse_nvidia_topology("\x1b[0m" + TOPOLOGY_8X_5060_TI) + assert topology.devices == tuple(range(8)) + assert topology.link(0, 1) == "PHB" + assert topology.link(0, 2) == "NODE" + + with pytest.raises(ValueError, match="row mismatch"): + parse_nvidia_topology( + TOPOLOGY_8X_5060_TI.replace("GPU7 NODE", "CPU7 NODE") + ) + + +def test_parse_single_gpu_topology() -> None: + topology = parse_nvidia_topology("GPU0 CPU Affinity NUMA Affinity\nGPU0 X 0-31 0\n") + assert topology.devices == (0,) + assert topology.link(0, 0) == "X" + + +def test_sp_fastest_group_layout_matches_dlo_dp_precedence() -> None: + order = (0, 2, 4, 6, 1, 3, 5, 7) + dp_groups, sp_groups = build_parallel_groups(order, dp_size=2, sp_size=4) + assert dp_groups == ((0, 1), (2, 3), (4, 5), (6, 7)) + assert sp_groups == ((0, 2, 4, 6), (1, 3, 5, 7)) + + plan = select_dlo_plan( + parse_nvidia_topology(TOPOLOGY_8X_5060_TI), + concurrent_requests=2, + ) + assert plan.dlo_group_kind == "dp" + assert plan.dp_size == 2 + assert plan.sp_size == 4 + assert plan.dlo_groups == dp_groups + + +def test_single_request_is_dp1_sp8_and_uses_sp_for_dlo() -> None: + plan = select_dlo_plan( + parse_nvidia_topology(TOPOLOGY_8X_5060_TI), concurrent_requests=1 + ) + assert (plan.dp_size, plan.sp_size) == (1, 8) + assert plan.dlo_group_kind == "sp" + assert plan.dlo_groups == (tuple(range(8)),) + + with pytest.raises(ValueError, match="concurrent request"): + select_dlo_plan( + parse_nvidia_topology(TOPOLOGY_8X_5060_TI), + concurrent_requests=1, + dp_size=2, + sp_size=4, + ) + + +def test_measurement_overrides_topology_label_heuristic() -> None: + topology = parse_nvidia_topology(TOPOLOGY_8X_5060_TI) + heuristic = select_dlo_plan(topology, concurrent_requests=2) + natural = tuple(range(8)) + measured = select_dlo_plan( + topology, + concurrent_requests=2, + measurements=( + DLOMeasurement(2, 4, heuristic.device_order, 4.0, 1.0), + DLOMeasurement(2, 4, natural, 1.0, 1.0), + ), + ) + assert heuristic.device_order != natural + assert measured.device_order == natural + assert measured.selection_source == "measured-collectives" + assert measured.score == 2.0 + + +@pytest.mark.parametrize( + ("value", "local_world_size"), + [("0,2,1,3", 4), ("0", 1)], +) +def test_parse_device_order(value: str, local_world_size: int) -> None: + assert parse_device_order(value, local_world_size) == tuple( + int(item) for item in value.split(",") + ) + + +@pytest.mark.parametrize("value", ["0,1,1,3", "0,1", "0,-1,2,3", "0,1,2,4", "0,x,2,3"]) +def test_parse_device_order_rejects_invalid_mapping(value: str) -> None: + with pytest.raises(ValueError): + parse_device_order(value, 4) diff --git a/tests/parallel/test_topology_cli.py b/tests/parallel/test_topology_cli.py new file mode 100644 index 00000000..51f39cff --- /dev/null +++ b/tests/parallel/test_topology_cli.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import json + +from click.testing import CliRunner + +from scripts.tools.plan_dlo_topology import main +from tests.parallel.test_topology import TOPOLOGY_8X_5060_TI + + +def test_plan_cli_emits_astrai_and_distributed_dlo_launch_contract(tmp_path) -> None: + topology_path = tmp_path / "topology.txt" + topology_path.write_text(TOPOLOGY_8X_5060_TI, encoding="utf-8") + + result = CliRunner().invoke( + main, + ["--topology-file", str(topology_path), "--concurrent-requests", "1"], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + assert payload["dlo_group_kind"] == "sp" + assert payload["launch"]["astrai_env"]["ASTRAI_DEVICE_ORDER"] == ("0,1,2,3,4,5,6,7") + assert ( + "--enable-distributed-layerwise-offload" in payload["launch"]["vllm_omni_args"] + ) + assert payload["launch"]["vllm_omni_args"][:8] == [ + "--num-gpus", + "8", + "--data-parallel-size", + "1", + "--usp", + "8", + "--ring", + "1", + ]