From af94b07a11eee39f2c6f27e8124616973ced10ab Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Thu, 20 Aug 2026 19:35:56 +0000 Subject: [PATCH 01/15] Init ultra config Signed-off-by: Frankie Siino --- .../vllm_configs/nemotron_3_ultra.sh | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh new file mode 100644 index 0000000000..9ccd64102f --- /dev/null +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Nemotron 3 Ultra BF16 baseline for disaggregated prefill/decode on 4-GPU +# GB200 nodes. This launcher's TP4-per-node layout needs four coupled DP ranks +# per tier so expert parallelism can shard the 512 experts over 16 GPUs. + +# InstantTensor's io_uring loader returned EIO while 24 ranks concurrently read +# the Lustre-hosted checkpoint in Run 001. Keep the loader selectable for later +# experiments, but use vLLM's standard safetensors path for the stable baseline. +ULTRA_LOAD_FORMAT="${ULTRA_LOAD_FORMAT:-safetensors}" +ULTRA_GPU_MEMORY_UTILIZATION="${ULTRA_GPU_MEMORY_UTILIZATION:-0.95}" +ULTRA_PREFILL_GPU_MEMORY_UTILIZATION="${ULTRA_PREFILL_GPU_MEMORY_UTILIZATION:-$ULTRA_GPU_MEMORY_UTILIZATION}" +ULTRA_DECODE_GPU_MEMORY_UTILIZATION="${ULTRA_DECODE_GPU_MEMORY_UTILIZATION:-$ULTRA_GPU_MEMORY_UTILIZATION}" +ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS="${ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS:-16384}" +ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS="${ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS:-8192}" +ULTRA_MAX_NUM_SEQS="${ULTRA_MAX_NUM_SEQS:-64}" +ULTRA_ENABLE_PREFIX_CACHING="${ULTRA_ENABLE_PREFIX_CACHING:-0}" +ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS="${ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS:-0}" +ULTRA_DECODE_CUDAGRAPH_MODE="${ULTRA_DECODE_CUDAGRAPH_MODE:-FULL_DECODE_ONLY}" +ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS="${ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS:-0}" +ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE="${ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE:-0}" +export SAFETENSORS_FAST_GPU=1 + +case "$ULTRA_DECODE_CUDAGRAPH_MODE" in + FULL_DECODE_ONLY | PIECEWISE | NONE) ;; + *) + echo "ERROR: ULTRA_DECODE_CUDAGRAPH_MODE must be FULL_DECODE_ONLY, PIECEWISE, or NONE; got '$ULTRA_DECODE_CUDAGRAPH_MODE'." >&2 + exit 2 + ;; +esac + +case "$ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS" in + 0) + ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS_JSON=false + ;; + 1) + ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS_JSON=true + ;; + *) + echo "ERROR: ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS must be 0 or 1; got '$ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS'." >&2 + exit 2 + ;; +esac + +case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in + 0) + ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=false + ;; + 1) + ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=true + ;; + *) + echo "ERROR: ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS must be 0 or 1; got '$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS'." >&2 + exit 2 + ;; +esac + +case "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" in + 0 | 1) ;; + *) + echo "ERROR: ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE must be 0 or 1; got '$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE'." >&2 + exit 2 + ;; +esac + +VLLM_COMMON_ARGS=( + --disable-uvicorn-access-log + --trust-remote-code + --dtype bfloat16 + --distributed-executor-backend mp + --data-parallel-backend mp + --max-model-len 262144 + --enable-auto-tool-choice + --tool-call-parser qwen3_coder + --reasoning-parser nemotron_v3 + --enable-chunked-prefill + --kv-cache-dtype fp8 + --no-disable-hybrid-kv-cache-manager + --no-async-scheduling + --block-size 128 + --mamba-cache-mode align + --mamba-ssm-cache-dtype float16 + --mamba-backend flashinfer + --enable-mamba-cache-stochastic-rounding + --mamba-cache-philox-rounds 5 + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 96}' + --load-format "$ULTRA_LOAD_FORMAT" + --enable-expert-parallel + --distributed-timeout-seconds 3600 +) + +if [[ "$ULTRA_ENABLE_PREFIX_CACHING" == "1" ]]; then + VLLM_COMMON_ARGS+=(--enable-prefix-caching) +else + # vLLM 0.25.1's NIXL pull connector cannot reconcile multiple locally + # prefix-cached Mamba/SSM blocks with transferred prefill state. + VLLM_COMMON_ARGS+=(--no-enable-prefix-caching) +fi + +VLLM_PREFILL_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}' + --compilation-config "{\"cudagraph_copy_inputs\":$ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" + --gpu-memory-utilization "$ULTRA_PREFILL_GPU_MEMORY_UTILIZATION" + --max-num-batched-tokens "$ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS" + --max-num-seqs "$ULTRA_MAX_NUM_SEQS" + --data-parallel-size-local 1 + --tensor-parallel-size 4 +) + +VLLM_DECODE_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}' + --compilation-config "{\"cudagraph_mode\":\"$ULTRA_DECODE_CUDAGRAPH_MODE\",\"cudagraph_copy_inputs\":$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" + --gpu-memory-utilization "$ULTRA_DECODE_GPU_MEMORY_UTILIZATION" + --max-num-batched-tokens "$ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS" + --max-num-seqs "$ULTRA_MAX_NUM_SEQS" + --data-parallel-size-local 1 + --tensor-parallel-size 4 +) + +if [[ "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" == "1" ]]; then + # Isolate custom all-reduce from the graph-enabled distributed decode path. + # vLLM falls back to its standard NCCL-backed tensor-parallel reductions. + VLLM_DECODE_ARGS+=(--disable-custom-all-reduce) +fi + +if [[ "${ULTRA_ENABLE_MTP:-1}" == "1" ]]; then + # Producer and consumer must expose matching cache layouts for NIXL KV transfer. + VLLM_PREFILL_ARGS+=(--speculative-config '{"method":"nemotron_h_mtp","num_speculative_tokens":5}') + VLLM_DECODE_ARGS+=(--speculative-config '{"method":"nemotron_h_mtp","num_speculative_tokens":5}') +fi From 62d9f78ff84872bd16e28c3e8a04f9e00b307e71 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 21 Aug 2026 00:46:25 +0000 Subject: [PATCH 02/15] Tune Nemotron 3 Ultra vLLM config Signed-off-by: Frankie Siino --- .../vllm_configs/nemotron_3_ultra.sh | 66 +++++++++++++++++-- 1 file changed, 61 insertions(+), 5 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh index 9ccd64102f..78e5774474 100644 --- a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -10,17 +10,24 @@ # the Lustre-hosted checkpoint in Run 001. Keep the loader selectable for later # experiments, but use vLLM's standard safetensors path for the stable baseline. ULTRA_LOAD_FORMAT="${ULTRA_LOAD_FORMAT:-safetensors}" -ULTRA_GPU_MEMORY_UTILIZATION="${ULTRA_GPU_MEMORY_UTILIZATION:-0.95}" -ULTRA_PREFILL_GPU_MEMORY_UTILIZATION="${ULTRA_PREFILL_GPU_MEMORY_UTILIZATION:-$ULTRA_GPU_MEMORY_UTILIZATION}" -ULTRA_DECODE_GPU_MEMORY_UTILIZATION="${ULTRA_DECODE_GPU_MEMORY_UTILIZATION:-$ULTRA_GPU_MEMORY_UTILIZATION}" +# Run 021's all-eager P4/D4 baseline needs more prefill activation headroom, +# while decode can retain a larger KV cache. ULTRA_GPU_MEMORY_UTILIZATION is +# preserved as an optional global override; tier-specific overrides take priority. +ULTRA_GPU_MEMORY_UTILIZATION="${ULTRA_GPU_MEMORY_UTILIZATION:-}" +ULTRA_PREFILL_GPU_MEMORY_UTILIZATION="${ULTRA_PREFILL_GPU_MEMORY_UTILIZATION:-${ULTRA_GPU_MEMORY_UTILIZATION:-0.90}}" +ULTRA_DECODE_GPU_MEMORY_UTILIZATION="${ULTRA_DECODE_GPU_MEMORY_UTILIZATION:-${ULTRA_GPU_MEMORY_UTILIZATION:-0.95}}" ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS="${ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS:-16384}" ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS="${ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS:-8192}" ULTRA_MAX_NUM_SEQS="${ULTRA_MAX_NUM_SEQS:-64}" ULTRA_ENABLE_PREFIX_CACHING="${ULTRA_ENABLE_PREFIX_CACHING:-0}" ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS="${ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS:-0}" -ULTRA_DECODE_CUDAGRAPH_MODE="${ULTRA_DECODE_CUDAGRAPH_MODE:-FULL_DECODE_ONLY}" +ULTRA_PREFILL_ENFORCE_EAGER="${ULTRA_PREFILL_ENFORCE_EAGER:-1}" +ULTRA_DECODE_CUDAGRAPH_MODE="${ULTRA_DECODE_CUDAGRAPH_MODE:-NONE}" ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS="${ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS:-0}" +ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES="${ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES:-}" ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE="${ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE:-0}" +ULTRA_DECODE_ALL2ALL_BACKEND="${ULTRA_DECODE_ALL2ALL_BACKEND:-}" +ULTRA_DECODE_ENFORCE_EAGER="${ULTRA_DECODE_ENFORCE_EAGER:-1}" export SAFETENSORS_FAST_GPU=1 case "$ULTRA_DECODE_CUDAGRAPH_MODE" in @@ -44,6 +51,14 @@ case "$ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS" in ;; esac +case "$ULTRA_PREFILL_ENFORCE_EAGER" in + 0 | 1) ;; + *) + echo "ERROR: ULTRA_PREFILL_ENFORCE_EAGER must be 0 or 1; got '$ULTRA_PREFILL_ENFORCE_EAGER'." >&2 + exit 2 + ;; +esac + case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in 0) ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=false @@ -57,6 +72,15 @@ case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in ;; esac +ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES_JSON="" +if [[ -n "$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES" ]]; then + if [[ ! "$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES" =~ ^\[[0-9]+(,[0-9]+)*\]$ ]]; then + echo "ERROR: ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES must be a compact JSON integer array; got '$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES'." >&2 + exit 2 + fi + ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES_JSON=",\"cudagraph_capture_sizes\":$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES" +fi + case "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" in 0 | 1) ;; *) @@ -65,6 +89,22 @@ case "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" in ;; esac +case "$ULTRA_DECODE_ALL2ALL_BACKEND" in + "" | allgather_reducescatter | deepep_high_throughput | deepep_low_latency | deepep_v2 | flashinfer_all2allv | flashinfer_nvlink_one_sided | flashinfer_nvlink_two_sided | mori_high_throughput | mori_low_latency | nixl_ep) ;; + *) + echo "ERROR: unsupported ULTRA_DECODE_ALL2ALL_BACKEND '$ULTRA_DECODE_ALL2ALL_BACKEND'." >&2 + exit 2 + ;; +esac + +case "$ULTRA_DECODE_ENFORCE_EAGER" in + 0 | 1) ;; + *) + echo "ERROR: ULTRA_DECODE_ENFORCE_EAGER must be 0 or 1; got '$ULTRA_DECODE_ENFORCE_EAGER'." >&2 + exit 2 + ;; +esac + VLLM_COMMON_ARGS=( --disable-uvicorn-access-log --trust-remote-code @@ -109,9 +149,14 @@ VLLM_PREFILL_ARGS=( --tensor-parallel-size 4 ) +if [[ "$ULTRA_PREFILL_ENFORCE_EAGER" == "1" ]]; then + # Bypass torch.compile as well as CUDA graphs on the prefill tier. + VLLM_PREFILL_ARGS+=(--enforce-eager) +fi + VLLM_DECODE_ARGS=( --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}' - --compilation-config "{\"cudagraph_mode\":\"$ULTRA_DECODE_CUDAGRAPH_MODE\",\"cudagraph_copy_inputs\":$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" + --compilation-config "{\"cudagraph_mode\":\"$ULTRA_DECODE_CUDAGRAPH_MODE\",\"cudagraph_copy_inputs\":$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" --gpu-memory-utilization "$ULTRA_DECODE_GPU_MEMORY_UTILIZATION" --max-num-batched-tokens "$ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS" --max-num-seqs "$ULTRA_MAX_NUM_SEQS" @@ -125,6 +170,17 @@ if [[ "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" == "1" ]]; then VLLM_DECODE_ARGS+=(--disable-custom-all-reduce) fi +if [[ -n "$ULTRA_DECODE_ALL2ALL_BACKEND" ]]; then + # Select the collective used to dispatch tokens to the expert-parallel + # decode ranks and combine the expert outputs. Prefill stays on the default. + VLLM_DECODE_ARGS+=(--all2all-backend "$ULTRA_DECODE_ALL2ALL_BACKEND") +fi + +if [[ "$ULTRA_DECODE_ENFORCE_EAGER" == "1" ]]; then + # Bypass torch.compile as well as CUDA graphs on the decode tier. + VLLM_DECODE_ARGS+=(--enforce-eager) +fi + if [[ "${ULTRA_ENABLE_MTP:-1}" == "1" ]]; then # Producer and consumer must expose matching cache layouts for NIXL KV transfer. VLLM_PREFILL_ARGS+=(--speculative-config '{"method":"nemotron_h_mtp","num_speculative_tokens":5}') From cd9eb6aacab70b1d83a243df85dab325c9a153a9 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Thu, 3 Sep 2026 22:58:10 +0000 Subject: [PATCH 03/15] Add configurable MTP depth for Nemotron 3 Ultra tuning Allow experiments to override the speculative-token count while preserving the current default. Signed-off-by: Frankie Siino --- .../vllm_configs/nemotron_3_ultra.sh | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh index 78e5774474..d690588459 100644 --- a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -28,6 +28,7 @@ ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES="${ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES:-}" ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE="${ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE:-0}" ULTRA_DECODE_ALL2ALL_BACKEND="${ULTRA_DECODE_ALL2ALL_BACKEND:-}" ULTRA_DECODE_ENFORCE_EAGER="${ULTRA_DECODE_ENFORCE_EAGER:-1}" +ULTRA_NUM_SPECULATIVE_TOKENS="${ULTRA_NUM_SPECULATIVE_TOKENS:-5}" export SAFETENSORS_FAST_GPU=1 case "$ULTRA_DECODE_CUDAGRAPH_MODE" in @@ -105,6 +106,11 @@ case "$ULTRA_DECODE_ENFORCE_EAGER" in ;; esac +if [[ ! "$ULTRA_NUM_SPECULATIVE_TOKENS" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: ULTRA_NUM_SPECULATIVE_TOKENS must be a positive integer; got '$ULTRA_NUM_SPECULATIVE_TOKENS'." >&2 + exit 2 +fi + VLLM_COMMON_ARGS=( --disable-uvicorn-access-log --trust-remote-code @@ -183,6 +189,7 @@ fi if [[ "${ULTRA_ENABLE_MTP:-1}" == "1" ]]; then # Producer and consumer must expose matching cache layouts for NIXL KV transfer. - VLLM_PREFILL_ARGS+=(--speculative-config '{"method":"nemotron_h_mtp","num_speculative_tokens":5}') - VLLM_DECODE_ARGS+=(--speculative-config '{"method":"nemotron_h_mtp","num_speculative_tokens":5}') + ULTRA_SPECULATIVE_CONFIG="{\"method\":\"nemotron_h_mtp\",\"num_speculative_tokens\":$ULTRA_NUM_SPECULATIVE_TOKENS}" + VLLM_PREFILL_ARGS+=(--speculative-config "$ULTRA_SPECULATIVE_CONFIG") + VLLM_DECODE_ARGS+=(--speculative-config "$ULTRA_SPECULATIVE_CONFIG") fi From a75b4d1d118afa2d5e5489c7629de745e3878ed6 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Sat, 5 Sep 2026 01:08:52 +0000 Subject: [PATCH 04/15] Add Ultra async scheduling and EPLB tuning controls Signed-off-by: Frankie Siino --- .../vllm_configs/nemotron_3_ultra.sh | 51 ++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh index d690588459..e952ea91c3 100644 --- a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -22,12 +22,16 @@ ULTRA_MAX_NUM_SEQS="${ULTRA_MAX_NUM_SEQS:-64}" ULTRA_ENABLE_PREFIX_CACHING="${ULTRA_ENABLE_PREFIX_CACHING:-0}" ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS="${ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS:-0}" ULTRA_PREFILL_ENFORCE_EAGER="${ULTRA_PREFILL_ENFORCE_EAGER:-1}" +ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING="${ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING:-0}" ULTRA_DECODE_CUDAGRAPH_MODE="${ULTRA_DECODE_CUDAGRAPH_MODE:-NONE}" ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS="${ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS:-0}" ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES="${ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES:-}" ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE="${ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE:-0}" ULTRA_DECODE_ALL2ALL_BACKEND="${ULTRA_DECODE_ALL2ALL_BACKEND:-}" +ULTRA_DECODE_ENABLE_EPLB="${ULTRA_DECODE_ENABLE_EPLB:-0}" +ULTRA_DECODE_EPLB_CONFIG="${ULTRA_DECODE_EPLB_CONFIG:-}" ULTRA_DECODE_ENFORCE_EAGER="${ULTRA_DECODE_ENFORCE_EAGER:-1}" +ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING="${ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING:-0}" ULTRA_NUM_SPECULATIVE_TOKENS="${ULTRA_NUM_SPECULATIVE_TOKENS:-5}" export SAFETENSORS_FAST_GPU=1 @@ -60,6 +64,14 @@ case "$ULTRA_PREFILL_ENFORCE_EAGER" in ;; esac +case "$ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING" in + 0 | 1) ;; + *) + echo "ERROR: ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING must be 0 or 1; got '$ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING'." >&2 + exit 2 + ;; +esac + case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in 0) ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=false @@ -98,6 +110,14 @@ case "$ULTRA_DECODE_ALL2ALL_BACKEND" in ;; esac +case "$ULTRA_DECODE_ENABLE_EPLB" in + 0 | 1) ;; + *) + echo "ERROR: ULTRA_DECODE_ENABLE_EPLB must be 0 or 1; got '$ULTRA_DECODE_ENABLE_EPLB'." >&2 + exit 2 + ;; +esac + case "$ULTRA_DECODE_ENFORCE_EAGER" in 0 | 1) ;; *) @@ -106,6 +126,14 @@ case "$ULTRA_DECODE_ENFORCE_EAGER" in ;; esac +case "$ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING" in + 0 | 1) ;; + *) + echo "ERROR: ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING must be 0 or 1; got '$ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING'." >&2 + exit 2 + ;; +esac + if [[ ! "$ULTRA_NUM_SPECULATIVE_TOKENS" =~ ^[1-9][0-9]*$ ]]; then echo "ERROR: ULTRA_NUM_SPECULATIVE_TOKENS must be a positive integer; got '$ULTRA_NUM_SPECULATIVE_TOKENS'." >&2 exit 2 @@ -124,7 +152,6 @@ VLLM_COMMON_ARGS=( --enable-chunked-prefill --kv-cache-dtype fp8 --no-disable-hybrid-kv-cache-manager - --no-async-scheduling --block-size 128 --mamba-cache-mode align --mamba-ssm-cache-dtype float16 @@ -155,6 +182,12 @@ VLLM_PREFILL_ARGS=( --tensor-parallel-size 4 ) +if [[ "$ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING" == "1" ]]; then + VLLM_PREFILL_ARGS+=(--async-scheduling) +else + VLLM_PREFILL_ARGS+=(--no-async-scheduling) +fi + if [[ "$ULTRA_PREFILL_ENFORCE_EAGER" == "1" ]]; then # Bypass torch.compile as well as CUDA graphs on the prefill tier. VLLM_PREFILL_ARGS+=(--enforce-eager) @@ -170,6 +203,13 @@ VLLM_DECODE_ARGS=( --tensor-parallel-size 4 ) +if [[ "$ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING" == "1" ]]; then + # Overlap CPU scheduling for the next decode step with current GPU work. + VLLM_DECODE_ARGS+=(--async-scheduling) +else + VLLM_DECODE_ARGS+=(--no-async-scheduling) +fi + if [[ "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" == "1" ]]; then # Isolate custom all-reduce from the graph-enabled distributed decode path. # vLLM falls back to its standard NCCL-backed tensor-parallel reductions. @@ -182,6 +222,15 @@ if [[ -n "$ULTRA_DECODE_ALL2ALL_BACKEND" ]]; then VLLM_DECODE_ARGS+=(--all2all-backend "$ULTRA_DECODE_ALL2ALL_BACKEND") fi +if [[ "$ULTRA_DECODE_ENABLE_EPLB" == "1" ]]; then + # Rebalance the existing decode experts from observed routing load. Keep + # the stable default all-to-all backend unless it is explicitly overridden. + VLLM_DECODE_ARGS+=(--enable-eplb) + if [[ -n "$ULTRA_DECODE_EPLB_CONFIG" ]]; then + VLLM_DECODE_ARGS+=(--eplb-config "$ULTRA_DECODE_EPLB_CONFIG") + fi +fi + if [[ "$ULTRA_DECODE_ENFORCE_EAGER" == "1" ]]; then # Bypass torch.compile as well as CUDA graphs on the decode tier. VLLM_DECODE_ARGS+=(--enforce-eager) From b894cbefedcd1e1b60bb22b43b16226494eeeeaa Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Tue, 8 Sep 2026 20:27:16 +0000 Subject: [PATCH 05/15] Support Ultra coupled prefill/decode serving Signed-off-by: Frankie Siino --- .../sbatch_external_vllm.sh | 300 +++++++++++++++--- .../vllm_configs/nemotron_3_ultra.sh | 3 +- 2 files changed, 251 insertions(+), 52 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index c970535ac8..eff0aab17a 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -10,11 +10,95 @@ MODEL_NAME="${MODEL_NAME:-$MODEL}" CONTAINER=$CONTAINER MOUNTS=$MOUNTS VLLM_CONFIG=$VLLM_CONFIG +SBATCH_EXCLUDE="${SBATCH_EXCLUDE:-}" +SBATCH_TIME="${SBATCH_TIME:-04:00:00}" +NUM_SAMPLES_IN_PARALLEL="${NUM_SAMPLES_IN_PARALLEL:-}" +NUM_SAMPLES_IN_PARALLEL_ARG="" +RESUME_EVAL_ON_REQUEUE="${RESUME_EVAL_ON_REQUEUE:-0}" +RESUME_FROM_CACHE_ARG="" +# Independent mode starts one complete TP model replica per node. Coupled mode +# forms one multi-node DP/EP engine per tier for models that cannot fit per node. +VLLM_PD_DEPLOYMENT_MODE="${VLLM_PD_DEPLOYMENT_MODE:-independent}" +# auto uses all allocated nodes when sbatch supports --segment; none omits it; +# a positive integer requests that explicit segment size. +VLLM_SLURM_SEGMENT="${VLLM_SLURM_SEGMENT:-auto}" SLURM_COMMENT="${SLURM_COMMENT:-}" OPENSANDBOX_DOMAIN="${OPENSANDBOX_DOMAIN:-}" OPENSANDBOX_API_KEY="${OPENSANDBOX_API_KEY:-}" OPENSANDBOX_PROTOCOL="${OPENSANDBOX_PROTOCOL:-http}" +if [[ -n "$NUM_SAMPLES_IN_PARALLEL" ]]; then + if [[ ! "$NUM_SAMPLES_IN_PARALLEL" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: NUM_SAMPLES_IN_PARALLEL must be a positive integer; got '$NUM_SAMPLES_IN_PARALLEL'." >&2 + exit 1 + fi + NUM_SAMPLES_IN_PARALLEL_ARG="++num_samples_in_parallel=$NUM_SAMPLES_IN_PARALLEL" +fi + +case "$RESUME_EVAL_ON_REQUEUE" in + 0) + ;; + 1) + RESUME_FROM_CACHE_ARG="++resume_from_cache=true" + ;; + *) + echo "ERROR: RESUME_EVAL_ON_REQUEUE must be 0 or 1; got '$RESUME_EVAL_ON_REQUEUE'." >&2 + exit 1 + ;; +esac + +case "$VLLM_PD_DEPLOYMENT_MODE" in + independent | coupled) + ;; + *) + echo "ERROR: VLLM_PD_DEPLOYMENT_MODE must be independent or coupled; got '$VLLM_PD_DEPLOYMENT_MODE'." >&2 + exit 1 + ;; +esac + +SBATCH_EXCLUDE_ARGS=() +if [[ -n "$SBATCH_EXCLUDE" ]]; then + # Revalidate saved exclusion lists so removed nodes cannot invalidate the submission. + if ! requested_excludes=$(scontrol show hostnames "$SBATCH_EXCLUDE"); then + echo "ERROR: Could not expand SBATCH_EXCLUDE='$SBATCH_EXCLUDE'." >&2 + exit 1 + fi + if ! cluster_nodes=$(scontrol show nodes --oneliner); then + echo "ERROR: Could not query Slurm's node inventory." >&2 + exit 1 + fi + + declare -A cluster_node_set=() + while IFS= read -r node_record; do + node=${node_record#NodeName=} + node=${node%% *} + [[ -n "$node" ]] && cluster_node_set["$node"]=1 + done <<< "$cluster_nodes" + + valid_excludes=() + stale_excludes=() + while IFS= read -r node; do + [[ -z "$node" ]] && continue + if [[ -n "${cluster_node_set[$node]:-}" ]]; then + valid_excludes+=("$node") + else + stale_excludes+=("$node") + fi + done <<< "$requested_excludes" + + if (( ${#stale_excludes[@]} )); then + stale_excludes_csv=$(IFS=,; echo "${stale_excludes[*]}") + echo "WARNING: Ignoring exclusions absent from Slurm's current node inventory: $stale_excludes_csv" >&2 + fi + if (( ${#valid_excludes[@]} )); then + valid_excludes_csv=$(IFS=,; echo "${valid_excludes[*]}") + echo "Using ${#valid_excludes[@]} valid node exclusions." + SBATCH_EXCLUDE_ARGS+=(--exclude="$valid_excludes_csv") + else + echo "WARNING: No requested node exclusions exist in Slurm's current node inventory." >&2 + fi +fi + should_run_eval=$(( $# > 0 )) if (( should_run_eval )); then EXPERIMENT_NAME=$EXPERIMENT_NAME @@ -34,6 +118,11 @@ DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT=5700 ROUTER_SERVER_PORT=8000 WORKER_SERVER_PORT=8001 +PREFILL_SERVER_PORT=8001 +DECODE_SERVER_PORT=8002 + +PREFILL_DP_RPC_PORT=13345 +DECODE_DP_RPC_PORT=13346 eval_command=$(cat <_aggregate_metrics.json from this, so the # default timestamped name makes the aggregate unfindable to anything that # did not watch the job run. Override it when results/ is already per-run. @@ -75,6 +169,7 @@ gym eval run \ ++split=benchmark \ ++use_absolute_ip=true \ ++reuse_existing_data_preparation=true \ + $RESUME_FROM_CACHE_ARG \ ++policy_base_url=http://\$(getent hosts "\$ROUTER_NODE" | awk 'NR == 1 {print \$1}'):$ROUTER_SERVER_PORT/v1 \ ++policy_api_key=dummy_api_key \ ++policy_model_name=$MODEL_NAME \ @@ -82,6 +177,7 @@ gym eval run \ ++global_aiohttp_connector_limit_per_host=16384 \ ++port_range_low=63000 \ ++port_range_high=64000 \ + $NUM_SAMPLES_IN_PARALLEL_ARG \ "\${GYM_MODEL_PARAMS[@]}" @@ -127,59 +223,159 @@ if [[ \$(ulimit -Hn) == "unlimited" ]] || [[ 65535 -lt \$(ulimit -Hn) ]]; then fi this_node_hostname=\$(hostname) -if (( SLURM_PROCID == 0 )); then - read -r -a nodes <<< "\$ALL_NODES" - - # @bxyu-nvidia: for --intra-node-data-parallel-size: Not sure what to set this to other than 1. I can't tell from the docs what is appropriate and 1 seems to work fine. - # Set a super long request timeout since some reasoning requests may take a long time to generate. - # Don't manually wait as vllm-router will wait for the URLs to come up - router_args=( \ - --prefill-policy cache_aware \ - --decode-policy cache_aware \ - --balance-abs-threshold 4 \ - --balance-rel-threshold 1.1 \ - --vllm-pd-disaggregation \ - --host \$this_node_hostname \ - --port $ROUTER_SERVER_PORT \ - --intra-node-data-parallel-size 1 \ - --request-timeout-secs 86400 \ - --log-level error - ) - - for (( i = 0; i < $NUM_PREFILL_NODES; i++ )); do - router_args+=(--prefill "http://\${nodes[i]}:$WORKER_SERVER_PORT") - done - for (( i = 0; i < $NUM_DECODE_NODES; i++ )); do - node_idx=\$(( $NUM_PREFILL_NODES + i )) - router_args+=(--decode "http://\${nodes[node_idx]}:$WORKER_SERVER_PORT") - done - - vllm-router "\${router_args[@]}" & - - router_pid=\$! - trap 'kill "\$router_pid" 2>/dev/null || true' EXIT -fi - -# Split nodes here by index -if (( SLURM_PROCID < $NUM_PREFILL_NODES )); then - # Prefill - VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ - VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ - vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ - --host \$this_node_hostname \ - --port $WORKER_SERVER_PORT +read -r -a nodes <<< "\$ALL_NODES" + +if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then + PREFILL_HEAD=\${nodes[0]} + DECODE_HEAD=\${nodes[$NUM_PREFILL_NODES]} + + wait_for_vllm_health() { + local role=\$1 + local url=\$2 + local pid=\$3 + + until curl -fs "\$url" >/dev/null; do + if ! kill -0 "\$pid" 2>/dev/null; then + local status=0 + wait "\$pid" || status=\$? + (( status != 0 )) || status=1 + echo "ERROR: \$role vLLM process exited before becoming healthy (status=\$status)." >&2 + return "\$status" + fi + sleep 5 + done + } + + if (( SLURM_PROCID == 0 )); then + # The first prefill rank owns its tier's API server. The remaining + # prefill ranks run headless so expert parallelism spans the tier. + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $PREFILL_SERVER_PORT \ + --data-parallel-size $NUM_PREFILL_NODES \ + --data-parallel-address \$PREFILL_HEAD \ + --data-parallel-rpc-port $PREFILL_DP_RPC_PORT \ + --api-server-count 1 \ + & + prefill_pid=\$! + trap 'kill "\$prefill_pid" 2>/dev/null || true' EXIT + + wait_for_vllm_health "prefill" "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT/health" "\$prefill_pid" + wait_for_vllm_health "decode" "http://\$DECODE_HEAD:$DECODE_SERVER_PORT/health" "\$prefill_pid" + + vllm-router \ + --prefill-policy cache_aware \ + --decode-policy cache_aware \ + --balance-abs-threshold 4 \ + --balance-rel-threshold 1.1 \ + --vllm-pd-disaggregation \ + --prefill "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT" \ + --decode "http://\$DECODE_HEAD:$DECODE_SERVER_PORT" \ + --host \$PREFILL_HEAD \ + --port $ROUTER_SERVER_PORT \ + --intra-node-data-parallel-size 1 \ + --request-timeout-secs 86400 \ + --log-level error + elif (( SLURM_PROCID < $NUM_PREFILL_NODES )); then + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ + --headless \ + --data-parallel-size $NUM_PREFILL_NODES \ + --data-parallel-start-rank \$SLURM_PROCID \ + --data-parallel-address \$PREFILL_HEAD \ + --data-parallel-rpc-port $PREFILL_DP_RPC_PORT + elif (( SLURM_PROCID == $NUM_PREFILL_NODES )); then + # Decode mirrors prefill with one API rank and headless ranks across + # the other decode nodes. + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $DECODE_SERVER_PORT \ + --data-parallel-size $NUM_DECODE_NODES \ + --data-parallel-address \$DECODE_HEAD \ + --data-parallel-rpc-port $DECODE_DP_RPC_PORT \ + --api-server-count 1 + else + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ + --headless \ + --data-parallel-size $NUM_DECODE_NODES \ + --data-parallel-start-rank \$(( SLURM_PROCID - $NUM_PREFILL_NODES )) \ + --data-parallel-address \$DECODE_HEAD \ + --data-parallel-rpc-port $DECODE_DP_RPC_PORT + fi else - # Decode - VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ - VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ - vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ - --host \$this_node_hostname \ - --port $WORKER_SERVER_PORT + # Preserve main's independent topology for models that fit one complete + # tensor-parallel replica on each node. + if (( SLURM_PROCID == 0 )); then + router_args=( \ + --prefill-policy cache_aware \ + --decode-policy cache_aware \ + --balance-abs-threshold 4 \ + --balance-rel-threshold 1.1 \ + --vllm-pd-disaggregation \ + --host \$this_node_hostname \ + --port $ROUTER_SERVER_PORT \ + --intra-node-data-parallel-size 1 \ + --request-timeout-secs 86400 \ + --log-level error + ) + + for (( i = 0; i < $NUM_PREFILL_NODES; i++ )); do + router_args+=(--prefill "http://\${nodes[i]}:$WORKER_SERVER_PORT") + done + for (( i = 0; i < $NUM_DECODE_NODES; i++ )); do + node_idx=\$(( $NUM_PREFILL_NODES + i )) + router_args+=(--decode "http://\${nodes[node_idx]}:$WORKER_SERVER_PORT") + done + + vllm-router "\${router_args[@]}" & + router_pid=\$! + trap 'kill "\$router_pid" 2>/dev/null || true' EXIT + fi + + if (( SLURM_PROCID < $NUM_PREFILL_NODES )); then + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $WORKER_SERVER_PORT + else + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $WORKER_SERVER_PORT + fi fi EOF ) NUM_NODES=$((NUM_PREFILL_NODES + NUM_DECODE_NODES)) +SBATCH_SEGMENT_ARGS=() +case "$VLLM_SLURM_SEGMENT" in + auto) + sbatch_help=$(sbatch --help 2>&1) + if [[ "$sbatch_help" == *"--segment"* ]]; then + SBATCH_SEGMENT_ARGS+=(--segment="$NUM_NODES") + fi + ;; + none) + ;; + *) + if [[ ! "$VLLM_SLURM_SEGMENT" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: VLLM_SLURM_SEGMENT must be auto, none, or a positive integer." >&2 + exit 2 + fi + SBATCH_SEGMENT_ARGS+=(--segment="$VLLM_SLURM_SEGMENT") + ;; +esac + batch_command=$(cat < 0 otherwise the engine will hang on the second or third engine step. +# This cluster needs --segment > 0 to avoid distributed engine hangs. Keep the +# setting caller-configurable because coupled tiers benefit from tier-sized segments. submit_dir=$(pwd -P) # An exported connection is sent as arguments; otherwise env.yaml is read. if [[ -n "$OPENSANDBOX_DOMAIN" ]]; then @@ -275,13 +472,14 @@ main_job_id=$( sbatch \ --parsable \ --nodes=$NUM_NODES \ - --time=04:00:00 \ + --time="$SBATCH_TIME" \ + "${SBATCH_EXCLUDE_ARGS[@]}" \ + "${SBATCH_SEGMENT_ARGS[@]}" \ --job-name=gym-$EXPERIMENT_NAME-$USER \ --output=slurm-logs/%j-%x.log \ --ntasks-per-node=1 \ --comment="$SLURM_COMMENT" \ --exclusive \ - --segment=$NUM_NODES \ --wrap 'exec bash -c "$batch_command"' ) main_job_id=${main_job_id%%;*} diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh index e952ea91c3..8f0c8b24c9 100644 --- a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -4,7 +4,8 @@ # Nemotron 3 Ultra BF16 baseline for disaggregated prefill/decode on 4-GPU # GB200 nodes. This launcher's TP4-per-node layout needs four coupled DP ranks -# per tier so expert parallelism can shard the 512 experts over 16 GPUs. +# per tier so expert parallelism can shard the 512 experts over 16 GPUs. Launch +# this config with VLLM_PD_DEPLOYMENT_MODE=coupled; it cannot fit in DP=1 mode. # InstantTensor's io_uring loader returned EIO while 24 ranks concurrently read # the Lustre-hosted checkpoint in Run 001. Keep the loader selectable for later From f6387b10b94bc5d6a23b02745f5b4f2d66d89eb4 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Wed, 9 Sep 2026 18:52:13 +0000 Subject: [PATCH 06/15] Clean up Ultra vLLM tuning config Remove obsolete experiment switches and retain the validated MTP3 defaults. Simplify Slurm argument handling and correct the coupled decode health check. Signed-off-by: Frankie Siino --- .../sbatch_external_vllm.sh | 86 ++------- .../vllm_configs/nemotron_3_ultra.sh | 180 ++++-------------- 2 files changed, 48 insertions(+), 218 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index eff0aab17a..3d6173d38b 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -10,7 +10,6 @@ MODEL_NAME="${MODEL_NAME:-$MODEL}" CONTAINER=$CONTAINER MOUNTS=$MOUNTS VLLM_CONFIG=$VLLM_CONFIG -SBATCH_EXCLUDE="${SBATCH_EXCLUDE:-}" SBATCH_TIME="${SBATCH_TIME:-04:00:00}" NUM_SAMPLES_IN_PARALLEL="${NUM_SAMPLES_IN_PARALLEL:-}" NUM_SAMPLES_IN_PARALLEL_ARG="" @@ -19,9 +18,9 @@ RESUME_FROM_CACHE_ARG="" # Independent mode starts one complete TP model replica per node. Coupled mode # forms one multi-node DP/EP engine per tier for models that cannot fit per node. VLLM_PD_DEPLOYMENT_MODE="${VLLM_PD_DEPLOYMENT_MODE:-independent}" -# auto uses all allocated nodes when sbatch supports --segment; none omits it; -# a positive integer requests that explicit segment size. -VLLM_SLURM_SEGMENT="${VLLM_SLURM_SEGMENT:-auto}" +# Empty uses one segment containing all allocated nodes, matching main's +# behavior. Coupled deployments can override this with their tier size. +VLLM_SLURM_SEGMENT="${VLLM_SLURM_SEGMENT:-}" SLURM_COMMENT="${SLURM_COMMENT:-}" OPENSANDBOX_DOMAIN="${OPENSANDBOX_DOMAIN:-}" OPENSANDBOX_API_KEY="${OPENSANDBOX_API_KEY:-}" @@ -56,49 +55,6 @@ case "$VLLM_PD_DEPLOYMENT_MODE" in ;; esac -SBATCH_EXCLUDE_ARGS=() -if [[ -n "$SBATCH_EXCLUDE" ]]; then - # Revalidate saved exclusion lists so removed nodes cannot invalidate the submission. - if ! requested_excludes=$(scontrol show hostnames "$SBATCH_EXCLUDE"); then - echo "ERROR: Could not expand SBATCH_EXCLUDE='$SBATCH_EXCLUDE'." >&2 - exit 1 - fi - if ! cluster_nodes=$(scontrol show nodes --oneliner); then - echo "ERROR: Could not query Slurm's node inventory." >&2 - exit 1 - fi - - declare -A cluster_node_set=() - while IFS= read -r node_record; do - node=${node_record#NodeName=} - node=${node%% *} - [[ -n "$node" ]] && cluster_node_set["$node"]=1 - done <<< "$cluster_nodes" - - valid_excludes=() - stale_excludes=() - while IFS= read -r node; do - [[ -z "$node" ]] && continue - if [[ -n "${cluster_node_set[$node]:-}" ]]; then - valid_excludes+=("$node") - else - stale_excludes+=("$node") - fi - done <<< "$requested_excludes" - - if (( ${#stale_excludes[@]} )); then - stale_excludes_csv=$(IFS=,; echo "${stale_excludes[*]}") - echo "WARNING: Ignoring exclusions absent from Slurm's current node inventory: $stale_excludes_csv" >&2 - fi - if (( ${#valid_excludes[@]} )); then - valid_excludes_csv=$(IFS=,; echo "${valid_excludes[*]}") - echo "Using ${#valid_excludes[@]} valid node exclusions." - SBATCH_EXCLUDE_ARGS+=(--exclude="$valid_excludes_csv") - else - echo "WARNING: No requested node exclusions exist in Slurm's current node inventory." >&2 - fi -fi - should_run_eval=$(( $# > 0 )) if (( should_run_eval )); then EXPERIMENT_NAME=$EXPERIMENT_NAME @@ -232,12 +188,12 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then wait_for_vllm_health() { local role=\$1 local url=\$2 - local pid=\$3 + local local_pid=\${3:-} until curl -fs "\$url" >/dev/null; do - if ! kill -0 "\$pid" 2>/dev/null; then + if [[ -n "\$local_pid" ]] && ! kill -0 "\$local_pid" 2>/dev/null; then local status=0 - wait "\$pid" || status=\$? + wait "\$local_pid" || status=\$? (( status != 0 )) || status=1 echo "ERROR: \$role vLLM process exited before becoming healthy (status=\$status)." >&2 return "\$status" @@ -263,7 +219,9 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then trap 'kill "\$prefill_pid" 2>/dev/null || true' EXIT wait_for_vllm_health "prefill" "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT/health" "\$prefill_pid" - wait_for_vllm_health "decode" "http://\$DECODE_HEAD:$DECODE_SERVER_PORT/health" "\$prefill_pid" + # The decode API is owned by another Slurm rank. The enclosing srun's + # --kill-on-bad-exit handles an early decode-process failure. + wait_for_vllm_health "decode" "http://\$DECODE_HEAD:$DECODE_SERVER_PORT/health" vllm-router \ --prefill-policy cache_aware \ @@ -357,24 +315,11 @@ EOF ) NUM_NODES=$((NUM_PREFILL_NODES + NUM_DECODE_NODES)) -SBATCH_SEGMENT_ARGS=() -case "$VLLM_SLURM_SEGMENT" in - auto) - sbatch_help=$(sbatch --help 2>&1) - if [[ "$sbatch_help" == *"--segment"* ]]; then - SBATCH_SEGMENT_ARGS+=(--segment="$NUM_NODES") - fi - ;; - none) - ;; - *) - if [[ ! "$VLLM_SLURM_SEGMENT" =~ ^[1-9][0-9]*$ ]]; then - echo "ERROR: VLLM_SLURM_SEGMENT must be auto, none, or a positive integer." >&2 - exit 2 - fi - SBATCH_SEGMENT_ARGS+=(--segment="$VLLM_SLURM_SEGMENT") - ;; -esac +VLLM_SLURM_SEGMENT="${VLLM_SLURM_SEGMENT:-$NUM_NODES}" +if [[ ! "$VLLM_SLURM_SEGMENT" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: VLLM_SLURM_SEGMENT must be a positive integer." >&2 + exit 2 +fi batch_command=$(cat <&2 - exit 2 - ;; -esac - -case "$ULTRA_PREFILL_ENFORCE_EAGER" in - 0 | 1) ;; - *) - echo "ERROR: ULTRA_PREFILL_ENFORCE_EAGER must be 0 or 1; got '$ULTRA_PREFILL_ENFORCE_EAGER'." >&2 - exit 2 - ;; -esac - -case "$ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING" in - 0 | 1) ;; - *) - echo "ERROR: ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING must be 0 or 1; got '$ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING'." >&2 - exit 2 - ;; -esac - case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in 0) ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=false @@ -86,39 +45,6 @@ case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in ;; esac -ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES_JSON="" -if [[ -n "$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES" ]]; then - if [[ ! "$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES" =~ ^\[[0-9]+(,[0-9]+)*\]$ ]]; then - echo "ERROR: ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES must be a compact JSON integer array; got '$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES'." >&2 - exit 2 - fi - ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES_JSON=",\"cudagraph_capture_sizes\":$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES" -fi - -case "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" in - 0 | 1) ;; - *) - echo "ERROR: ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE must be 0 or 1; got '$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE'." >&2 - exit 2 - ;; -esac - -case "$ULTRA_DECODE_ALL2ALL_BACKEND" in - "" | allgather_reducescatter | deepep_high_throughput | deepep_low_latency | deepep_v2 | flashinfer_all2allv | flashinfer_nvlink_one_sided | flashinfer_nvlink_two_sided | mori_high_throughput | mori_low_latency | nixl_ep) ;; - *) - echo "ERROR: unsupported ULTRA_DECODE_ALL2ALL_BACKEND '$ULTRA_DECODE_ALL2ALL_BACKEND'." >&2 - exit 2 - ;; -esac - -case "$ULTRA_DECODE_ENABLE_EPLB" in - 0 | 1) ;; - *) - echo "ERROR: ULTRA_DECODE_ENABLE_EPLB must be 0 or 1; got '$ULTRA_DECODE_ENABLE_EPLB'." >&2 - exit 2 - ;; -esac - case "$ULTRA_DECODE_ENFORCE_EAGER" in 0 | 1) ;; *) @@ -127,10 +53,10 @@ case "$ULTRA_DECODE_ENFORCE_EAGER" in ;; esac -case "$ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING" in +case "$ULTRA_ENABLE_MTP" in 0 | 1) ;; *) - echo "ERROR: ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING must be 0 or 1; got '$ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING'." >&2 + echo "ERROR: ULTRA_ENABLE_MTP must be 0 or 1; got '$ULTRA_ENABLE_MTP'." >&2 exit 2 ;; esac @@ -160,86 +86,46 @@ VLLM_COMMON_ARGS=( --enable-mamba-cache-stochastic-rounding --mamba-cache-philox-rounds 5 --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 96}' - --load-format "$ULTRA_LOAD_FORMAT" + --load-format safetensors --enable-expert-parallel --distributed-timeout-seconds 3600 + # NIXL-transferred Mamba state must not coexist with locally retained + # prefix-cache blocks; doing so triggers the multiple-local-block assertion. + --no-enable-prefix-caching ) -if [[ "$ULTRA_ENABLE_PREFIX_CACHING" == "1" ]]; then - VLLM_COMMON_ARGS+=(--enable-prefix-caching) -else - # vLLM 0.25.1's NIXL pull connector cannot reconcile multiple locally - # prefix-cached Mamba/SSM blocks with transferred prefill state. - VLLM_COMMON_ARGS+=(--no-enable-prefix-caching) -fi - VLLM_PREFILL_ARGS=( --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}' - --compilation-config "{\"cudagraph_copy_inputs\":$ULTRA_PREFILL_CUDAGRAPH_COPY_INPUTS_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" --gpu-memory-utilization "$ULTRA_PREFILL_GPU_MEMORY_UTILIZATION" --max-num-batched-tokens "$ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS" --max-num-seqs "$ULTRA_MAX_NUM_SEQS" --data-parallel-size-local 1 --tensor-parallel-size 4 + --no-async-scheduling + # Eager prefill avoids the compiled/CUDA-graph stalls observed during tuning. + --enforce-eager ) -if [[ "$ULTRA_PREFILL_ENABLE_ASYNC_SCHEDULING" == "1" ]]; then - VLLM_PREFILL_ARGS+=(--async-scheduling) -else - VLLM_PREFILL_ARGS+=(--no-async-scheduling) -fi - -if [[ "$ULTRA_PREFILL_ENFORCE_EAGER" == "1" ]]; then - # Bypass torch.compile as well as CUDA graphs on the prefill tier. - VLLM_PREFILL_ARGS+=(--enforce-eager) -fi - VLLM_DECODE_ARGS=( --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}' - --compilation-config "{\"cudagraph_mode\":\"$ULTRA_DECODE_CUDAGRAPH_MODE\",\"cudagraph_copy_inputs\":$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON$ULTRA_DECODE_CUDAGRAPH_CAPTURE_SIZES_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" + --compilation-config "{\"cudagraph_mode\":\"$ULTRA_DECODE_CUDAGRAPH_MODE\",\"cudagraph_copy_inputs\":$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" --gpu-memory-utilization "$ULTRA_DECODE_GPU_MEMORY_UTILIZATION" --max-num-batched-tokens "$ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS" --max-num-seqs "$ULTRA_MAX_NUM_SEQS" --data-parallel-size-local 1 --tensor-parallel-size 4 + --no-async-scheduling ) -if [[ "$ULTRA_DECODE_ENABLE_ASYNC_SCHEDULING" == "1" ]]; then - # Overlap CPU scheduling for the next decode step with current GPU work. - VLLM_DECODE_ARGS+=(--async-scheduling) -else - VLLM_DECODE_ARGS+=(--no-async-scheduling) -fi - -if [[ "$ULTRA_DISABLE_DECODE_CUSTOM_ALL_REDUCE" == "1" ]]; then - # Isolate custom all-reduce from the graph-enabled distributed decode path. - # vLLM falls back to its standard NCCL-backed tensor-parallel reductions. - VLLM_DECODE_ARGS+=(--disable-custom-all-reduce) -fi - -if [[ -n "$ULTRA_DECODE_ALL2ALL_BACKEND" ]]; then - # Select the collective used to dispatch tokens to the expert-parallel - # decode ranks and combine the expert outputs. Prefill stays on the default. - VLLM_DECODE_ARGS+=(--all2all-backend "$ULTRA_DECODE_ALL2ALL_BACKEND") -fi - -if [[ "$ULTRA_DECODE_ENABLE_EPLB" == "1" ]]; then - # Rebalance the existing decode experts from observed routing load. Keep - # the stable default all-to-all backend unless it is explicitly overridden. - VLLM_DECODE_ARGS+=(--enable-eplb) - if [[ -n "$ULTRA_DECODE_EPLB_CONFIG" ]]; then - VLLM_DECODE_ARGS+=(--eplb-config "$ULTRA_DECODE_EPLB_CONFIG") - fi -fi - if [[ "$ULTRA_DECODE_ENFORCE_EAGER" == "1" ]]; then - # Bypass torch.compile as well as CUDA graphs on the decode tier. + # This fallback disables compilation and CUDA graphs for decode. VLLM_DECODE_ARGS+=(--enforce-eager) fi -if [[ "${ULTRA_ENABLE_MTP:-1}" == "1" ]]; then - # Producer and consumer must expose matching cache layouts for NIXL KV transfer. - ULTRA_SPECULATIVE_CONFIG="{\"method\":\"nemotron_h_mtp\",\"num_speculative_tokens\":$ULTRA_NUM_SPECULATIVE_TOKENS}" +if [[ "$ULTRA_ENABLE_MTP" == "1" ]]; then + # Prefill and decode must use the same speculative width so the transferred + # cache layouts agree. + ULTRA_SPECULATIVE_CONFIG="{\"method\":\"mtp\",\"num_speculative_tokens\":$ULTRA_NUM_SPECULATIVE_TOKENS}" VLLM_PREFILL_ARGS+=(--speculative-config "$ULTRA_SPECULATIVE_CONFIG") VLLM_DECODE_ARGS+=(--speculative-config "$ULTRA_SPECULATIVE_CONFIG") fi From e1628e9ca3f0ee1ab4d8d0d1a180b04edb4755e1 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Wed, 9 Sep 2026 20:30:20 +0000 Subject: [PATCH 07/15] Fix vLLM startup monitoring and override precedence Detect prefill failures while waiting for decode readiness. Honor explicit concurrency and resume overrides over environment defaults. Signed-off-by: Frankie Siino --- .../sbatch_external_vllm.sh | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index 3d6173d38b..2ca3a182dc 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -46,6 +46,19 @@ case "$RESUME_EVAL_ON_REQUEUE" in ;; esac +# Environment controls are defaults; explicit Hydra overrides belong to Gym. +# Avoid adding a duplicate setting that could overwrite the caller's value. +for eval_arg in "$@"; do + case "$eval_arg" in + num_samples_in_parallel=* | +num_samples_in_parallel=* | ++num_samples_in_parallel=*) + NUM_SAMPLES_IN_PARALLEL_ARG="" + ;; + resume_from_cache=* | +resume_from_cache=* | ++resume_from_cache=*) + RESUME_FROM_CACHE_ARG="" + ;; + esac +done + case "$VLLM_PD_DEPLOYMENT_MODE" in independent | coupled) ;; @@ -189,15 +202,19 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then local role=\$1 local url=\$2 local local_pid=\${3:-} + local local_role=\${4:-\$role} - until curl -fs "\$url" >/dev/null; do + while true; do if [[ -n "\$local_pid" ]] && ! kill -0 "\$local_pid" 2>/dev/null; then local status=0 wait "\$local_pid" || status=\$? (( status != 0 )) || status=1 - echo "ERROR: \$role vLLM process exited before becoming healthy (status=\$status)." >&2 + echo "ERROR: \$local_role vLLM process exited while waiting for \$role health (status=\$status)." >&2 return "\$status" fi + if curl -fs "\$url" >/dev/null; then + return 0 + fi sleep 5 done } @@ -219,9 +236,9 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then trap 'kill "\$prefill_pid" 2>/dev/null || true' EXIT wait_for_vllm_health "prefill" "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT/health" "\$prefill_pid" - # The decode API is owned by another Slurm rank. The enclosing srun's - # --kill-on-bad-exit handles an early decode-process failure. - wait_for_vllm_health "decode" "http://\$DECODE_HEAD:$DECODE_SERVER_PORT/health" + # Keep watching the local prefill process while the remote decode API + # starts. The enclosing srun handles failures on the decode ranks. + wait_for_vllm_health "decode" "http://\$DECODE_HEAD:$DECODE_SERVER_PORT/health" "\$prefill_pid" "prefill" vllm-router \ --prefill-policy cache_aware \ From a6c2a443a4d2999faceed5b1af905cfd4a035f59 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Wed, 9 Sep 2026 21:07:06 +0000 Subject: [PATCH 08/15] Restore comment from main Signed-off-by: Frankie Siino --- benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index 2ca3a182dc..6f8bde87d5 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -288,6 +288,9 @@ else # Preserve main's independent topology for models that fit one complete # tensor-parallel replica on each node. if (( SLURM_PROCID == 0 )); then + # @bxyu-nvidia: for --intra-node-data-parallel-size: Not sure what to set this to other than 1. I can't tell from the docs what is appropriate and 1 seems to work fine. + # Set a super long request timeout since some reasoning requests may take a long time to generate. + # Don't manually wait as vllm-router will wait for the URLs to come up router_args=( \ --prefill-policy cache_aware \ --decode-policy cache_aware \ From b1da0cae1706b6ccdd56f5955f26ea8a3f228aa2 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Thu, 10 Sep 2026 22:08:57 +0000 Subject: [PATCH 09/15] Bound coupled vLLM startup health checks Add connection and request timeouts so stalled probes cannot block prefill failure detection. Signed-off-by: Frankie Siino --- benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index 6f8bde87d5..ed683795c6 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -212,7 +212,9 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then echo "ERROR: \$local_role vLLM process exited while waiting for \$role health (status=\$status)." >&2 return "\$status" fi - if curl -fs "\$url" >/dev/null; then + # Bound each probe so a stalled endpoint cannot block process checks. + # Timeouts retry below; they do not limit overall model startup time. + if curl -fs --connect-timeout 5 --max-time 10 "\$url" >/dev/null; then return 0 fi sleep 5 From 0ad44781e7a86e0ca09c73b8a2c1cc1dd164e174 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Thu, 10 Sep 2026 22:24:09 +0000 Subject: [PATCH 10/15] Detect coupled vLLM failures after startup Monitor prefill and router processes, propagate unexpected exits, and stop the surviving process. Signed-off-by: Frankie Siino --- .../sbatch_external_vllm.sh | 35 +++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index ed683795c6..a1579af5c3 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -235,7 +235,18 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then --api-server-count 1 \ & prefill_pid=\$! - trap 'kill "\$prefill_pid" 2>/dev/null || true' EXIT + coupled_pids=("\$prefill_pid") + cleanup_coupled_head() { + local status=\$? + trap - EXIT INT TERM + # Signal both local services without delaying failure propagation; + # the enclosing srun tears down the remaining distributed workers. + kill "\${coupled_pids[@]}" 2>/dev/null || true + exit "\$status" + } + trap cleanup_coupled_head EXIT + trap 'exit 130' INT + trap 'exit 143' TERM wait_for_vllm_health "prefill" "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT/health" "\$prefill_pid" # Keep watching the local prefill process while the remote decode API @@ -254,7 +265,27 @@ if [[ "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then --port $ROUTER_SERVER_PORT \ --intra-node-data-parallel-size 1 \ --request-timeout-secs 86400 \ - --log-level error + --log-level error & + router_pid=\$! + coupled_pids+=("\$router_pid") + + # Keep monitoring after readiness. Polling also catches children that + # exited before monitoring started, which wait -n can otherwise miss. + while kill -0 "\$prefill_pid" 2>/dev/null && kill -0 "\$router_pid" 2>/dev/null; do + sleep 1 + done + failed_role=prefill + failed_pid=\$prefill_pid + if kill -0 "\$prefill_pid" 2>/dev/null; then + failed_role=router + failed_pid=\$router_pid + fi + failed_status=0 + wait "\$failed_pid" || failed_status=\$? + # Neither service should exit by itself, even with a zero exit status. + (( failed_status != 0 )) || failed_status=1 + echo "ERROR: \$failed_role process exited after startup (status=\$failed_status)." >&2 + exit "\$failed_status" elif (( SLURM_PROCID < $NUM_PREFILL_NODES )); then VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ From 55933f19bdc1dd3349b841527059742747c5d4f1 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Thu, 10 Sep 2026 22:59:23 +0000 Subject: [PATCH 11/15] Add regression tests for the Super vLLM launcher Cover override precedence, resume paths, health-check timeouts, and coupled-mode failure detection and cancellation cleanup. Signed-off-by: Frankie Siino --- tests/unit_tests/test_super_vllm_launcher.py | 359 +++++++++++++++++++ 1 file changed, 359 insertions(+) create mode 100644 tests/unit_tests/test_super_vllm_launcher.py diff --git a/tests/unit_tests/test_super_vllm_launcher.py b/tests/unit_tests/test_super_vllm_launcher.py new file mode 100644 index 0000000000..2458740baf --- /dev/null +++ b/tests/unit_tests/test_super_vllm_launcher.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import signal +import subprocess +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + + +SCRIPT = Path(__file__).resolve().parents[2] / "benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh" + + +class TestSuperVllmLauncher(unittest.TestCase): + def setUp(self): + # Never inherit cluster credentials, tuning overrides, or real sbatch commands. + self.env = { + "PATH": os.defpath, + "USER": "launcher-test", + "MODEL": "/test/model", + "CONTAINER": "/test/image.sqsh", + "MOUNTS": "/test:/test", + "VLLM_CONFIG": "/dev/null", + "EXPERIMENT_NAME": "launcher-test", + "NUM_PREFILL_NODES": "4", + "NUM_DECODE_NODES": "4", + "SLURM_PROCID": "0", + "SLURM_JOB_ID": "12345", + "SLURM_JOB_USER": "launcher-test", + "ROUTER_NODE": "node0", + "ALL_NODES": "node0 node1 node2 node3 node4 node5 node6 node7", + } + + def run_shell(self, command, *args, env=None): + proc = subprocess.Popen( + ["bash", "--noprofile", "--norc", "-c", command, "launcher-test", *args], + env=self.env | (env or {}), + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=5) + except subprocess.TimeoutExpired: + # The failure-path tests must terminate instead of leaking a polling worker. + os.killpg(proc.pid, signal.SIGKILL) + stdout, stderr = proc.communicate() + self.fail(f"Launcher did not terminate. stdout={stdout!r}, stderr={stderr!r}") + return proc.returncode, stdout, stderr + + def generate_commands(self, *overrides, env=None): + # Capture both commands using a shell function; never submit a Slurm job. + stub = r""" +sbatch() { + if [[ -n "${vllm_command:-}" ]]; then + printf '%s\0' "$eval_command" "$vllm_command" >&2 + fi + printf '12345\n' +} +source "$@" +""" + status, _, captured = self.run_shell(stub, str(SCRIPT), "--config", "benchmark.yaml", *overrides, env=env) + self.assertEqual(status, 0, captured) + eval_command, pd_command, terminator = captured.split("\0") + self.assertEqual(terminator, "") + return eval_command, pd_command + + def eval_arguments(self, *overrides, env=None): + command, _ = self.generate_commands(*overrides, env=env) + command = command.replace("source /opt/Gym_venv/bin/activate", ":").replace("cd /opt/Gym\n", ":\n") + stubs = r""" +gym() { + if [[ "$2" == run ]]; then printf '%s\0' "$@"; fi +} +date() { printf '20260909_120000\n'; } +getent() { printf '10.0.0.1 node0\n'; } +""" + status, stdout, stderr = self.run_shell(stubs + command, env=env) + self.assertEqual(status, 0, stderr) + return stdout.rstrip("\0").split("\0") + + def settings(self, args, key): + return [arg for arg in args if arg.lstrip("+").startswith(key + "=")] + + def test_environment_fallbacks_are_preserved(self): + """Apply concurrency and resume environment defaults, including the stable results path.""" + args = self.eval_arguments(env={"NUM_SAMPLES_IN_PARALLEL": "16", "RESUME_EVAL_ON_REQUEUE": "1"}) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), ["++num_samples_in_parallel=16"]) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn("++output_jsonl_fpath=results/launcher-test/resumable.jsonl", args) + + def test_no_environment_defaults_leave_gym_settings_untouched(self): + """Preserve Gym settings and timestamped output naming when environment overrides are unset.""" + args = self.eval_arguments() + self.assertEqual(self.settings(args, "num_samples_in_parallel"), []) + self.assertEqual(self.settings(args, "resume_from_cache"), []) + self.assertIn("++output_jsonl_fpath=results/launcher-test/slurm_job_id_12345/date_20260909_120000.jsonl", args) + + def test_explicit_concurrency_overrides_environment(self): + """Let explicit concurrency arguments override the environment without duplicate settings.""" + for prefix in ("", "+", "++"): + with self.subTest(prefix=prefix): + override = prefix + "num_samples_in_parallel=512" + args = self.eval_arguments(override, env={"NUM_SAMPLES_IN_PARALLEL": "16"}) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), [override]) + + def test_explicit_resume_overrides_environment(self): + """Honor explicit resume arguments while retaining the switch's stable output naming.""" + for prefix in ("", "+", "++"): + for value in ("true", "false"): + with self.subTest(prefix=prefix, value=value): + override = prefix + "resume_from_cache=" + value + args = self.eval_arguments(override, env={"RESUME_EVAL_ON_REQUEUE": "1"}) + self.assertEqual(self.settings(args, "resume_from_cache"), [override]) + self.assertIn("++output_jsonl_fpath=results/launcher-test/resumable.jsonl", args) + + def test_resume_switch_keeps_output_path_across_job_ids(self): + """Keep resumable output and W&B names stable when the Slurm job ID changes.""" + for job_id in ("12345", "12346"): + with self.subTest(job_id=job_id): + args = self.eval_arguments(env={"RESUME_EVAL_ON_REQUEUE": "1", "SLURM_JOB_ID": job_id}) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn("++output_jsonl_fpath=results/launcher-test/resumable.jsonl", args) + self.assertIn("+wandb_name=launcher-test/resumable", args) + + def test_invalid_resume_switch_is_rejected_before_submission(self): + """Reject an invalid resume switch before any Slurm job is submitted.""" + stub = r""" +sbatch() { printf 'unexpected-submission\n'; } +source "$@" +""" + status, stdout, stderr = self.run_shell(stub, str(SCRIPT), env={"RESUME_EVAL_ON_REQUEUE": "invalid"}) + self.assertNotEqual(status, 0) + self.assertNotIn("unexpected-submission", stdout) + self.assertIn("RESUME_EVAL_ON_REQUEUE must be 0 or 1", stderr) + + def test_explicit_resume_keeps_output_path_across_job_ids(self): + """Preserve a manually supplied resume path while keeping the default per-job run naming.""" + for job_id in ("12345", "12346"): + with self.subTest(job_id=job_id): + args = self.eval_arguments( + "++resume_from_cache=true", + env={"ROLLOUTS_FPATH": "results/existing/resumable.jsonl", "SLURM_JOB_ID": job_id}, + ) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn("++output_jsonl_fpath=results/existing/resumable.jsonl", args) + self.assertIn(f"+wandb_name=launcher-test/slurm_job_id_{job_id}/date_20260909_120000", args) + + def test_unrelated_override_does_not_suppress_concurrency_default(self): + """Keep global environment defaults when overrides only target nested agent settings.""" + args = self.eval_arguments( + "++agent.num_samples_in_parallel=8", + "++agent.resume_from_cache=false", + env={"NUM_SAMPLES_IN_PARALLEL": "16", "RESUME_EVAL_ON_REQUEUE": "1"}, + ) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), ["++num_samples_in_parallel=16"]) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn("++agent.num_samples_in_parallel=8", args) + self.assertIn("++agent.resume_from_cache=false", args) + + def test_prefill_exit_is_detected_during_both_health_checks(self): + """Detect prefill exits during either startup health check, including after request timeouts.""" + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +vllm() { command sleep 0.1; return "$TEST_PREFILL_STATUS"; } +curl() { + if [[ "$TEST_WAIT_ROLE" == decode && "$*" == *':8001/health'* ]]; then + return 0 + fi + return "$TEST_HEALTH_STATUS" +} +sleep() { command sleep 0.01; } +hostname() { printf 'node0\n'; } +vllm-router() { printf 'unexpected-router-start\n'; } +""" + for role in ("prefill", "decode"): + for exit_status in (0, 7): + # curl returns 28 for a timed-out request. That must not hide the + # prefill exit status or abort the polling loop prematurely. + for health_status in (1, 28): + with self.subTest(role=role, exit_status=exit_status, health_status=health_status): + status, stdout, stderr = self.run_shell( + stubs + command, + env={ + "TEST_WAIT_ROLE": role, + "TEST_PREFILL_STATUS": str(exit_status), + "TEST_HEALTH_STATUS": str(health_status), + }, + ) + expected_status = exit_status or 1 + self.assertEqual(status, expected_status, stderr) + self.assertNotIn("unexpected-router-start", stdout) + self.assertIn( + f"prefill vLLM process exited while waiting for {role} health (status={expected_status})", + stderr, + ) + + def test_coupled_health_probes_have_timeouts_and_retry_until_ready(self): + """Bound both tiers' health probes and retry timed-out requests until the servers are ready.""" + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +vllm() { while true; do command sleep 0.01; done; } +probe_count=0 +curl() { + printf '%s\0' "$@" >&2 + printf '\0' >&2 + probe_count=$((probe_count + 1)) + # Both tiers time out once, then become healthy while prefill stays alive. + if (( probe_count % 2 )); then return 28; fi + return 0 +} +sleep() { :; } +hostname() { printf 'node0\n'; } +vllm-router() { printf 'router-started\n'; kill -TERM "$$"; } +""" + status, stdout, stderr = self.run_shell(stubs + command) + self.assertEqual(status, 143, stderr) + self.assertEqual(stdout, "router-started\n") + probes = [probe.split("\0") for probe in stderr.rstrip("\0").split("\0\0")] + self.assertEqual(len(probes), 4) + self.assertEqual( + [args[-1] for args in probes], + ["http://node0:8001/health"] * 2 + ["http://node4:8002/health"] * 2, + ) + for args in probes: + self.assertIn("--connect-timeout", args) + self.assertEqual(args[args.index("--connect-timeout") + 1], "5") + self.assertIn("--max-time", args) + self.assertEqual(args[args.index("--max-time") + 1], "10") + + def test_healthy_coupled_servers_start_router(self): + """Start the router with the correct prefill and decode URLs once both servers are healthy.""" + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +vllm() { while true; do command sleep 0.01; done; } +curl() { return 0; } +hostname() { printf 'node0\n'; } +vllm-router() { printf '%s\0' "$@"; kill -TERM "$$"; } +""" + status, stdout, stderr = self.run_shell(stubs + command) + self.assertEqual(status, 143, stderr) + args = stdout.rstrip("\0").split("\0") + self.assertEqual(args[args.index("--prefill") + 1], "http://node0:8001") + self.assertEqual(args[args.index("--decode") + 1], "http://node4:8002") + + def run_coupled_lifecycle( + self, *, exit_role="", exit_status=0, shutdown_signal="", startup_shutdown=False, exit_before_monitoring=False + ): + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +run_service() { + local role=$1 + trap 'printf "%s-stopped\n" "$role"; exit 0' TERM + trap 'exit "$TEST_EXIT_STATUS"' USR1 + printf '%s-started\n' "$role" + touch "$TEST_STATE_DIR/$role-ready" + if [[ "$role" == "$TEST_EXIT_ROLE" ]]; then + # Synchronize on real process startup, not an assumed sleep duration. + while [[ ! -f "$TEST_STATE_DIR/prefill-ready" || ! -f "$TEST_STATE_DIR/router-ready" ]]; do + command sleep 0.01 + done + printf '%s-exited\n' "$role" + return "$TEST_EXIT_STATUS" + fi + if [[ "$role" == router && -n "$TEST_SHUTDOWN_SIGNAL" ]]; then + # $$ remains the supervising shell's PID inside a background function. + kill -s "$TEST_SHUTDOWN_SIGNAL" "$$" + fi + while true; do command sleep 0.01; done +} +vllm() { run_service prefill; } +vllm-router() { run_service router; } +curl() { + [[ -f "$TEST_STATE_DIR/prefill-ready" ]] || return 1 + if [[ "$TEST_STARTUP_SHUTDOWN" == 1 ]]; then + kill -s "$TEST_SHUTDOWN_SIGNAL" "$$" + fi + if [[ "$TEST_EXIT_BEFORE_MONITORING" == 1 && "$*" == *':8002/health'* ]]; then + # Make prefill exit before the final health request returns success. + kill -USR1 "$prefill_pid" + wait "$prefill_pid" || true + fi + return 0 +} +sleep() { command sleep 0.01; } +hostname() { printf 'node0\n'; } +""" + with TemporaryDirectory(prefix="gym-coupled-lifecycle-") as state_dir: + return self.run_shell( + stubs + command, + env={ + "TEST_STATE_DIR": state_dir, + "TEST_EXIT_ROLE": exit_role, + "TEST_EXIT_STATUS": str(exit_status), + "TEST_SHUTDOWN_SIGNAL": shutdown_signal, + "TEST_STARTUP_SHUTDOWN": str(int(startup_shutdown)), + "TEST_EXIT_BEFORE_MONITORING": str(int(exit_before_monitoring)), + }, + ) + + def test_prefill_exit_after_router_start_stops_router(self): + """Treat a prefill exit after router startup as a failure and stop the router.""" + for exit_status in (0, 7): + with self.subTest(exit_status=exit_status): + status, stdout, stderr = self.run_coupled_lifecycle(exit_role="prefill", exit_status=exit_status) + expected_status = exit_status or 1 + self.assertEqual(status, expected_status, stderr) + self.assertIn("router-started\n", stdout) + self.assertIn("prefill-exited\n", stdout) + self.assertIn("router-stopped\n", stdout) + self.assertIn(f"prefill process exited after startup (status={expected_status})", stderr) + + def test_router_exit_after_startup_stops_prefill(self): + """Propagate unexpected router exits as failures and stop prefill.""" + for exit_status in (0, 9): + with self.subTest(exit_status=exit_status): + status, stdout, stderr = self.run_coupled_lifecycle(exit_role="router", exit_status=exit_status) + expected_status = exit_status or 1 + self.assertEqual(status, expected_status, stderr) + self.assertIn("router-exited\n", stdout) + self.assertIn("prefill-stopped\n", stdout) + self.assertIn(f"router process exited after startup (status={expected_status})", stderr) + + def test_coupled_shutdown_signals_stop_both_processes(self): + """Stop both services on SIGTERM or SIGINT while preserving the cancellation exit status.""" + for signal_name, expected_status in (("TERM", 143), ("INT", 130)): + with self.subTest(signal=signal_name): + status, stdout, stderr = self.run_coupled_lifecycle(shutdown_signal=signal_name) + self.assertEqual(status, expected_status, stderr) + self.assertIn("prefill-stopped\n", stdout) + self.assertIn("router-stopped\n", stdout) + self.assertNotIn("ERROR:", stderr) + + def test_coupled_shutdown_during_readiness_stops_prefill(self): + """Stop prefill on cancellation during readiness without starting the router.""" + for signal_name, expected_status in (("TERM", 143), ("INT", 130)): + with self.subTest(signal=signal_name): + status, stdout, stderr = self.run_coupled_lifecycle(shutdown_signal=signal_name, startup_shutdown=True) + self.assertEqual(status, expected_status, stderr) + self.assertIn("prefill-stopped\n", stdout) + self.assertNotIn("router-started\n", stdout) + self.assertNotIn("ERROR:", stderr) + + def test_prefill_exit_before_runtime_monitoring_is_detected(self): + """Catch a prefill exit between the final health response and runtime monitoring.""" + status, _, stderr = self.run_coupled_lifecycle(exit_status=7, exit_before_monitoring=True) + self.assertEqual(status, 7, stderr) + self.assertIn("prefill process exited after startup (status=7)", stderr) From 195962d187da0c6e6f7f37d5fb8308972792f483 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Thu, 10 Sep 2026 23:26:48 +0000 Subject: [PATCH 12/15] Expand Super vLLM launcher regression coverage Test independent and coupled node roles, existing model settings, submission overrides, invalid inputs, and serving-only execution. Signed-off-by: Frankie Siino --- tests/unit_tests/test_super_vllm_launcher.py | 293 ++++++++++++++++++- 1 file changed, 286 insertions(+), 7 deletions(-) diff --git a/tests/unit_tests/test_super_vllm_launcher.py b/tests/unit_tests/test_super_vllm_launcher.py index 2458740baf..c8f98b6b5f 100644 --- a/tests/unit_tests/test_super_vllm_launcher.py +++ b/tests/unit_tests/test_super_vllm_launcher.py @@ -51,21 +51,30 @@ def run_shell(self, command, *args, env=None): self.fail(f"Launcher did not terminate. stdout={stdout!r}, stderr={stderr!r}") return proc.returncode, stdout, stderr - def generate_commands(self, *overrides, env=None): - # Capture both commands using a shell function; never submit a Slurm job. + def capture_submission(self, *eval_args, env=None): + # Capture generated commands and both sbatch calls without submitting jobs. stub = r""" sbatch() { if [[ -n "${vllm_command:-}" ]]; then - printf '%s\0' "$eval_command" "$vllm_command" >&2 + printf '%s\0' "$eval_command" "$vllm_command" "$batch_command" >&2 fi + printf '%s\0' "$@" >&2 + printf '\0' >&2 printf '12345\n' } -source "$@" +launcher_script=$1 +shift +source "$launcher_script" "$@" """ - status, _, captured = self.run_shell(stub, str(SCRIPT), "--config", "benchmark.yaml", *overrides, env=env) + status, _, captured = self.run_shell(stub, str(SCRIPT), *eval_args, env=env) self.assertEqual(status, 0, captured) - eval_command, pd_command, terminator = captured.split("\0") - self.assertEqual(terminator, "") + eval_command, pd_command, batch_command, submissions = captured.split("\0", 3) + self.assertTrue(submissions.endswith("\0\0")) + calls = [call.split("\0") for call in submissions.removesuffix("\0\0").split("\0\0")] + return eval_command, pd_command, batch_command, calls + + def generate_commands(self, *overrides, env=None): + eval_command, pd_command, _, _ = self.capture_submission("--config", "benchmark.yaml", *overrides, env=env) return eval_command, pd_command def eval_arguments(self, *overrides, env=None): @@ -85,6 +94,276 @@ def eval_arguments(self, *overrides, env=None): def settings(self, args, key): return [arg for arg in args if arg.lstrip("+").startswith(key + "=")] + def serving_arguments(self, command, *, rank, coupled_head=False, env=None): + # Record argv separately for vLLM and the router; marker files synchronize startup. + stubs = r""" +VLLM_COMMON_ARGS=(--common-test 'value with spaces') +VLLM_PREFILL_ARGS=(--prefill-test producer) +VLLM_DECODE_ARGS=(--decode-test consumer) +vllm() { + printf '%s\0' "$VLLM_NIXL_SIDE_CHANNEL_HOST" "$VLLM_NIXL_SIDE_CHANNEL_PORT" "$@" + touch "$TEST_STATE_DIR/service-ready" + if [[ "$TEST_COUPLED_HEAD" == 1 ]]; then + while true; do command sleep 0.01; done + fi + if (( SLURM_PROCID == 0 )); then + while [[ ! -f "$TEST_STATE_DIR/router-ready" ]]; do command sleep 0.01; done + fi +} +vllm-router() { + printf '%s\0' "$@" >&2 + touch "$TEST_STATE_DIR/router-ready" + if [[ "$TEST_COUPLED_HEAD" == 1 ]]; then kill -TERM "$$"; fi +} +curl() { [[ -f "$TEST_STATE_DIR/service-ready" ]]; } +sleep() { command sleep 0.01; } +hostname() { printf 'node%s\n' "$SLURM_PROCID"; } +""" + with TemporaryDirectory(prefix="gym-serving-args-") as state_dir: + status, stdout, stderr = self.run_shell( + stubs + command, + env=(env or {}) + | { + "SLURM_PROCID": str(rank), + "TEST_STATE_DIR": state_dir, + "TEST_COUPLED_HEAD": str(int(coupled_head)), + }, + ) + self.assertEqual(status, 143 if coupled_head else 0, stderr) + host, nixl_port, *vllm_args = stdout.rstrip("\0").split("\0") + router_args = stderr.rstrip("\0").split("\0") if stderr else [] + return host, nixl_port, vllm_args, router_args + + def assert_router_arguments(self, args, prefill_urls, decode_urls): + expected = [ + "--prefill-policy", + "cache_aware", + "--decode-policy", + "cache_aware", + "--balance-abs-threshold", + "4", + "--balance-rel-threshold", + "1.1", + "--vllm-pd-disaggregation", + "--host", + "node0", + "--port", + "8000", + "--intra-node-data-parallel-size", + "1", + "--request-timeout-secs", + "86400", + "--log-level", + "error", + ] + # URL options may precede or follow the common options; retain their tier order. + for flag, urls in (("--prefill", prefill_urls), ("--decode", decode_urls)): + actual_urls = [] + remaining = [] + i = 0 + while i < len(args): + if args[i] == flag: + actual_urls.append(args[i + 1]) + i += 2 + else: + remaining.append(args[i]) + i += 1 + self.assertEqual(actual_urls, urls) + args = remaining + self.assertEqual(args, expected) + + def test_independent_mode_preserves_per_node_engines(self): + """Keep one engine per node and all router destinations with default or explicit independent mode.""" + for mode in (None, "independent"): + for prefill_count, decode_count in ((1, 1), (1, 4), (4, 4)): + env = {"NUM_PREFILL_NODES": str(prefill_count), "NUM_DECODE_NODES": str(decode_count)} + if mode is not None: + env["VLLM_PD_DEPLOYMENT_MODE"] = mode + _, command = self.generate_commands(env=env) + for rank in range(prefill_count + decode_count): + with self.subTest(mode=mode, prefill=prefill_count, decode=decode_count, rank=rank): + host, nixl_port, args, router = self.serving_arguments(command, rank=rank, env=env) + is_prefill = rank < prefill_count + self.assertEqual(host, f"node{rank}") + self.assertEqual(nixl_port, "5600" if is_prefill else "5700") + self.assertEqual( + args, + [ + "serve", + "/test/model", + "--served-model-name", + "/test/model", + "--common-test", + "value with spaces", + "--prefill-test" if is_prefill else "--decode-test", + "producer" if is_prefill else "consumer", + "--host", + f"node{rank}", + "--port", + "8001", + ], + ) + if rank == 0: + self.assert_router_arguments( + router, + [f"http://node{i}:8001" for i in range(prefill_count)], + [f"http://node{i}:8001" for i in range(prefill_count, prefill_count + decode_count)], + ) + else: + self.assertEqual(router, []) + + def test_coupled_nodes_use_correct_tier_roles_and_ranks(self): + """Assign the correct API heads, headless ranks, tier sizes, and ports for each coupled node.""" + for prefill_count, decode_count in ((1, 1), (1, 4), (2, 3), (4, 4)): + env = { + "VLLM_PD_DEPLOYMENT_MODE": "coupled", + "NUM_PREFILL_NODES": str(prefill_count), + "NUM_DECODE_NODES": str(decode_count), + "MODEL_NAME": "served-model-alias", + } + _, command = self.generate_commands(env=env) + for rank in range(prefill_count + decode_count): + with self.subTest(prefill=prefill_count, decode=decode_count, rank=rank): + host, nixl_port, args, router = self.serving_arguments( + command, rank=rank, coupled_head=rank == 0, env=env + ) + is_prefill = rank < prefill_count + local_rank = rank if is_prefill else rank - prefill_count + head = "node0" if is_prefill else f"node{prefill_count}" + self.assertEqual(host, f"node{rank}") + self.assertEqual(nixl_port, "5600" if is_prefill else "5700") + expected = [ + "serve", + "/test/model", + "--served-model-name", + "served-model-alias", + "--common-test", + "value with spaces", + "--prefill-test" if is_prefill else "--decode-test", + "producer" if is_prefill else "consumer", + ] + if local_rank == 0: + expected += ["--host", host, "--port", "8001" if is_prefill else "8002"] + else: + expected += ["--headless"] + expected += ["--data-parallel-size", str(prefill_count if is_prefill else decode_count)] + if local_rank != 0: + expected += ["--data-parallel-start-rank", str(local_rank)] + expected += [ + "--data-parallel-address", + head, + "--data-parallel-rpc-port", + "13345" if is_prefill else "13346", + ] + if local_rank == 0: + expected += ["--api-server-count", "1"] + self.assertEqual(args, expected) + if rank == 0: + self.assert_router_arguments( + router, ["http://node0:8001"], [f"http://node{prefill_count}:8002"] + ) + else: + self.assertEqual(router, []) + + def test_existing_recipes_preserve_sampling_overrides(self): + """Pass through existing models' sampling parameters without injecting Ultra evaluation defaults.""" + for name in ("inkling_small.sh", "qwen3.5-122b-a10b.sh", "nemotron_3.5_super.sh", "nemotron_3.5_lightning.sh"): + with self.subTest(recipe=name): + args = self.eval_arguments(env={"VLLM_CONFIG": str(SCRIPT.parent / "vllm_configs" / name)}) + expected = ( + [] + if name == "inkling_small.sh" + else [ + "++policy_model.responses_api_models.vllm_model.sampling_overrides.temperature=1.0", + "++policy_model.responses_api_models.vllm_model.sampling_overrides.top_p=0.95", + ] + ) + self.assertEqual([arg for arg in args if ".sampling_overrides." in arg], expected) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), []) + self.assertEqual(self.settings(args, "resume_from_cache"), []) + + def test_submission_defaults_preserve_time_and_full_allocation_segment(self): + """Keep the four-hour default and one segment spanning every allocated node.""" + for prefill_count, decode_count in ((1, 1), (1, 4), (4, 4)): + with self.subTest(prefill=prefill_count, decode=decode_count): + _, _, _, calls = self.capture_submission( + "--config", + "benchmark.yaml", + env={"NUM_PREFILL_NODES": str(prefill_count), "NUM_DECODE_NODES": str(decode_count)}, + ) + self.assertEqual(len(calls), 2) + self.assertIn(f"--nodes={prefill_count + decode_count}", calls[0]) + self.assertIn(f"--segment={prefill_count + decode_count}", calls[0]) + self.assertIn("--time=04:00:00", calls[0]) + self.assertIn("--ntasks-per-node=1", calls[0]) + self.assertIn("--exclusive", calls[0]) + self.assertIn("--dependency=afterany:12345", calls[1]) + + def test_submission_overrides_do_not_change_cleanup_allocation(self): + """Apply custom walltime and segment size only to the main job, leaving cleanup CPU-only and short.""" + _, _, _, calls = self.capture_submission( + "--config", "benchmark.yaml", env={"SBATCH_TIME": "7-00:00:00", "VLLM_SLURM_SEGMENT": "4"} + ) + self.assertEqual(len(calls), 2) + self.assertIn("--time=7-00:00:00", calls[0]) + self.assertIn("--segment=4", calls[0]) + self.assertIn("--time=00:30:00", calls[1]) + self.assertIn("--nodes=1", calls[1]) + self.assertIn("--partition=cpu", calls[1]) + self.assertIn("--gres=none", calls[1]) + self.assertFalse(any(arg.startswith("--segment=") for arg in calls[1])) + + def test_invalid_launcher_controls_are_rejected_before_submission(self): + """Reject invalid concurrency, deployment modes, and segment sizes without submitting any job.""" + stub = 'sbatch() { printf "unexpected-submission\\n"; }; source "$@"' + cases = { + "NUM_SAMPLES_IN_PARALLEL": (("0", "-1", "1.5", "bad"), 1), + "VLLM_PD_DEPLOYMENT_MODE": (("bad", "COUPLED"), 1), + "VLLM_SLURM_SEGMENT": (("0", "-1", "1.5", "bad"), 2), + } + for name, (values, expected_status) in cases.items(): + for value in values: + with self.subTest(name=name, value=value): + status, stdout, stderr = self.run_shell(stub, str(SCRIPT), env={name: value}) + self.assertEqual(status, expected_status, stderr) + self.assertNotIn("unexpected-submission", stdout) + self.assertIn(name, stderr) + + def test_serving_only_skips_evaluation_and_cleanup_submission(self): + """Run only the serving step without eval arguments and preserve its success or failure status.""" + stubs = r""" +scontrol() { printf '%s\n' node0 node1 node2 node3 node4 node5 node6 node7; } +srun() { printf '%s\0' "$@"; printf '\0'; return "$TEST_SERVER_STATUS"; } +""" + for mode in ("independent", "coupled"): + for server_status in (0, 7): + with self.subTest(mode=mode, server_status=server_status): + _, _, batch_command, calls = self.capture_submission( + env={"VLLM_PD_DEPLOYMENT_MODE": mode, "EXPERIMENT_NAME": "", "EXPORT_TO_CSV": "1"} + ) + self.assertEqual(len(calls), 1) + self.assertIn("--job-name=gym-vllm_only-launcher-test", calls[0]) + status, stdout, stderr = self.run_shell( + stubs + batch_command, + env={ + "SLURM_JOB_NODELIST": "test-nodes", + "SLURM_SUBMIT_DIR": "/test", + "SLURM_CPUS_ON_NODE": "64", + "vllm_command": "fake-serving-command", + "eval_command": "fake-eval-command", + "TEST_SERVER_STATUS": str(server_status), + }, + ) + self.assertEqual(status, server_status, stderr) + steps = stdout.removesuffix("\0\0").split("\0\0") + self.assertEqual(len(steps), 1) + args = steps[0].split("\0") + self.assertIn("--nodes=8", args) + self.assertIn("--ntasks=8", args) + self.assertIn("--kill-on-bad-exit=1", args) + self.assertIn("fake-serving-command", args) + self.assertNotIn("--overlap", args) + def test_environment_fallbacks_are_preserved(self): """Apply concurrency and resume environment defaults, including the stable results path.""" args = self.eval_arguments(env={"NUM_SAMPLES_IN_PARALLEL": "16", "RESUME_EVAL_ON_REQUEUE": "1"}) From 21a0a1019a4a4d0481c7403fdaf20816671c07f9 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Fri, 11 Sep 2026 16:57:43 +0000 Subject: [PATCH 13/15] Prefer Hydra overrides for resuming and concurrency Signed-off-by: Frankie Siino --- .../sbatch_external_vllm.sh | 45 +------- tests/unit_tests/test_super_vllm_launcher.py | 102 +++++++----------- 2 files changed, 41 insertions(+), 106 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index de7e5a94cd..4e35e3d098 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -11,10 +11,6 @@ CONTAINER=$CONTAINER MOUNTS=$MOUNTS VLLM_CONFIG=$VLLM_CONFIG SBATCH_TIME="${SBATCH_TIME:-04:00:00}" -NUM_SAMPLES_IN_PARALLEL="${NUM_SAMPLES_IN_PARALLEL:-}" -NUM_SAMPLES_IN_PARALLEL_ARG="" -RESUME_EVAL_ON_REQUEUE="${RESUME_EVAL_ON_REQUEUE:-0}" -RESUME_FROM_CACHE_ARG="" # Independent mode starts one complete TP model replica per node. Coupled mode # forms one multi-node DP/EP engine per tier for models that cannot fit per node. VLLM_PD_DEPLOYMENT_MODE="${VLLM_PD_DEPLOYMENT_MODE:-independent}" @@ -26,39 +22,6 @@ OPENSANDBOX_DOMAIN="${OPENSANDBOX_DOMAIN:-}" OPENSANDBOX_API_KEY="${OPENSANDBOX_API_KEY:-}" OPENSANDBOX_PROTOCOL="${OPENSANDBOX_PROTOCOL:-http}" -if [[ -n "$NUM_SAMPLES_IN_PARALLEL" ]]; then - if [[ ! "$NUM_SAMPLES_IN_PARALLEL" =~ ^[1-9][0-9]*$ ]]; then - echo "ERROR: NUM_SAMPLES_IN_PARALLEL must be a positive integer; got '$NUM_SAMPLES_IN_PARALLEL'." >&2 - exit 1 - fi - NUM_SAMPLES_IN_PARALLEL_ARG="++num_samples_in_parallel=$NUM_SAMPLES_IN_PARALLEL" -fi - -case "$RESUME_EVAL_ON_REQUEUE" in - 0) - ;; - 1) - RESUME_FROM_CACHE_ARG="++resume_from_cache=true" - ;; - *) - echo "ERROR: RESUME_EVAL_ON_REQUEUE must be 0 or 1; got '$RESUME_EVAL_ON_REQUEUE'." >&2 - exit 1 - ;; -esac - -# Environment controls are defaults; explicit Hydra overrides belong to Gym. -# Avoid adding a duplicate setting that could overwrite the caller's value. -for eval_arg in "$@"; do - case "$eval_arg" in - num_samples_in_parallel=* | +num_samples_in_parallel=* | ++num_samples_in_parallel=*) - NUM_SAMPLES_IN_PARALLEL_ARG="" - ;; - resume_from_cache=* | +resume_from_cache=* | ++resume_from_cache=*) - RESUME_FROM_CACHE_ARG="" - ;; - esac -done - case "$VLLM_PD_DEPLOYMENT_MODE" in independent | coupled) ;; @@ -112,11 +75,7 @@ source "$VLLM_CONFIG" gym eval prepare $@ +use_cached_prepared_benchmarks=true -if (( $RESUME_EVAL_ON_REQUEUE )); then - experiment_name=$EXPERIMENT_NAME/resumable -else - experiment_name=$EXPERIMENT_NAME/slurm_job_id_\$SLURM_JOB_ID/date_\$(date +%Y%m%d_%H%M%S) -fi +experiment_name=$EXPERIMENT_NAME/slurm_job_id_\$SLURM_JOB_ID/date_\$(date +%Y%m%d_%H%M%S) # export_to_csv.py derives _aggregate_metrics.json from this, so the # default timestamped name makes the aggregate unfindable to anything that # did not watch the job run. Override it when results/ is already per-run. @@ -142,7 +101,6 @@ gym eval run \ ++split=benchmark \ ++use_absolute_ip=true \ ++reuse_existing_data_preparation=true \ - $RESUME_FROM_CACHE_ARG \ ++policy_base_url=http://\$(getent hosts "\$ROUTER_NODE" | awk 'NR == 1 {print \$1}'):$ROUTER_SERVER_PORT/v1 \ ++policy_api_key=dummy_api_key \ ++policy_model_name=$MODEL_NAME \ @@ -150,7 +108,6 @@ gym eval run \ ++global_aiohttp_connector_limit_per_host=16384 \ ++port_range_low=63000 \ ++port_range_high=64000 \ - $NUM_SAMPLES_IN_PARALLEL_ARG \ "\${GYM_MODEL_PARAMS[@]}" diff --git a/tests/unit_tests/test_super_vllm_launcher.py b/tests/unit_tests/test_super_vllm_launcher.py index a3c4979f04..3619de8f06 100644 --- a/tests/unit_tests/test_super_vllm_launcher.py +++ b/tests/unit_tests/test_super_vllm_launcher.py @@ -84,7 +84,7 @@ def eval_arguments(self, *overrides, env=None): gym() { if [[ "$2" == run ]]; then printf '%s\0' "$@"; fi } -date() { printf '20260909_120000\n'; } +date() { printf '%s\n' "${TEST_DATE:-20260909_120000}"; } getent() { printf '10.0.0.1 node0\n'; } """ status, stdout, stderr = self.run_shell(stubs + command, env=env) @@ -361,10 +361,9 @@ def test_submission_overrides_do_not_change_cleanup_allocation(self): self.assertFalse(any(arg.startswith("--segment=") for arg in calls[1])) def test_invalid_launcher_controls_are_rejected_before_submission(self): - """Reject invalid concurrency, deployment modes, and segment sizes without submitting any job.""" + """Reject invalid deployment modes and segment sizes without submitting any job.""" stub = 'sbatch() { printf "unexpected-submission\\n"; }; source "$@"' cases = { - "NUM_SAMPLES_IN_PARALLEL": (("0", "-1", "1.5", "bad"), 1), "VLLM_PD_DEPLOYMENT_MODE": (("bad", "COUPLED"), 1), "VLLM_SLURM_SEGMENT": (("0", "-1", "1.5", "bad"), 2), } @@ -411,81 +410,60 @@ def test_serving_only_skips_evaluation_and_cleanup_submission(self): self.assertIn("fake-serving-command", args) self.assertNotIn("--overlap", args) - def test_environment_fallbacks_are_preserved(self): - """Apply concurrency and resume environment defaults, including the stable results path.""" - args = self.eval_arguments(env={"NUM_SAMPLES_IN_PARALLEL": "16", "RESUME_EVAL_ON_REQUEUE": "1"}) - self.assertEqual(self.settings(args, "num_samples_in_parallel"), ["++num_samples_in_parallel=16"]) - self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) - self.assertIn("++output_jsonl_fpath=results/launcher-test/resumable.jsonl", args) - - def test_no_environment_defaults_leave_gym_settings_untouched(self): - """Preserve Gym settings and timestamped output naming when environment overrides are unset.""" + def test_default_evaluation_settings_are_unchanged(self): + """Leave concurrency and resume to Gym and preserve timestamped output naming by default.""" args = self.eval_arguments() self.assertEqual(self.settings(args, "num_samples_in_parallel"), []) self.assertEqual(self.settings(args, "resume_from_cache"), []) self.assertIn("++output_jsonl_fpath=results/launcher-test/slurm_job_id_12345/date_20260909_120000.jsonl", args) - def test_explicit_concurrency_overrides_environment(self): - """Let explicit concurrency arguments override the environment without duplicate settings.""" + def test_explicit_concurrency_arguments_are_preserved(self): + """Pass explicit concurrency values through to Gym without replacing or duplicating them.""" for prefix in ("", "+", "++"): - with self.subTest(prefix=prefix): - override = prefix + "num_samples_in_parallel=512" - args = self.eval_arguments(override, env={"NUM_SAMPLES_IN_PARALLEL": "16"}) - self.assertEqual(self.settings(args, "num_samples_in_parallel"), [override]) + for value in (16, 32, 512): + with self.subTest(prefix=prefix, value=value): + override = f"{prefix}num_samples_in_parallel={value}" + args = self.eval_arguments(override) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), [override]) - def test_explicit_resume_overrides_environment(self): - """Honor explicit resume arguments while retaining the switch's stable output naming.""" + def test_explicit_resume_arguments_are_preserved(self): + """Pass explicit resume values through to Gym without changing the supplied output path.""" for prefix in ("", "+", "++"): for value in ("true", "false"): with self.subTest(prefix=prefix, value=value): override = prefix + "resume_from_cache=" + value - args = self.eval_arguments(override, env={"RESUME_EVAL_ON_REQUEUE": "1"}) + args = self.eval_arguments(override, env={"ROLLOUTS_FPATH": "results/existing/resumable.jsonl"}) self.assertEqual(self.settings(args, "resume_from_cache"), [override]) - self.assertIn("++output_jsonl_fpath=results/launcher-test/resumable.jsonl", args) + self.assertIn("++output_jsonl_fpath=results/existing/resumable.jsonl", args) - def test_resume_switch_keeps_output_path_across_job_ids(self): - """Keep resumable output and W&B names stable when the Slurm job ID changes.""" - for job_id in ("12345", "12346"): - with self.subTest(job_id=job_id): - args = self.eval_arguments(env={"RESUME_EVAL_ON_REQUEUE": "1", "SLURM_JOB_ID": job_id}) + def test_resume_without_fixed_path_uses_restart_timestamp(self): + """Show that enabling resume alone still selects a new output path when the same job restarts.""" + for timestamp in ("20260909_120000", "20260909_180000"): + with self.subTest(timestamp=timestamp): + args = self.eval_arguments("++resume_from_cache=true", env={"TEST_DATE": timestamp}) self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) - self.assertIn("++output_jsonl_fpath=results/launcher-test/resumable.jsonl", args) - self.assertIn("+wandb_name=launcher-test/resumable", args) - - def test_invalid_resume_switch_is_rejected_before_submission(self): - """Reject an invalid resume switch before any Slurm job is submitted.""" - stub = r""" -sbatch() { printf 'unexpected-submission\n'; } -source "$@" -""" - status, stdout, stderr = self.run_shell(stub, str(SCRIPT), env={"RESUME_EVAL_ON_REQUEUE": "invalid"}) - self.assertNotEqual(status, 0) - self.assertNotIn("unexpected-submission", stdout) - self.assertIn("RESUME_EVAL_ON_REQUEUE must be 0 or 1", stderr) + self.assertIn( + f"++output_jsonl_fpath=results/launcher-test/slurm_job_id_12345/date_{timestamp}.jsonl", args + ) - def test_explicit_resume_keeps_output_path_across_job_ids(self): - """Preserve a manually supplied resume path while keeping the default per-job run naming.""" + def test_explicit_resume_keeps_output_path_across_restarts(self): + """Keep the supplied resume path across requeues and new job IDs, with timestamped logs and W&B names.""" for job_id in ("12345", "12346"): - with self.subTest(job_id=job_id): - args = self.eval_arguments( - "++resume_from_cache=true", - env={"ROLLOUTS_FPATH": "results/existing/resumable.jsonl", "SLURM_JOB_ID": job_id}, - ) - self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) - self.assertIn("++output_jsonl_fpath=results/existing/resumable.jsonl", args) - self.assertIn(f"+wandb_name=launcher-test/slurm_job_id_{job_id}/date_20260909_120000", args) - - def test_unrelated_override_does_not_suppress_concurrency_default(self): - """Keep global environment defaults when overrides only target nested agent settings.""" - args = self.eval_arguments( - "++agent.num_samples_in_parallel=8", - "++agent.resume_from_cache=false", - env={"NUM_SAMPLES_IN_PARALLEL": "16", "RESUME_EVAL_ON_REQUEUE": "1"}, - ) - self.assertEqual(self.settings(args, "num_samples_in_parallel"), ["++num_samples_in_parallel=16"]) - self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) - self.assertIn("++agent.num_samples_in_parallel=8", args) - self.assertIn("++agent.resume_from_cache=false", args) + for timestamp in ("20260909_120000", "20260909_180000"): + with self.subTest(job_id=job_id, timestamp=timestamp): + args = self.eval_arguments( + "++resume_from_cache=true", + env={ + "ROLLOUTS_FPATH": "results/existing/resumable.jsonl", + "SLURM_JOB_ID": job_id, + "TEST_DATE": timestamp, + }, + ) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn("++output_jsonl_fpath=results/existing/resumable.jsonl", args) + experiment_name = f"launcher-test/slurm_job_id_{job_id}/date_{timestamp}" + self.assertIn(f"+wandb_name={experiment_name}", args) + self.assertIn(f"+nemo_gym_log_dir=results/{experiment_name}/logs", args) def test_prefill_exit_is_detected_during_both_health_checks(self): """Detect prefill exits during either startup health check, including after request timeouts.""" From b4d8cf511346662b6246cd06f79b6b79081de3e6 Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Sat, 12 Sep 2026 01:18:31 +0000 Subject: [PATCH 14/15] Pin Ultra's non-eager MTP3 tuning baseline Enable prefix caching and piecewise graphs with async prefill and synchronous decode. Replace tuning overrides and eager fallback with settings tested in run 7091385. Signed-off-by: Frankie Siino --- .../vllm_configs/nemotron_3_ultra.sh | 95 +++---------------- 1 file changed, 15 insertions(+), 80 deletions(-) diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh index 318045de35..9b60dcf830 100644 --- a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -7,65 +7,13 @@ # parallelism can shard the 512 experts over 16 GPUs. Launch this config with # VLLM_PD_DEPLOYMENT_MODE=coupled and VLLM_SLURM_SEGMENT=4. -# The defaults reproduce the full-suite MTP3 configuration validated in Run -# 060: eager prefill and piecewise decode CUDA graphs with graph-owned inputs. -ULTRA_PREFILL_GPU_MEMORY_UTILIZATION="${ULTRA_PREFILL_GPU_MEMORY_UTILIZATION:-0.90}" -ULTRA_DECODE_GPU_MEMORY_UTILIZATION="${ULTRA_DECODE_GPU_MEMORY_UTILIZATION:-0.95}" -ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS="${ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS:-16384}" -ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS="${ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS:-8192}" -ULTRA_MAX_NUM_SEQS="${ULTRA_MAX_NUM_SEQS:-64}" -ULTRA_DECODE_CUDAGRAPH_MODE="${ULTRA_DECODE_CUDAGRAPH_MODE:-PIECEWISE}" -ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS="${ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS:-1}" -ULTRA_DECODE_ENFORCE_EAGER="${ULTRA_DECODE_ENFORCE_EAGER:-0}" -ULTRA_ENABLE_MTP="${ULTRA_ENABLE_MTP:-1}" -ULTRA_NUM_SPECULATIVE_TOKENS="${ULTRA_NUM_SPECULATIVE_TOKENS:-3}" +# Both tiers use piecewise CUDA graphs with graph-owned inputs and MTP3. +# Settings are fixed in this recipe; asynchronous scheduling is prefill-only. # Standard safetensors loading avoided the InstantTensor io_uring failures seen # against the Lustre-hosted checkpoint. export SAFETENSORS_FAST_GPU=1 -case "$ULTRA_DECODE_CUDAGRAPH_MODE" in - FULL_DECODE_ONLY | PIECEWISE | NONE) ;; - *) - echo "ERROR: ULTRA_DECODE_CUDAGRAPH_MODE must be FULL_DECODE_ONLY, PIECEWISE, or NONE; got '$ULTRA_DECODE_CUDAGRAPH_MODE'." >&2 - exit 2 - ;; -esac - -case "$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS" in - 0) - ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=false - ;; - 1) - ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON=true - ;; - *) - echo "ERROR: ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS must be 0 or 1; got '$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS'." >&2 - exit 2 - ;; -esac - -case "$ULTRA_DECODE_ENFORCE_EAGER" in - 0 | 1) ;; - *) - echo "ERROR: ULTRA_DECODE_ENFORCE_EAGER must be 0 or 1; got '$ULTRA_DECODE_ENFORCE_EAGER'." >&2 - exit 2 - ;; -esac - -case "$ULTRA_ENABLE_MTP" in - 0 | 1) ;; - *) - echo "ERROR: ULTRA_ENABLE_MTP must be 0 or 1; got '$ULTRA_ENABLE_MTP'." >&2 - exit 2 - ;; -esac - -if [[ ! "$ULTRA_NUM_SPECULATIVE_TOKENS" =~ ^[1-9][0-9]*$ ]]; then - echo "ERROR: ULTRA_NUM_SPECULATIVE_TOKENS must be a positive integer; got '$ULTRA_NUM_SPECULATIVE_TOKENS'." >&2 - exit 2 -fi - VLLM_COMMON_ARGS=( --disable-uvicorn-access-log --trust-remote-code @@ -89,43 +37,30 @@ VLLM_COMMON_ARGS=( --load-format safetensors --enable-expert-parallel --distributed-timeout-seconds 3600 - # NIXL-transferred Mamba state must not coexist with locally retained - # prefix-cache blocks; doing so triggers the multiple-local-block assertion. - --no-enable-prefix-caching + --enable-prefix-caching + # Both tiers need the same speculative width for compatible cache layouts. + --speculative-config '{"method":"mtp","num_speculative_tokens":3}' ) VLLM_PREFILL_ARGS=( --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}' - --gpu-memory-utilization "$ULTRA_PREFILL_GPU_MEMORY_UTILIZATION" - --max-num-batched-tokens "$ULTRA_PREFILL_MAX_NUM_BATCHED_TOKENS" - --max-num-seqs "$ULTRA_MAX_NUM_SEQS" + --gpu-memory-utilization 0.90 + --max-num-batched-tokens 16384 + --max-num-seqs 64 --data-parallel-size-local 1 --tensor-parallel-size 4 - --no-async-scheduling - # Eager prefill avoids the compiled/CUDA-graph stalls observed during tuning. - --enforce-eager + --async-scheduling + --compilation-config '{"cudagraph_mode":"PIECEWISE","cudagraph_copy_inputs":true,"pass_config":{"fuse_allreduce_rms":false}}' ) VLLM_DECODE_ARGS=( --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}' - --compilation-config "{\"cudagraph_mode\":\"$ULTRA_DECODE_CUDAGRAPH_MODE\",\"cudagraph_copy_inputs\":$ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS_JSON,\"pass_config\":{\"fuse_allreduce_rms\":false}}" - --gpu-memory-utilization "$ULTRA_DECODE_GPU_MEMORY_UTILIZATION" - --max-num-batched-tokens "$ULTRA_DECODE_MAX_NUM_BATCHED_TOKENS" - --max-num-seqs "$ULTRA_MAX_NUM_SEQS" + --compilation-config '{"cudagraph_mode":"PIECEWISE","cudagraph_copy_inputs":true,"cudagraph_capture_sizes": + [1,2,3,4,5,8,10,12,15,16,20,24,25,28,30,32,35,40,45,50,55,60,65,70,75,80,128,256,512],"pass_config":{"fuse_allreduce_rms":false}}' + --gpu-memory-utilization 0.95 + --max-num-batched-tokens 8192 + --max-num-seqs 64 --data-parallel-size-local 1 --tensor-parallel-size 4 --no-async-scheduling ) - -if [[ "$ULTRA_DECODE_ENFORCE_EAGER" == "1" ]]; then - # This fallback disables compilation and CUDA graphs for decode. - VLLM_DECODE_ARGS+=(--enforce-eager) -fi - -if [[ "$ULTRA_ENABLE_MTP" == "1" ]]; then - # Prefill and decode must use the same speculative width so the transferred - # cache layouts agree. - ULTRA_SPECULATIVE_CONFIG="{\"method\":\"mtp\",\"num_speculative_tokens\":$ULTRA_NUM_SPECULATIVE_TOKENS}" - VLLM_PREFILL_ARGS+=(--speculative-config "$ULTRA_SPECULATIVE_CONFIG") - VLLM_DECODE_ARGS+=(--speculative-config "$ULTRA_SPECULATIVE_CONFIG") -fi From 11f2d9097a2471e307384416b4d5a20cfddd9a6b Mon Sep 17 00:00:00 2001 From: Frankie Siino Date: Sat, 19 Sep 2026 03:20:37 +0000 Subject: [PATCH 15/15] revert Signed-off-by: Frankie Siino --- tests/unit_tests/test_super_vllm_launcher.py | 104 +++---------------- 1 file changed, 13 insertions(+), 91 deletions(-) diff --git a/tests/unit_tests/test_super_vllm_launcher.py b/tests/unit_tests/test_super_vllm_launcher.py index 2df4ea6102..9983044096 100644 --- a/tests/unit_tests/test_super_vllm_launcher.py +++ b/tests/unit_tests/test_super_vllm_launcher.py @@ -94,16 +94,14 @@ def eval_arguments(self, *overrides, env=None): def settings(self, args, key): return [arg for arg in args if arg.lstrip("+").startswith(key + "=")] - def serving_arguments( - self, command: str, *, rank: int, coupled_head: bool = False, env: dict[str, str] | None = None - ) -> tuple[str, str, list[str], list[str]]: + def serving_arguments(self, command, *, rank, coupled_head=False, env=None): # Record argv separately for vLLM and the router; marker files synchronize startup. stubs = r""" VLLM_COMMON_ARGS=(--common-test 'value with spaces') VLLM_PREFILL_ARGS=(--prefill-test producer) VLLM_DECODE_ARGS=(--decode-test consumer) vllm() { - printf '%s\0' "${VLLM_NIXL_SIDE_CHANNEL_HOST:-}" "${VLLM_NIXL_SIDE_CHANNEL_PORT:-}" "$@" + printf '%s\0' "$VLLM_NIXL_SIDE_CHANNEL_HOST" "$VLLM_NIXL_SIDE_CHANNEL_PORT" "$@" touch "$TEST_STATE_DIR/service-ready" if [[ "$TEST_COUPLED_HEAD" == 1 ]]; then while true; do command sleep 0.01; done @@ -229,63 +227,6 @@ def test_aggregated_submission_preserves_node_count(self) -> None: self.assertIn(f"--segment={count or 1}", calls[0]) self.assertIn(f"--ntasks={count or 1}", batch) - def test_aggregated_workers_strip_transfers_and_route_all_nodes(self) -> None: - """Standalone workers receive intact arguments with P/D cache-transfer options removed.""" - for tuning, expected in ( - (None, ["--prefill-test", "prefill value"]), - ("()", []), - ("(--aggregated-test 'custom value' --kv-transfer-config '{}')", ["--aggregated-test", "custom value"]), - ): - env = { - "VLLM_MODE": "aggregated", - "NUM_NODES": "3", - "ALL_NODES": "node0 node1 node2", - "ROUTER_POLICY": "round_robin", - } - _, command = self.generate_commands(env=env) - config = """ -VLLM_COMMON_ARGS=(--common-test 'common value' --kv-transfer-config '{}') -VLLM_PREFILL_ARGS=(--prefill-test 'prefill value' '--kv-transfer-config={}') -""" - if tuning is not None: - config += f"VLLM_AGGREGATED_ARGS={tuning}\n" - for rank in range(3): - with self.subTest(tuning=tuning, rank=rank): - host, port, args, router = self.serving_arguments(config + command, rank=rank, env=env) - self.assertEqual((host, port), ("", "")) - self.assertEqual( - args, - ["serve", "/test/model", "--served-model-name", "/test/model", "--common-test", "common value"] - + expected - + ["--host", f"node{rank}", "--port", "8001"], - ) - if rank == 0: - self.assertEqual( - router, - [ - "--host", - "node0", - "--port", - "8000", - "--intra-node-data-parallel-size", - "1", - "--request-timeout-secs", - "86400", - "--worker-startup-timeout-secs", - "1200", - "--log-level", - "error", - "--policy", - "round_robin", - "--worker-urls", - "http://node0:8001", - "http://node1:8001", - "http://node2:8001", - ], - ) - else: - self.assertEqual(router, []) - def test_aggregated_coupled_combination_is_rejected(self) -> None: """Reject unsupported multi-node aggregated execution before submitting jobs.""" status, stdout, stderr = self.run_shell( @@ -306,22 +247,9 @@ def test_generated_scripts_have_valid_syntax(self) -> None: result = subprocess.run(["bash", "-n"], input=command, text=True, capture_output=True, timeout=5) self.assertEqual(result.returncode, 0, result.stderr) - def test_ultra_runtime_overrides_preserve_other_models_defaults(self) -> None: - """Source the recipe after launcher defaults and inspect the effective worker environment.""" - defaults = { - "VLLM_USE_V2_MODEL_RUNNER": "unset", - "VLLM_SSM_CONV_STATE_LAYOUT": "unset", - "UCX_TLS": "rc_x,rc,dc_x,dc,cuda_copy,cuda_ipc", - "UCX_NET_DEVICES": "all", - "UCX_IB_ADDR_TYPE": "unset", - "UCX_RNDV_SCHEME": "get_zcopy", - "UCX_RNDV_THRESH": "0", - "NCCL_CUMEM_ENABLE": "1", - "NCCL_MNNVL_ENABLE": "1", - "NCCL_NVLS_ENABLE": "1", - "VLLM_HTTP_TIMEOUT_KEEP_ALIVE": "180", - } - ultra = defaults | { + def test_ultra_runtime_settings_override_launcher_defaults(self) -> None: + """Ultra configures its model runner, state layout, and communication environment.""" + expected = { "VLLM_USE_V2_MODEL_RUNNER": "0", "VLLM_SSM_CONV_STATE_LAYOUT": "DS", "UCX_TLS": "rc_x,rc,cuda_copy,cuda_ipc", @@ -331,18 +259,13 @@ def test_ultra_runtime_overrides_preserve_other_models_defaults(self) -> None: "NCCL_MNNVL_ENABLE": "unset", "NCCL_NVLS_ENABLE": "unset", } - for recipe, expected in ( - ("/dev/null", defaults), - (str(SCRIPT.parent / "vllm_configs/nemotron_3_ultra.sh"), ultra), - (str(SCRIPT.parent / "vllm_configs/inkling_small.sh"), defaults | {"VLLM_USE_V2_MODEL_RUNNER": "1"}), - ): - with self.subTest(recipe=recipe): - _, command = self.generate_commands(env={"VLLM_CONFIG": recipe}) - setup = command.split("this_node_hostname=", 1)[0] - inspect = '\nfor name in "$@"; do printf "%s\\0" "${!name-unset}"; done\n' - status, stdout, stderr = self.run_shell(setup + inspect, *expected) - self.assertEqual(status, 0, stderr) - self.assertEqual(dict(zip(expected, stdout.removesuffix("\0").split("\0"), strict=True)), expected) + recipe = str(SCRIPT.parent / "vllm_configs/nemotron_3_ultra.sh") + _, command = self.generate_commands(env={"VLLM_CONFIG": recipe}) + setup = command.split("this_node_hostname=", 1)[0] + inspect = '\nfor name in "$@"; do printf "%s\\0" "${!name-unset}"; done\n' + status, stdout, stderr = self.run_shell(setup + inspect, *expected) + self.assertEqual(status, 0, stderr) + self.assertEqual(dict(zip(expected, stdout.removesuffix("\0").split("\0"), strict=True)), expected) def test_coupled_nodes_use_correct_tier_roles_and_ranks(self): """Assign coupled tier roles, ranks, and ports while leaving router balancing thresholds at defaults.""" @@ -535,11 +458,10 @@ def test_walltime_alias_precedence_and_cleanup_isolation(self) -> None: self.assertIn("--time=00:30:00", calls[1]) self.assertFalse(any(arg.startswith("--segment=") for arg in calls[1])) - def test_invalid_launcher_controls_are_rejected_before_submission(self) -> None: + def test_invalid_launcher_controls_are_rejected_before_submission(self): """Reject invalid deployment modes and segment sizes without submitting any job.""" stub = 'sbatch() { printf "unexpected-submission\\n"; }; source "$@"' cases = { - "VLLM_MODE": (("bad", "PD"), 1), "VLLM_PD_DEPLOYMENT_MODE": (("bad", "COUPLED"), 1), "VLLM_SLURM_SEGMENT": (("0", "-1", "1.5", "bad"), 2), "SEGMENT": (("0", "-1", "1.5", "bad"), 2),