Use ROLLOUT_VISIBLE_DEVICE env var for vLLM GPU placement; rename vllm_dtype to engine_dtype

Author: delock

training/opsd/README.md

@@ -86,9 +86,11 @@ NUM_TRAIN_GPUS=6 INCLUDE_GPUS=0,1,2,3,4,5 \
     bash scripts/train_opsd_vllm.sh configs/opsd_vllm_disjoint.json
 ```
 
-vLLM gets dedicated GPUs (`rollout.gpus` in the config). Training rank 0
-constructs the `LLM` handle; other training ranks receive generated token
-ids via NCCL broadcast.
+vLLM gets dedicated GPUs via the `ROLLOUT_VISIBLE_DEVICE` environment
+variable (comma-separated CUDA device indices, e.g.
+`ROLLOUT_VISIBLE_DEVICE=6,7`). Training rank 0 spawns the vLLM server as
+a subprocess with `CUDA_VISIBLE_DEVICES` set to those devices; other
+training ranks receive generated token ids via NCCL broadcast.
 
 ### Smoke tests (5 steps, small models)
 
@@ -100,7 +102,8 @@ end-to-end before scaling up.
 cd examples/opsd
 deepspeed --num_gpus 2 main.py --config configs/smoke_hybrid.json
 # For vLLM (uses GPUs 0,1 for training and 2,3 for vLLM):
-NUM_TRAIN_GPUS=2 INCLUDE_GPUS=0,1 deepspeed --num_gpus 2 --include localhost:0,1 \
+NUM_TRAIN_GPUS=2 INCLUDE_GPUS=0,1 ROLLOUT_VISIBLE_DEVICE=2,3 \
+    deepspeed --num_gpus 2 --include localhost:0,1 \
     main.py --config configs/smoke_vllm.json
 ```
 
@@ -133,6 +136,13 @@ python -m pytest tests/ -v
 See `configs/opsd_hybrid_engine.json` and `configs/opsd_vllm_disjoint.json`
 for fully-populated examples.
 
+**GPU placement for vLLM rollout:** The GPUs available to the vLLM server
+are controlled by the `ROLLOUT_VISIBLE_DEVICE` environment variable
+(comma-separated CUDA device indices, e.g. `ROLLOUT_VISIBLE_DEVICE=6,7`),
+not by a field in the JSON config. This keeps the vLLM device assignment
+decoupled from the DeepSpeed launcher's own `CUDA_VISIBLE_DEVICES` /
+`--include` flags, which control only the training ranks.
+
 ## Adding a new model architecture
 
 No special steps are needed for new model architectures. vLLM's RLHF weight

training/opsd/configs/opsd_vllm_disjoint.json

@@ -18,10 +18,9 @@
         "top_p": 1.0,
         "top_k": -1,
         "n_samples_per_prompt": 1,
-        "gpus": [6, 7],
         "tensor_parallel_size": 2,
         "gpu_memory_utilization": 0.85,
-        "vllm_dtype": "bfloat16",
+        "engine_dtype": "bfloat16",
         "weight_sync_interval": 4,
         "vllm_min_version": "0.6.4",
         "vllm_port": 8000

training/opsd/configs/smoke_vllm.json

@@ -18,10 +18,9 @@
         "top_p": 1.0,
         "top_k": -1,
         "n_samples_per_prompt": 1,
-        "gpus": [],
         "tensor_parallel_size": 1,
         "gpu_memory_utilization": 0.3,
-        "vllm_dtype": "bfloat16",
+        "engine_dtype": "bfloat16",
         "weight_sync_interval": 2,
         "vllm_min_version": "0.6.4",
         "vllm_enforce_eager": true,

training/opsd/scripts/train_opsd_vllm.sh

@@ -8,17 +8,20 @@
 #
 # The vLLM server is started **lazily** as a subprocess by training rank 0
 # on first use, so no separate vLLM launch step is required.  The GPUs
-# listed in ``rollout.gpus`` in the config are assigned to the vLLM server
-# via ``CUDA_VISIBLE_DEVICES`` in the subprocess environment.
+# assigned to the vLLM server are controlled by the ROLLOUT_VISIBLE_DEVICE
+# environment variable (comma-separated CUDA device indices).  The training
+# ranks must run on a *different* set of GPUs so the two don't contend for
+# memory.
 #
-# Default config assumes 8 GPUs: ranks 0..5 train (ZeRO-3), devices 6-7
-# run vLLM with TP=2.  Adjust configs/opsd_vllm_disjoint.json::rollout.gpus
-# and NUM_TRAIN_GPUS to match your topology.
+# Default topology: ranks 0..5 train on GPUs 0-5 (ZeRO-3), devices 6-7
+# run vLLM with TP=2.  Override via:
+#   ROLLOUT_VISIBLE_DEVICE=... NUM_TRAIN_GPUS=.. INCLUDE_GPUS=.. bash ...
 set -euo pipefail
 
 CONFIG="${1:-configs/opsd_vllm_disjoint.json}"
 NUM_TRAIN_GPUS="${NUM_TRAIN_GPUS:-6}"
 INCLUDE_GPUS="${INCLUDE_GPUS:-0,1,2,3,4,5}"
+export ROLLOUT_VISIBLE_DEVICE="${ROLLOUT_VISIBLE_DEVICE:-6,7}"
 
 deepspeed --num_gpus "${NUM_TRAIN_GPUS}" --include "localhost:${INCLUDE_GPUS}" \
     main.py --config "${CONFIG}"