diff --git a/vidur-alibabacloud/tests/test_pd_network_config_sweep.py b/vidur-alibabacloud/tests/test_pd_network_config_sweep.py new file mode 100644 index 000000000..48553bccb --- /dev/null +++ b/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) diff --git a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/__init__.py b/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/__init__.py index fed3ddd3a..c53f8c556 100644 --- a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/__init__.py +++ b/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, ] diff --git a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/config.py b/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/config.py index 4a073fedd..75b2396ea 100644 --- a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/config.py +++ b/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 diff --git a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/debug_config.yml b/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/debug_config.yml index f3dc54b97..e60da5cb6 100644 --- a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config/debug_config.yml +++ b/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 diff --git a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config_explorer.py b/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config_explorer.py index e3edf60f0..5c0f7c21f 100644 --- a/vidur-alibabacloud/vidur/config_optimizer/config_explorer/config_explorer.py +++ b/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 = (