diff --git a/fern/versions/latest/pages/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx index b086c60ee..ca37c3074 100644 --- a/fern/versions/latest/pages/slurm/benchmarks.mdx +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -2,16 +2,50 @@ Use a benchmark to compare concurrency and deployment topology while keeping each case as an ordinary Data Designer Slurm run. No benchmark controller stays resident after submission. -Complete [Getting Started](/slurm/getting-started) first so the profile, images, and base run are ready. +Complete [Getting Started](/slurm/getting-started) first so the profile, serving image, and builder config are ready. + +## Prepare a compatible base run + +This example requires two nodes with eight GPUs per node. Confirm that the selected profile or its detected GPU count requests eight GPUs per node. Reuse `builder.yaml` and the registered `vllm` image from Getting Started, and select a model that supports TP8 and fits in eight GPUs. The TP1 run from Getting Started cannot span two nodes with the default PP1. + +Create `benchmark-run.yaml` next to `builder.yaml`: + +```yaml +schema_version: 1 +name: benchmark-base-run +builder: + source: builder.yaml +invocation: + num_records: 8 + dataset_name: greetings + model_concurrency: + generator: 8 +deployments: + - model_alias: generator + served_model_name: example/generator + model: example/generator + server: + type: vllm + image: + name: vllm + startup_timeout: 15m + distributed_init_timeout: 10m + topology: + tensor_parallel: 8 +submission: + time_limit: 01:00:00 +``` + +The repository ships this base run as `packages/data-designer-slurm/examples/benchmark-run.yaml`, alongside the matching `benchmark.yaml`. Keep PP omitted to use its default of one stage. For other GPU counts, choose a TP size that distributes evenly across both one-node and two-node groups and fits the GPUs available in each group. ## Define cases -Create `benchmark.yaml` next to an existing `run.yaml`: +Create `benchmark.yaml` next to `benchmark-run.yaml`: ```yaml schema_version: 1 name: generator-scaling -base_run: run.yaml +base_run: benchmark-run.yaml model_aliases: all concurrency_values: - 32 @@ -38,7 +72,9 @@ The compiler expands the authored order deterministically. For this example, the Each benchmark case runs as one Slurm array task, even when the base run uses multiple tasks. The record policy sets the sample size for that task. This measures the throughput of one allocation, which analysis uses to estimate how many such tasks are needed for `target_total_records` within `target_runtime`. -`nodes_per_replica` controls the number of nodes hosting each replica. When `pipeline_parallel` is omitted, it defaults to `nodes_per_replica`, and the resolved replica count is `(nodes / nodes_per_replica) * (gpus_per_node / tensor_parallel)`. With an explicit `pipeline_parallel`, each replica uses `tensor_parallel * pipeline_parallel` GPUs, distributed evenly across its nodes. Benchmark deployment cases change `nodes` and `nodes_per_replica`; they keep each deployment's `tensor_parallel` and explicit `pipeline_parallel` from the base run. Invalid combinations fail before submission. See [GPU topology and replica count](/slurm/run-configuration#gpu-topology-and-replica-count) for concrete eight-GPU examples. +`nodes_per_replica` controls the number of nodes hosting each replica. `pipeline_parallel` defaults to `1`, independently of the node count. Each replica uses `tensor_parallel * pipeline_parallel` GPUs, distributed evenly across its nodes; the resolved replica count is `(nodes * gpus_per_node) / (tensor_parallel * pipeline_parallel)`. Benchmark deployment cases change `nodes` and `nodes_per_replica`; they keep each deployment's TP and PP sizes from the base run. Changing the node count does not add pipeline stages. Each child plan is validated before its submission. An invalid later case can leave earlier cases submitted, so choose a base topology compatible with every case. See [GPU topology and replica count](/slurm/run-configuration#gpu-topology-and-replica-count) for concrete eight-GPU examples. + +Changing only `nodes_per_replica` changes placement, not GPUs per replica or replica count. For example, with TP8, default PP1, and eight GPUs per node, both cases above create two eight-GPU replicas. The first places each replica on one node; the second spreads each replica across both nodes with four GPUs per node. Use a different base run to compare different TP or PP sizes. ## Run diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index 18efc38b0..8b79a1e9d 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -91,7 +91,7 @@ For runs with local deployments, `gpus_per_node` is both the GPU count used for Multi-node generation jobs request exclusive nodes, including in `gres` mode. This keeps serving ports isolated because Slurm reports batch-job GPU IDs for the first node only. A partial-node `gpus_per_node` setting still limits the GPUs used per node, but other jobs cannot share those nodes during the run. -The profile decides how many GPUs the job requests **per node**. Each deployment's `tensor_parallel`, optional `pipeline_parallel`, and `nodes_per_replica` decide how those GPUs form vLLM replicas. For example, one node with `gpus_per_node: 8` and `tensor_parallel: 1` runs eight one-GPU replicas; it does not request only one GPU. See [GPU topology and replica count](/slurm/run-configuration#gpu-topology-and-replica-count). +The profile decides how many GPUs the job requests **per node**. Each deployment's `tensor_parallel`, `pipeline_parallel` (default: `1`), and `nodes_per_replica` decide how those GPUs form vLLM replicas. Increasing `nodes_per_replica` does not increase PP. For example, one node with `gpus_per_node: 8` and `tensor_parallel: 1` runs eight one-GPU replicas; it does not request only one GPU. See [GPU topology and replica count](/slurm/run-configuration#gpu-topology-and-replica-count). When selecting multiple partitions with `gpus_per_node: auto`, every selected partition must report GPU resources and all eligible nodes must report the same GPU count. Choose partitions that can also satisfy the run's GPU topology, model, and image requirements. diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 90180df47..62754c3d6 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -219,51 +219,88 @@ resources: topology: tensor_parallel: 8 nodes_per_replica: 2 + pipeline_parallel: 2 ``` -In an eight-GPU profile, this starts two replicas, each spanning two nodes with 16 GPUs total. When `pipeline_parallel` is omitted, it defaults to `nodes_per_replica`, preserving this pipeline-parallel layout. +In an eight-GPU profile, this starts two replicas, each spanning two nodes with 16 GPUs total. Each replica has two pipeline stages, each using TP8. TP communication stays within each node, and intermediate activations pass between the stages across nodes. ### GPU topology and replica count -`tensor_parallel` and `pipeline_parallel` specify vLLM's tensor and pipeline parallel sizes for one replica. `nodes_per_replica` specifies how many nodes host that replica. Omitting `pipeline_parallel` defaults it to `nodes_per_replica`, so existing configurations continue to use tensor parallelism within each node and pipeline parallelism across nodes. +One replica serves requests using a group of GPU workers. Choose how those workers cooperate: -To spread tensor parallelism across nodes while keeping PP=1, set `pipeline_parallel: 1` explicitly. For a profile requesting **four GPUs per node**, this configuration runs one TP16/PP1 replica across four nodes, with expert parallelism enabled: +- **Tensor parallelism (TP)** splits work within model layers across GPUs. A TP group can span multiple nodes. +- **Pipeline parallelism (PP)** splits model layers into sequential stages. A replica can use TP, PP, or both. +- **Independent replicas** serve different requests with separate copies of the model. The planner derives their count from the GPU budget and the topology of each replica. + +`tensor_parallel` and `pipeline_parallel` specify the TP and PP sizes for one replica. `nodes_per_replica` specifies how many nodes host that replica; it does not select the parallelism method. + + +`pipeline_parallel` defaults to `1`, regardless of `nodes_per_replica`. Spreading a replica across nodes does not add pipeline stages. Set `pipeline_parallel: 2` explicitly when you want two stages. TP and PP do not need to be equal; their product determines the GPU count per replica. + + +For a profile requesting **four GPUs per node**, this configuration runs one TP16 replica across four nodes. PP defaults to 1, and expert parallelism defaults to disabled: ```yaml resources: nodes: 4 topology: tensor_parallel: 16 - pipeline_parallel: 1 nodes_per_replica: 4 server: type: vllm image: name: my-serving-image - enable_expert_parallel: true ``` -The serving image must support the selected model and parallel settings. The runtime starts one vLLM process on each replica node, with one API server and headless followers sharing a rendezvous. EP operates within each replica; independent replica lanes do not form a coordinated data-parallel expert group. +The serving image must support the selected model and parallel settings. The runtime starts one vLLM process on each replica node, with one API server and headless followers sharing a rendezvous. + +### One TP16 replica across two eight-GPU nodes + +For a profile requesting **eight GPUs per node**, the following deployment fragment creates one replica using all 16 GPUs: + +```yaml +resources: + nodes: 2 +topology: + tensor_parallel: 16 + nodes_per_replica: 2 +``` + +Node 0 hosts the API server and eight GPU workers; node 1 hosts a headless follower and eight GPU workers. Both nodes cooperate on the same requests. The native Data Designer client uses a separate CPU step and sends requests to the replica through its logical endpoint. + +This topology describes GPU placement, not whether the model fits in memory. Check the serving runtime's model, quantization, memory, and interconnect requirements separately. If you configure CPU offloading through `server.extra_args`, allow enough host memory and increase model request timeouts to account for the additional latency. + +### Optional expert parallelism + +For a Mixture-of-Experts (MoE) model, `server.enable_expert_parallel: true` asks vLLM to distribute routed experts across its expert-parallel ranks instead of using tensor parallelism to split each expert. Attention layers still use the configured TP layout. See [vLLM's expert parallelism guide](https://docs.vllm.ai/en/latest/serving/expert_parallel_deployment/) for the layer behavior and backend requirements. + +EP defaults to `false` and is not required for a TP group to span nodes. Enabling it does not multiply the GPU count: the TP16/PP1 example still uses 16 GPUs in one replica. Choose EP based on the model, serving image, communication backend, and measured workload behavior. The planner does not create coordinated data-parallel expert groups across independent replicas. + +### Calculate replica counts The planner distributes each replica's GPU workers evenly across its nodes: ```text -PP = pipeline_parallel, or nodes_per_replica when omitted +PP = pipeline_parallel (default: 1) GPUs per replica = tensor_parallel * PP GPUs per replica node = GPUs per replica / nodes_per_replica replicas = (resources.nodes / nodes_per_replica) * (gpus_per_node / GPUs per replica node) ``` -For a profile that requests eight GPUs per node, the following are separate choices for **one deployment in one array task**, with `pipeline_parallel` omitted: +For a profile that requests eight GPUs per node, the following are separate choices for **one deployment in one array task**. PP is explicit in this table: + +| `resources.nodes` | TP | PP | `nodes_per_replica` | vLLM replicas | GPUs per replica | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1 | 1 | 1 | 8 | 1 | +| 1 | 2 | 4 | 1 | 1 | 8 | +| 2 | 8 | 1 | 1 | 2 | 8 | +| 2 | 8 | 2 | 2 | 1 | 16 | +| 2 | 16 | 1 | 2 | 1 | 16 | +| 2 | 4 | 4 | 2 | 1 | 16 | +| 4 | 16 | 1 | 2 | 2 | 16 | +| 4 | 1 | 1 | 1 | 32 | 1 | -| `resources.nodes` | `tensor_parallel` | `nodes_per_replica` | vLLM replicas | GPUs per replica | -| ---: | ---: | ---: | ---: | ---: | -| 1 | 1 | 1 | 8 | 1 | -| 4 | 1 | 1 | 32 | 1 | -| 4 | 2 | 1 | 16 | 2 | -| 4 | 4 | 1 | 8 | 4 | -| 4 | 8 | 1 | 4 | 8 | -| 4 | 8 | 2 | 2 | 16 | +These are valid planner layouts; each model and serving image must also support the selected TP and PP sizes. In the four-node TP16/PP1 example, each pair of nodes hosts an independent 16-GPU replica. Enabling EP keeps those two replicas separate. Increasing `resources.nodes` at TP=1 creates more independent one-GPU servers; it does not combine their GPUs into one larger server. The client uses one logical endpoint per model alias, backed by a local proxy that routes requests among that alias's replicas. `invocation.model_concurrency` limits in-flight model calls **per alias within each array task**; it does not set the replica count or the number of array tasks. diff --git a/packages/data-designer-slurm/examples/benchmark-run.yaml b/packages/data-designer-slurm/examples/benchmark-run.yaml new file mode 100644 index 000000000..4b29cee68 --- /dev/null +++ b/packages/data-designer-slurm/examples/benchmark-run.yaml @@ -0,0 +1,23 @@ +schema_version: 1 +name: benchmark-base-run +builder: + source: builder.yaml +invocation: + num_records: 8 + dataset_name: greetings + model_concurrency: + generator: 8 +deployments: + - model_alias: generator + served_model_name: example/generator + model: example/generator + server: + type: vllm + image: + name: vllm + startup_timeout: 15m + distributed_init_timeout: 10m + topology: + tensor_parallel: 8 +submission: + time_limit: 01:00:00 diff --git a/packages/data-designer-slurm/examples/benchmark.yaml b/packages/data-designer-slurm/examples/benchmark.yaml index f3608b15e..535f365f5 100644 --- a/packages/data-designer-slurm/examples/benchmark.yaml +++ b/packages/data-designer-slurm/examples/benchmark.yaml @@ -1,6 +1,6 @@ schema_version: 1 name: generator-scaling -base_run: run.yaml +base_run: benchmark-run.yaml model_aliases: all concurrency_values: - 32 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py index 368654bed..62baa8fd3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -275,7 +275,7 @@ def validate_shared_node(self) -> DeploymentResources: class DeploymentTopology(AuthoredConfig): tensor_parallel: PositiveInt = 1 nodes_per_replica: PositiveInt = 1 - pipeline_parallel: PositiveInt | None = Field(default=None, exclude_if=lambda value: value is None) + pipeline_parallel: PositiveInt = Field(default=1, exclude_if=lambda value: value == 1) class ServerDeploymentConfig(AuthoredConfig): diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py index 52a771632..05d1f82de 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -265,14 +265,13 @@ def derive( gpus_per_node: int, tensor_parallel: int, nodes_per_replica: int, - pipeline_parallel: int | None = None, + pipeline_parallel: int = 1, ) -> ResolvedTopology: """Derive the canonical v1 topology from placement and GPU resources.""" if min(node_count, gpus_per_node, tensor_parallel, nodes_per_replica) <= 0: raise ValueError("topology derivation inputs must be positive") if node_count % nodes_per_replica: raise ValueError("nodes_per_replica must divide the deployment node count") - pipeline_parallel = nodes_per_replica if pipeline_parallel is None else pipeline_parallel if pipeline_parallel <= 0: raise ValueError("pipeline_parallel must be positive") gpus_per_replica = tensor_parallel * pipeline_parallel diff --git a/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py index 3c1c2a3b1..019bdefb5 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py +++ b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py @@ -3,10 +3,109 @@ from __future__ import annotations +import re +from pathlib import Path + import pytest +import yaml from data_designer.slurm.benchmark import BenchmarkCompiler -from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + load_benchmark_config, + load_run_config, +) +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.resolution import resolve_slurm_config + + +@pytest.mark.parametrize("source", ["shipped", "guide"]) +def test_benchmark_examples_compile_every_child_plan( + source: str, + tmp_path: Path, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + examples = Path(__file__).parents[2] / "examples" + benchmark_file = examples / "benchmark.yaml" + if source == "guide": + guide = Path(__file__).parents[4] / "fern/versions/latest/pages/slurm/benchmarks.mdx" + blocks = re.findall(r"```yaml\n(.*?)```", guide.read_text(), flags=re.DOTALL) + assert len(blocks) == 2 + (tmp_path / "benchmark-run.yaml").write_text(blocks[0]) + benchmark_file = tmp_path / "benchmark.yaml" + benchmark_file.write_text(blocks[1]) + config = load_benchmark_config(benchmark_file) + assert config.base_run.source is not None + base_run = load_run_config(benchmark_file.parent / config.base_run.source) + assert base_run.builder.source is not None + builder_payload = yaml.safe_load((examples / base_run.builder.source).read_text()) + compiled = BenchmarkCompiler.compile(config, base_run) + image = single_node_plan.deployments[0].image.model_copy( + update={"authored_ref": base_run.deployments[0].server.image} + ) + + assert len(compiled.cases) == 4 + for case in compiled.cases: + plan = SlurmRunCompiler.compile( + resolve_slurm_config( + case.child_run_config, + selected_profile=single_node_plan.selected_profile, + client_runtime=single_node_plan.client.runtime, + deployment_images=(image,), + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=case.child_run_id, + package_version=single_node_plan.package_version, + resolved_gpus_per_node=8, + builder_payload=builder_payload, + ) + ) + topology = plan.deployments[0].topology + assert topology.tensor_parallel == 8 + assert topology.pipeline_parallel == 1 + assert topology.gpus_per_replica == 8 + assert topology.replica_count == 2 + assert topology.nodes_per_replica == case.child_run_config.deployments[0].topology.nodes_per_replica + + +@pytest.mark.parametrize("pipeline_parallel", [None, 2], ids=["default-pp1", "explicit-pp2"]) +def test_benchmark_node_placement_preserves_pipeline_size_and_replica_count( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + pipeline_parallel: int | None, +) -> None: + payload = authored_run_single.model_dump(mode="json") + topology = payload["deployments"][0]["topology"] + topology["tensor_parallel"] = 8 if pipeline_parallel is None else 4 + topology.pop("pipeline_parallel", None) + if pipeline_parallel is not None: + topology["pipeline_parallel"] = pipeline_parallel + base_run = DataDesignerSlurmConfig.model_validate(payload) + compiled = BenchmarkCompiler.compile(benchmark_config, base_run) + + for case in compiled.cases: + plan = SlurmRunCompiler.compile( + resolve_slurm_config( + case.child_run_config, + selected_profile=single_node_plan.selected_profile, + client_runtime=single_node_plan.client.runtime, + deployment_images=(single_node_plan.deployments[0].image,), + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=case.child_run_id, + package_version=single_node_plan.package_version, + resolved_gpus_per_node=8, + ) + ) + resolved = plan.deployments[0].topology + assert resolved.pipeline_parallel == (1 if pipeline_parallel is None else pipeline_parallel) + assert resolved.gpus_per_replica == 8 + assert resolved.replica_count == 2 def test_compiler_expands_stable_ordered_ordinary_runs( diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py index 795875881..90cac1c75 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -611,8 +611,8 @@ def test_analysis_accepts_saved_benchmark_from_previous_compiler( ) -> None: config = _inline_config(benchmark_config, authored_run) legacy = BenchmarkCompiler._compile(config, authored_run, legacy=True) - assert legacy.benchmark_id == "benchmark-996c47990258ceb4a3ca7b245ddee971" - assert legacy.cases[0].child_run_id == "run-996c47990258ceb4-0000-89d3e0eccd48" + assert legacy.benchmark_id == "benchmark-88f3e365c1ef64ce47624bebeb80d408" + assert legacy.cases[0].child_run_id == "run-88f3e365c1ef64ce-0000-26e0bbf3b4ad" assert all(case.child_run_config.array_tasks.count == 2 for case in legacy.cases) store, manifest = _manifest(tmp_path, config, legacy) store.publish(config, authored_run, manifest) diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json index 6bcdc0505..484fc082c 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -75,6 +75,7 @@ }, "topology": { "nodes_per_replica": 2, + "pipeline_parallel": 2, "tensor_parallel": 8 } }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json index 2b912541e..23d0c3ef1 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-001/authored-config.json", - "sha256": "8c453e1064bf37c675b8ed62f1fe7b89234c7854e2c2f0955ae0c83938c73506" + "sha256": "0d63492595d41062ba2c2278dd36f7775fca8dacb04d6d8b41f77b6bbbb1605b" }, "builder": { "authored_source": null, @@ -148,6 +148,7 @@ }, "topology": { "nodes_per_replica": 2, + "pipeline_parallel": 2, "tensor_parallel": 8 } }, diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index c57fa9048..6d89774c8 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -560,6 +560,36 @@ def test_vllm_rejects_negative_launch_timing(field: str) -> None: VllmServerConfig.model_validate({"type": "vllm", "image": {"name": "vllm"}, field: "-1s"}) +@pytest.mark.parametrize("nodes_per_replica", [1, 2, 4]) +def test_deployment_defaults_to_one_pipeline_stage_without_expert_parallelism(nodes_per_replica: int) -> None: + deployment = ServerDeploymentConfig.model_validate( + { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + "resources": {"nodes": nodes_per_replica}, + "topology": {"tensor_parallel": 16, "nodes_per_replica": nodes_per_replica}, + } + ) + + assert deployment.topology.pipeline_parallel == 1 + assert deployment.server.enable_expert_parallel is False + assert ServerDeploymentConfig.model_validate_json(deployment.model_dump_json()) == deployment + + +@pytest.mark.parametrize("pipeline_parallel", [None, 0, -1]) +def test_deployment_rejects_invalid_pipeline_parallel(pipeline_parallel: int | None) -> None: + with pytest.raises(ValidationError, match="pipeline_parallel"): + ServerDeploymentConfig.model_validate( + { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + "topology": {"pipeline_parallel": pipeline_parallel}, + } + ) + + def test_deployment_rejects_invalid_topology() -> None: payload = { "model_alias": "generator", diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index 6601b8db9..459d82f8f 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -35,11 +35,11 @@ def test_resolved_topology_derives_all_resource_fields() -> None: assert topology == ResolvedTopology( tensor_parallel=4, nodes_per_replica=2, - pipeline_parallel=2, + pipeline_parallel=1, node_group_count=2, - replicas_per_node_group=2, - replica_count=4, - gpus_per_replica=8, + replicas_per_node_group=4, + replica_count=8, + gpus_per_replica=4, ) diff --git a/packages/data-designer-slurm/tests/integration/test_cross_node_parallelism.py b/packages/data-designer-slurm/tests/integration/test_cross_node_parallelism.py index 0b21a9e62..3016aeed2 100644 --- a/packages/data-designer-slurm/tests/integration/test_cross_node_parallelism.py +++ b/packages/data-designer-slurm/tests/integration/test_cross_node_parallelism.py @@ -30,6 +30,8 @@ def _compile( pp: int | None, nodes: int, groups: int = 1, + ep: bool | None = None, + gpus_per_node: int = 4, ) -> ResolvedSlurmRunPlan: payload = authored.model_dump(mode="json") deployment = payload["deployments"][0] @@ -37,9 +39,11 @@ def _compile( deployment["topology"] = {"tensor_parallel": tp, "nodes_per_replica": nodes} if pp is not None: deployment["topology"]["pipeline_parallel"] = pp - deployment["server"]["enable_expert_parallel"] = True + deployment["server"].pop("enable_expert_parallel", None) + if ep is not None: + deployment["server"]["enable_expert_parallel"] = ep authored = DataDesignerSlurmConfig.model_validate(payload) - profile = injected_profile(template.selected_profile.profile.model_copy(update={"gpus_per_node": 4})) + profile = injected_profile(template.selected_profile.profile.model_copy(update={"gpus_per_node": gpus_per_node})) return SlurmRunCompiler.compile( resolve_slurm_config( authored, @@ -50,7 +54,7 @@ def _compile( runtime_bundle=template.runtime_bundle, run_id=template.run_id, package_version=template.package_version, - resolved_gpus_per_node=4, + resolved_gpus_per_node=gpus_per_node, ) ) @@ -62,12 +66,13 @@ def _compile( (8, 1, 2, 1, 1), (4, 1, 2, 2, 4), (8, 2, 4, 1, 1), - (4, None, 2, 1, 1), + (8, None, 2, 1, 1), (2, 2, 1, 1, 1), (1, 4, 1, 1, 1), ], - ids=["tp16-ep", "tp8-ep", "two-lanes-two-groups", "hybrid-tp-pp", "legacy-pp", "local-pp", "local-pp-tp1"], + ids=["tp16", "tp8", "two-lanes-two-groups", "hybrid-tp-pp", "default-pp1", "local-pp", "local-pp-tp1"], ) +@pytest.mark.parametrize("ep", [None, False, True], ids=["ep-omitted", "ep-disabled", "ep-enabled"]) def test_compiled_parallel_topology_drives_node_gpu_ownership_and_commands( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -77,14 +82,17 @@ def test_compiled_parallel_topology_drives_node_gpu_ownership_and_commands( nodes: int, groups: int, replicas: int, + ep: bool | None, ) -> None: plan = _compile( - authored_run_single, dependency_lock_single, single_node_plan, tp=tp, pp=pp, nodes=nodes, groups=groups + authored_run_single, dependency_lock_single, single_node_plan, tp=tp, pp=pp, nodes=nodes, groups=groups, ep=ep ) assert ResolvedSlurmRunPlan.model_validate_json(plan.serialize_json()).compute_sha256() == plan.compute_sha256() server = resolve_vllm_server(plan, plan.deployments[0].deployment_id) assert ResolvedVllmServerDeployment.model_validate_json(server.model_dump_json()) == server - effective_pp = nodes if pp is None else pp + effective_pp = 1 if pp is None else pp + assert plan.deployments[0].authored.topology.pipeline_parallel == effective_pp + assert plan.deployments[0].authored.server.enable_expert_parallel is bool(ep) assert server.topology.replica_count == replicas assert server.topology.gpus_per_replica == tp * effective_pp assert len(server.processes) == nodes * replicas @@ -105,7 +113,7 @@ def test_compiled_parallel_topology_drives_node_gpu_ownership_and_commands( command = build_vllm_process_command(server, process, plan, layout) assert command[command.index("--tensor-parallel-size") + 1] == str(tp) assert command[command.index("--distributed-executor-backend") + 1] == "mp" - assert "--enable-expert-parallel" in command + assert ("--enable-expert-parallel" in command) is bool(ep) if nodes > 1: assert command[command.index("--pipeline-parallel-size") + 1] == str(effective_pp) assert command[command.index("--nnodes") + 1] == str(nodes) @@ -119,6 +127,36 @@ def test_compiled_parallel_topology_drives_node_gpu_ownership_and_commands( assert "--nnodes" not in command +@pytest.mark.parametrize(("nodes", "replicas"), [(2, 1), (4, 2)]) +def test_tp16_on_eight_gpu_nodes_defaults_to_pp1_without_ep( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + nodes: int, + replicas: int, +) -> None: + plan = _compile( + authored_run_single, + dependency_lock_single, + single_node_plan, + tp=16, + pp=None, + nodes=nodes, + gpus_per_node=8, + ) + server = resolve_vllm_server(plan, plan.deployments[0].deployment_id) + layout = AllocationLayout(tuple(f"gpu-{index}" for index in range(nodes))) + + assert server.topology.pipeline_parallel == 1 + assert server.topology.gpus_per_replica == 16 + assert server.topology.replica_count == replicas + for process in server.processes: + assert len(process.gpu_indices) == 16 // nodes + command = build_vllm_process_command(server, process, plan, layout) + assert command[command.index("--pipeline-parallel-size") + 1] == "1" + assert "--enable-expert-parallel" not in command + + @pytest.mark.parametrize(("tp", "pp", "nodes"), [(3, 1, 2), (32, 1, 4), (12, 1, 2)]) def test_invalid_local_gpu_counts_fail_before_submission( authored_run_single: DataDesignerSlurmConfig, diff --git a/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index 0f6f6d66a..bdcc16941 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -345,6 +345,7 @@ def _independent_replica_placement(plan: ResolvedSlurmRunPlan) -> ResolvedDeploy payload = plan.deployments[0].model_dump(mode="json") payload["authored"]["server"]["enable_expert_parallel"] = True payload["authored"]["topology"]["nodes_per_replica"] = 1 + payload["authored"]["topology"]["pipeline_parallel"] = 1 payload["topology"].update( nodes_per_replica=1, pipeline_parallel=1, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 7a2687367..cabc4ee5e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -18,7 +18,7 @@ export PYTHONNOUSERSITE=1 readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="cc5b1fd8d7f3ad664ee56098a02747fa69888f0979646f67c78796d46674044b" +readonly DD_PLAN_SHA256="07cc32a1cf15b7ec7a6159319346d4a901c765758d2ead29900ac91f012b1881" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index d1bf2461e..80418eed8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -24,7 +24,7 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="c185ed76ac5e95e778d2ef8f442c046aa0fb982ef318a4404a9f5124c1a6484d", + expected_fixture_sha256="a4e6342dad6606b7dd6a2011b03e8566d017a5b30adfad337526e5a4b0567fea", )