Add P/D network dimension to Vidur config explorer

vidur-alibabacloud/tests/test_pd_network_config_sweep.py

@@ -0,0 +1,118 @@
+"""Tests for P/D network design-space generation."""
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from vidur.config_optimizer.config_explorer.config import JobConfig
+
+
+def _base_config():
+    return {
+        "models": [{"name": "toy", "identifier": "toy/model"}],
+        "traces": [
+            {
+                "name": "chat",
+                "trace_file": "chat.csv",
+                "max_seq_len": 4096,
+                "num_requests": 10,
+                "start_qps": 1,
+            }
+        ],
+        "clusters": [{"device": "h100", "num_gpus": 8, "gpus_per_node": 8}],
+        "schedulers": [{"scheduler": "vllm"}],
+        "tp_dimensions": [1],
+        "pp_dimensions": [1],
+        "batch_sizes": [32],
+    }
+
+
+def test_pd_network_dimension_expands_jobs_and_replica_config():
+    config = _base_config()
+    config["pd_networks"] = [
+        {
+            "name": "pd",
+            "pd_node_ratio": 0.5,
+            "pd_p2p_comm_bandwidth": 100,
+            "rdma_bandwidth": 200,
+            "nvlink_bandwidth": 900,
+            "pd_p2p_comm_dtype": "float16",
+        },
+        {
+            "name": "pd",
+            "pd_node_ratio": 0.5,
+            "pd_p2p_comm_bandwidth": 100,
+            "rdma_bandwidth": 200,
+            "nvlink_bandwidth": 900,
+            "pd_p2p_comm_dtype": "fp8",
+        },
+    ]
+
+    jobs = JobConfig.generate_job_configs(config)
+
+    assert len(jobs) == 2
+    assert len({job.get_key() for job in jobs}) == 2
+    assert {
+        job.to_config_dict()["replica_config_pd_p2p_comm_dtype"] for job in jobs
+    } == {"float16", "fp8"}
+
+
+def test_legacy_config_preserves_job_key():
+    job = JobConfig.generate_job_configs(_base_config())[0]
+
+    assert job.get_key() == "toy_chat_tk4096_rq10_h100_vllm_tp1_pp1_bsz32"
+    assert job.to_config_dict()["replica_config_pd_node_ratio"] == 1.0
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"pd_node_ratio": 0},
+        {"pd_node_ratio": 1.1},
+        {"pd_p2p_comm_bandwidth": 0},
+        {"rdma_bandwidth": -1},
+        {"nvlink_bandwidth": 0},
+        {"pd_p2p_comm_dtype": "int8"},
+    ],
+)
+def test_invalid_pd_network_point_is_rejected(overrides):
+    config = _base_config()
+    network = {
+        "name": "invalid",
+        "pd_node_ratio": 0.5,
+        "pd_p2p_comm_bandwidth": 100,
+        "rdma_bandwidth": 200,
+        "nvlink_bandwidth": 900,
+        "pd_p2p_comm_dtype": "float16",
+    }
+    network.update(overrides)
+    config["pd_networks"] = [network]
+
+    with pytest.raises(ValueError):
+        JobConfig.generate_job_configs(config)
+
+
+def test_pd_split_with_zero_phase_replicas_is_skipped():
+    config = _base_config()
+    config["pd_networks"] = [{"name": "invalid-split", "pd_node_ratio": 0.01}]
+
+    assert JobConfig.generate_job_configs(config) == []
+
+
+def test_empty_or_duplicate_pd_network_dimension_is_rejected():
+    config = _base_config()
+    config["pd_networks"] = []
+    with pytest.raises(ValueError):
+        JobConfig.generate_job_configs(config)
+
+    point = {"name": "duplicate", "pd_node_ratio": 1}
+    config["pd_networks"] = [point, point]
+    with pytest.raises(ValueError):
+        JobConfig.generate_job_configs(config)
+
+    config["pd_networks"] = [{"pd_node_ratio": 0.5}]
+    with pytest.raises(ValueError):
+        JobConfig.generate_job_configs(config)

vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/__init__.py

@@ -2,6 +2,7 @@
     ClusterConfig,
     JobConfig,
     ModelConfig,
+    PDNetworkConfig,
     SchedulerConfig,
     SimulationConfig,
     TraceConfig,
@@ -14,4 +15,5 @@
     SchedulerConfig,
     ClusterConfig,
     ModelConfig,
+    PDNetworkConfig,
 ]

vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/config.py

@@ -1,6 +1,8 @@
 import hashlib
+import math
 from dataclasses import dataclass
 from itertools import product
+from numbers import Real
 from typing import List, Optional
 
 
@@ -94,7 +96,106 @@ def to_config_dict(self):
         }
 
 
