Multi-node multi-instance vLLM via gym eval submit

examples/slurm_vllm_ray_multi_node.yaml

@@ -0,0 +1,41 @@
+compute:
+  cluster:
+    type: slurm
+    hostname: myslurmcluster
+    walltime: "01:00:00"
+    account: myaccount
+    node_pools:
+      compute:
+        partition: batch
+        nodes: 2
+        ntasks_per_node: 1
+        gpus_per_node: 8
+
+services:
+  vllm_model:
+    container: vllm/vllm-openai:v0.9.0
+    type: vllm
+    model: Qwen/Qwen2.5-72B-Instruct
+    trust_remote_code: true
+    tensor_parallel_size: 2      # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs)
+    number_of_instances: 8       # footprint exceeds one node, so the ray backend is used automatically
+    port: 8000
+    health_check:
+      timeout_seconds: 1200
+    mounts:
+      - /lustre/datasets:/data
+
+driver:
+  policy_model: vllm_model
+  container: python:3.13
+  gym_install:
+    ref: main
+  benchmarks:
+    gpqa:
+      run:
+        split: benchmark
+        config_paths:
+          - benchmarks/gpqa/config.yaml
+
+job:
+  output_path: /lustre/fsw/my-path

nemo_gym/orchestration/api.py

@@ -49,26 +49,12 @@ class BaseModelServiceConfig(BaseServiceConfig):
     port: int = 8000
 
 
-class VllmServiceDistributedBackend(_StrictModel):
-    """Use vLLM's native data-parallel multi-instance (--data-parallel-size N)."""
-
-    type: Literal["mp"] = "mp"
-
-
-# Future backends: add Annotated[RayServeDistributedBackend, Tag("ray_serve")], etc.
-DistributedBackendConfig = Annotated[
-    Annotated[VllmServiceDistributedBackend, Tag("mp")],
-    Discriminator("type"),
-]
-
-
 class VllmServiceConfig(BaseModelServiceConfig):
     type: Literal["vllm"]
     tensor_parallel_size: int = 1
     pipeline_parallel_size: int = 1
     trust_remote_code: bool = False
     number_of_instances: int = 1
-    distributed_backend: DistributedBackendConfig | None = None
 
     @field_validator("number_of_instances")
     @classmethod
@@ -77,14 +63,6 @@ def _validate_number_of_instances(cls, v: int) -> int:
             raise ValueError(f"number_of_instances must be >= 1, got {v}")
         return v
 
-    @model_validator(mode="after")
-    def _validate_distributed_backend(self) -> "VllmServiceConfig":
-        if self.number_of_instances == 1 and self.distributed_backend is not None:
-            raise ValueError("distributed_backend should not be set when number_of_instances == 1")
-        if self.number_of_instances > 1 and self.distributed_backend is None:
-            self.distributed_backend = VllmServiceDistributedBackend()
-        return self
-
     @model_validator(mode="after")
     def _default_health_check(self) -> "VllmServiceConfig":
         # vLLM always exposes /health on its serving port; set it automatically
@@ -181,6 +159,11 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
             raise ValueError(f"Multiple compute resources are not supported yet ({', '.join(sorted(compute_names))}).")
 
         sole_compute = next(iter(compute_names))
+        compute = self.compute[sole_compute]
+        total_nodes = (
+            sum(p.nodes for p in compute.node_pools.values()) if isinstance(compute, SlurmComputeConfig) else 1
+        )
+        is_multi_node = total_nodes > 1
 
         for service_name, service in self.services.items():
             if service.placement is None:
@@ -191,8 +174,20 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
                     f"({', '.join(sorted(compute_names))})."
                 )
 
+            if (
+                is_multi_node
+                and isinstance(service, VllmServiceConfig)
+                and service.number_of_instances > 1
+                and service.number_of_instances % total_nodes != 0
+            ):
+                raise ValueError(
+                    f"Service '{service_name}' has number_of_instances={service.number_of_instances}, which must "
+                    f"be evenly divisible by the number of nodes ({total_nodes}) for multi-node data-parallel "
+                    "deployment - each node hosts an equal share of the data-parallel replicas."
+                )
+
             if isinstance(service, VllmServiceConfig):