+@dataclass
+class PDNetworkConfig:
+    name: str = "mixed"
+    pd_node_ratio: float = 1.0
+    pd_p2p_comm_bandwidth: int = 800
+    rdma_bandwidth: int = 800
+    nvlink_bandwidth: int = 1600
+    pd_p2p_comm_dtype: str = "float16"
+
+    def get_key(self):
+        pd_node_ratio = repr(float(self.pd_node_ratio))
+        return (
+            f"{self.name}_pdr{pd_node_ratio}"
+            f"_pdbw{self.pd_p2p_comm_bandwidth}"
+            f"_rdbw{self.rdma_bandwidth}"
+            f"_nvbw{self.nvlink_bandwidth}"
+            f"_dtype{self.pd_p2p_comm_dtype}"
+        )
+
+    def validate_values(self):
+        if not isinstance(self.name, str) or not self.name.strip():
+            raise ValueError("P/D network point name must not be empty")
+        if (
+            not isinstance(self.pd_node_ratio, Real)
+            or isinstance(self.pd_node_ratio, bool)
+            or not math.isfinite(self.pd_node_ratio)
+            or not 0 < self.pd_node_ratio <= 1
+        ):
+            raise ValueError(
+                f"Invalid pd_node_ratio {self.pd_node_ratio}; expected 0 < ratio <= 1"
+            )
+        bandwidths = (
+            self.pd_p2p_comm_bandwidth,
+            self.rdma_bandwidth,
+            self.nvlink_bandwidth,
+        )
+        if any(
+            not isinstance(bandwidth, int)
+            or isinstance(bandwidth, bool)
+            or bandwidth <= 0
+            for bandwidth in bandwidths
+        ):
+            raise ValueError("P/D, RDMA, and NVLink bandwidths must be positive")
+        if self.pd_p2p_comm_dtype not in ("fp8", "float16", "float32"):
+            raise ValueError(
+                f"Invalid P/D transfer dtype {self.pd_p2p_comm_dtype!r}"
+            )
+
+    def validate(self, num_replicas: int):
+        self.validate_values()
+        if self.pd_node_ratio < 1:
+            num_prefill_replicas = int(num_replicas * self.pd_node_ratio)
+            num_decode_replicas = num_replicas - num_prefill_replicas
+            if num_prefill_replicas == 0 or num_decode_replicas == 0:
+                raise ValueError(
+                    f"P/D network point {self.name!r} produces "
+                    f"{num_prefill_replicas} P and {num_decode_replicas} D replicas"
+                )
+
+    def is_valid(self, num_replicas: int):
+        try:
+            self.validate(num_replicas)
+        except (TypeError, ValueError):
+            return False
+        return True
+
+    def to_config_dict(self):
+        return {
+            "replica_config_pd_node_ratio": self.pd_node_ratio,
+            "replica_config_pd_p2p_comm_bandwidth": self.pd_p2p_comm_bandwidth,
+            "replica_config_rdma_bandwidth": self.rdma_bandwidth,
+            "replica_config_nvlink_bandwidth": self.nvlink_bandwidth,
+            "replica_config_pd_p2p_comm_dtype": self.pd_p2p_comm_dtype,
+        }
+
+
 class JobConfig:
+    @staticmethod
+    def _get_pd_network_configs(config: dict):
+        pd_networks = config.get("pd_networks")
+        if pd_networks is None:
+            return [None]
+        if not pd_networks:
+            raise ValueError("pd_networks must contain at least one point")
+
+        if any(
+            not isinstance(pd_network, dict) or "name" not in pd_network
+            for pd_network in pd_networks
+        ):
+            raise ValueError("Each P/D network point must have a name")
+        pd_network_configs = [
+            PDNetworkConfig(**pd_network) for pd_network in pd_networks
+        ]
+        for pd_network in pd_network_configs:
+            pd_network.validate_values()
+        keys = [pd_network.get_key() for pd_network in pd_network_configs]
+        if len(keys) != len(set(keys)):
+            raise ValueError("pd_networks contains duplicate points")
+        return pd_network_configs
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -104,6 +205,7 @@ def __init__(
         num_tensor_parallel_workers: int,
         num_pipeline_stages: int,
         batch_size: int,
+        pd_network_config: Optional[PDNetworkConfig] = None,
     ):
         self.model_config = model_config
         self.trace_config = trace_config