-                self._validate_vllm_gpu_footprint(service_name, service)
+                self._validate_vllm_gpu_footprint(service_name, service, total_nodes)
 
         if self.driver.policy_model is not None:
             if self.driver.policy_model not in self.services:
@@ -218,7 +213,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
 
         return self
 
-    def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceConfig") -> None:
+    def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceConfig", total_nodes: int) -> None:
         compute = self.compute[service.placement]
         if not isinstance(compute, SlurmComputeConfig):
             return
@@ -230,25 +225,68 @@ def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceC
             return
 
         max_gpus_per_node = max(gpus_per_node_values)
-        gpus_needed = service.tensor_parallel_size * service.pipeline_parallel_size * service.number_of_instances
+        tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size
+
+        if total_nodes > 1 and service.number_of_instances > 1:
+            if tp_pp > max_gpus_per_node:
+                # Each instance's own TP/PP footprint already exceeds a single node's GPU count, so
+                # spreading multiple such instances across nodes would require every instance to
+                # itself span multiple nodes. That's not supported: multi-node data-parallel only
+                # distributes whole instances across nodes with tensor/pipeline parallelism kept
+                # local to each node (see _build_vllm_multi_instance_multi_node_command).
+                raise ValueError(
+                    f"Service '{service_name}' sets number_of_instances={service.number_of_instances} with "
+                    f"tensor_parallel_size={service.tensor_parallel_size} x "
+                    f"pipeline_parallel_size={service.pipeline_parallel_size}={tp_pp}, which exceeds a single "
+                    f"node's gpus_per_node ({max_gpus_per_node}). Multiple instances where each instance's own "
+                    "tensor/pipeline-parallel footprint spans multiple nodes is not supported - reduce "
+                    "tensor_parallel_size/pipeline_parallel_size to fit within one node, or set "
+                    "number_of_instances=1 to let a single instance span nodes."
+                )
+            # Multi-node data-parallel: each node runs its own equal share of the replicas with
+            # local tensor/pipeline parallelism (see _build_vllm_multi_instance_multi_node_command);
+            # the per-node share, not the total footprint, has to fit in that node's GPU count.
+            instances_per_node = service.number_of_instances // total_nodes
+            gpus_needed = tp_pp * instances_per_node
+            gpus_available = max_gpus_per_node
+            footprint = (
+                f"{instances_per_node} local replica(s) per node (number_of_instances="
+                f"{service.number_of_instances} / {total_nodes} nodes) x tensor_parallel_size="
+                f"{service.tensor_parallel_size} x pipeline_parallel_size={service.pipeline_parallel_size}"
+            )
+            scope = f"a single node's gpus_per_node ({max_gpus_per_node})"
+        elif total_nodes > 1:
+            # Single instance's TP/PP footprint spans the whole allocation via the ray backend.
+            gpus_needed = tp_pp
+            gpus_available = sum(
+                pool.nodes * pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node
+            )
+            footprint = (
+                f"tensor_parallel_size={service.tensor_parallel_size} x "
+                f"pipeline_parallel_size={service.pipeline_parallel_size}"
+            )
+            scope = f"the total GPUs across all nodes ({gpus_available})"
+        else:
+            gpus_needed = tp_pp * service.number_of_instances
+            gpus_available = max_gpus_per_node
+            footprint = (
+                f"tensor_parallel_size={service.tensor_parallel_size} x "
+                f"pipeline_parallel_size={service.pipeline_parallel_size} x "
+                f"number_of_instances={service.number_of_instances}"
+            )
+            scope = f"the node pool's gpus_per_node ({max_gpus_per_node})"
 
-        if gpus_needed > max_gpus_per_node:
+        if gpus_needed > gpus_available:
             raise ValueError(
-                f"Service '{service_name}' requires {gpus_needed} GPUs "
-                f"(tensor_parallel_size={service.tensor_parallel_size} x "
-                f"pipeline_parallel_size={service.pipeline_parallel_size} x "
-                f"number_of_instances={service.number_of_instances}), which exceeds the largest available "
-                f"node pool's gpus_per_node ({max_gpus_per_node}) on compute '{service.placement}'. "
-                "Multi-node vLLM services are not supported yet by the 'mp' distributed backend; "
-                "reduce number_of_instances/tensor_parallel_size/pipeline_parallel_size to fit on a single node."
+                f"Service '{service_name}' requires {gpus_needed} GPUs ({footprint}), which exceeds {scope} "
+                f"on compute '{service.placement}'. Reduce number_of_instances/tensor_parallel_size/"
+                "pipeline_parallel_size, or add more nodes/GPUs."
             )
-        elif gpus_needed < max_gpus_per_node:
+        elif gpus_needed < gpus_available:
             warnings.warn(
-                f"Service '{service_name}' requires {gpus_needed} GPUs "
-                f"(tensor_parallel_size={service.tensor_parallel_size} x "
-                f"pipeline_parallel_size={service.pipeline_parallel_size} x "
-                f"number_of_instances={service.number_of_instances}) but compute '{service.placement}' allocates "
-                f"nodes with {max_gpus_per_node} GPUs each, leaving {max_gpus_per_node - gpus_needed} GPU(s) idle. "
-                "Increase number_of_instances/tensor_parallel_size or reduce gpus_per_node to use the full node.",
+                f"Service '{service_name}' requires {gpus_needed} GPUs ({footprint}) but compute "
+                f"'{service.placement}' provides {scope}, leaving {gpus_available - gpus_needed} GPU(s) idle. "
+                "Increase number_of_instances/tensor_parallel_size or reduce gpus_per_node to use the full "
+                "allocation.",
                 stacklevel=2,
             )

nemo_gym/orchestration/executors/script_templates.py

@@ -17,6 +17,35 @@
 import shlex
 
 
+_RAY_PRELUDE = """\
+# Resolve the head node IP for multi-node vLLM services (spanning nodes via Ray).
+nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+nodes_array=($nodes)
+head_node_hostname=${nodes_array[0]}
+head_node_ip=$(getent hosts "$head_node_hostname" | awk '{print $1}')
+export HEAD_NODE_IP="$head_node_ip"
+export RAY_HEAD_NODE_IP="$head_node_ip:6379"
+echo "Head node IP address: $HEAD_NODE_IP\""""
+
+
+_VLLM_RAY_SYMMETRIC_RUN = """\
+bash -lc '
+    command -v ray >/dev/null 2>&1 || pip install -q "ray[default]"
+    if ray symmetric-run --help >/dev/null 2>&1; then
+        ray symmetric-run \\
+            --address "$RAY_HEAD_NODE_IP" \\
+            --min-nodes {total_nodes} \\
+            {resource_flags} \\
+            -- {inner_cmd}
+    elif [ "$SLURM_NODEID" = "0" ]; then
+        ray start --head --port=6379 {resource_flags}
+        {inner_cmd}
+    else
+        ray start --address="$RAY_HEAD_NODE_IP" {resource_flags} --block
+    fi
+'"""
+
+
 _HEALTH_WAIT_MULTI = """\
 # Wait for {name} (try multiple health endpoints)
 echo "Waiting for {name} at {url}..."
@@ -44,6 +73,22 @@ def bash_var(name: str) -> str:
     return re.sub(r"[^A-Z0-9]", "_", name.upper())
 
 
+def render_ray_prelude() -> str:
+    return _RAY_PRELUDE
+
+
+def render_vllm_ray_symmetric_run(inner_cmd: str, total_nodes: int, resource_flags: str) -> str:
+    """Render the Ray head/worker bootstrap that wraps a single vLLM instance's TP/PP command so
+    it spans multiple Slurm nodes.
+
+    Uses `ray symmetric-run` when available (Ray >= 2.50), which starts/joins a Ray cluster across
+    every task and runs the entrypoint only on the elected head node. Containers with an older Ray
+    pin fall back to manually starting head/worker Ray processes, keyed on Slurm's per-node task
+    rank ($SLURM_NODEID).
+    """
+    return _VLLM_RAY_SYMMETRIC_RUN.format(total_nodes=total_nodes, resource_flags=resource_flags, inner_cmd=inner_cmd)
+
+
 def render_health_check(name: str, port: int, path: str, timeout: int) -> str:
     return _HEALTH_WAIT_MULTI.format(
         name=name,