@@ -114,29 +216,36 @@ def __init__(
         self.num_workers = self.num_tensor_parallel_workers * self.num_pipeline_stages
         self.batch_size = batch_size * num_pipeline_stages
         self.num_replicas = self.cluster_config.num_gpus // self.num_workers
+        self.pd_network_config = pd_network_config or PDNetworkConfig()
+        self._has_pd_network_dimension = pd_network_config is not None
 
         self.start_qps = self.trace_config.start_qps
 
     def is_valid(self):
-        return (
+        is_valid = (
             self.num_replicas > 0
             and self.model_config.is_tensor_parallel_degree_valid(
                 self.num_tensor_parallel_workers
             )
             and self.num_tensor_parallel_workers <= self.cluster_config.gpus_per_node
         )
+        return is_valid and self.pd_network_config.is_valid(self.num_replicas)
 
     def get_key(self):
-        return (
+        key = (
             f"{self.model_config.name}_{self.trace_config.get_key()}_{self.cluster_config.get_key()}_{self.scheduler_config.get_key()}"
             f"_tp{self.num_tensor_parallel_workers}_pp{self.num_pipeline_stages}_bsz{self.batch_size}"
         )
+        if self._has_pd_network_dimension:
+            key += f"_{self.pd_network_config.get_key()}"
+        return key
 
     def get_human_readable_name(self):
         return (
             f"Model: {self.model_config.name}, Trace: {self.trace_config.name}, Cluster: {self.cluster_config.device}, "
             f"Scheduler: {self.scheduler_config.scheduler}, TP: {self.num_tensor_parallel_workers}, "
-            f"PP: {self.num_pipeline_stages}, BSZ: {self.batch_size}, CS: {self.scheduler_config.chunk_size}, Hash: {self.get_hash()}"
+            f"PP: {self.num_pipeline_stages}, BSZ: {self.batch_size}, CS: {self.scheduler_config.chunk_size}, "
+            f"PD Network: {self.pd_network_config.name}, Hash: {self.get_hash()}"
         )
 
     def get_hash(self):
@@ -148,6 +257,7 @@ def to_config_dict(self):
             **self.trace_config.to_config_dict(),
             **self.cluster_config.to_config_dict(),
             **self.scheduler_config.to_config_dict(),
+            **self.pd_network_config.to_config_dict(),
             "replica_config_tensor_parallel_size": self.num_tensor_parallel_workers,
             "replica_config_num_pipeline_stages": self.num_pipeline_stages,
             "vllm_scheduler_config_batch_size_cap": self.batch_size,
@@ -169,6 +279,7 @@ def generate_job_configs(cls, config: dict):
             tp_dimension,
             pp_dimension,
             batch_size,
+            pd_network_config,
         ) in product(
             config["models"],
             config["traces"],
@@ -177,6 +288,7 @@ def generate_job_configs(cls, config: dict):
             config["tp_dimensions"],
             config["pp_dimensions"],
             config["batch_sizes"],
+            cls._get_pd_network_configs(config),
         ):
             job_config = cls(
                 ModelConfig(**model_config),
@@ -186,6 +298,7 @@ def generate_job_configs(cls, config: dict):
                 tp_dimension,
                 pp_dimension,
                 batch_size,
+                pd_network_config,
             )
             if not job_config.is_valid():
                 continue
@@ -204,6 +317,7 @@ def generate_unique_model_job_configs(cls, config: dict, num_requests: int = 32)
         batch_size = config["batch_sizes"][0]
         # set pp_dimensions to 2 because it covers all the options
         pp_dimensions = [2]
+        pd_network_config = cls._get_pd_network_configs(config)[0]
 
         for model_config, cluster_config, tp_dimension, pp_dimension in product(
             config["models"],
@@ -219,6 +333,7 @@ def generate_unique_model_job_configs(cls, config: dict, num_requests: int = 32)
                 tp_dimension,
                 pp_dimension,
                 batch_size,
+                pd_network_config,
             )
             if not job_config.is_valid():
                 continue

vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/debug_config.yml

@@ -29,6 +29,15 @@ batch_sizes: [32, 64]
 tp_dimensions: [1, 2]
 pp_dimensions: [1, 2]
 
+# Optional P/D deployment and network sweep (可选的 P/D 部署和网络扫描):
+# pd_networks:
+#   - name: pd-balanced-100g
+#     pd_node_ratio: 0.5
+#     pd_p2p_comm_bandwidth: 100
+#     rdma_bandwidth: 100
+#     nvlink_bandwidth: 900
+#     pd_p2p_comm_dtype: float16
+
 models:
   - name: phi-2
     identifier: microsoft/phi-2

vidur-alibabacloud/vidur/config_optimizer/config_explorer/config_explorer.py

@@ -55,11 +55,11 @@ def _warmup_cache(self):
             ), "All nodes should have the same result"
 
     def run(self):
+        job_configs = JobConfig.generate_job_configs(self.config)
+
         if not self.args.skip_cache_warmup:
             self._warmup_cache()
 
-        job_configs = JobConfig.generate_job_configs(self.config)
-
         ray_parallel_runner = RayParallelRunner()
 
         remote_func = (