feat: fix the vLLM DP path

examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-vllm-dp8.yaml

@@ -0,0 +1,9 @@
+defaults: ./grpo-moonlight-16ba3b-4n8g-megatron.yaml
+policy:
+  generation:
+    vllm_cfg:
+      expert_parallel_size: 8
+      async_engine: true
+logger:
+  wandb:
+    name: grpo-moonlight-16ba3b-4n8g-megatron-vllm-dp

examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-vllm-tp2ep16.yaml

@@ -0,0 +1,12 @@
+defaults: ./grpo-moonlight-16ba3b-4n8g-megatron.yaml
+checkpointing:
+  checkpoint_dir: results/grpo-moonlight-16ba3b-4n8g-megatron-vllm-tp2ep16
+policy:
+  generation:
+    vllm_cfg:
+      tensor_parallel_size: 2  # test with model parallelism
+      expert_parallel_size: 16  # test with cross-node data parallelism
+logger:
+  wandb:
+    name: grpo-moonlight-16ba3b-4n8g-megatron-vllm-tp2ep16
+

nemo_rl/models/generation/vllm/vllm_generation.py

@@ -65,7 +65,6 @@ def __init__(
             f"Got world size {cluster.world_size()} and model parallel size (TP * PP) {self.model_parallel_size}."
         )
         self.dp_size = cluster.world_size() // self.model_parallel_size
-        self.vllm_dp_size = self.ep_size // self.tp_size
 
         if self.pp_size > 1:
             assert self.cfg["vllm_cfg"]["async_engine"], (
@@ -78,13 +77,7 @@ def __init__(
                 "When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP. "
                 "Please update your configuration to set expert_parallel_size to a multiple of tensor_parallel_size."
             )
-            if self.ep_size != self.tp_size:
-                # vLLM's EP = DP * TP, so here we need to use DP inside vLLM.
-                assert not self.cfg["vllm_cfg"]["async_engine"], (
-                    "vLLM async_engine has some issues when using DP inside vLLM. "
-                    "Please update your configuration to set `policy.generation.vllm_cfg.async_engine=false`. "
-                    "See https://github.com/NVIDIA-NeMo/RL/issues/1101 for more details."
-                )
+        self.vllm_dp_size = self.ep_size // self.tp_size if self.ep_size > 1 else 1
 
         # Validate sampling parameters early to avoid resource allocation with unsupported configs.
         top_k: int | None = self.cfg["top_k"]
@@ -163,10 +156,6 @@ def __init__(
                 "[INFO] NCCL_NVLS_ENABLE is set to 0 for non-colocated inference with cross-node model parallelism."
                 "See https://github.com/NVIDIA-NeMo/RL/issues/1352 for more details."
             )
-        # We should use vLLM DP if ep_size > tp_size since EP_SIZE = DP_SIZE * TP_SIZE in vLLM.
-        # See details in https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/data_parallel.py
-        if self.ep_size > self.tp_size:
-            env_vars["VLLM_DP_SIZE"] = str(self.vllm_dp_size)
 
         # Check if we need parallelism-aware worker group creation
         if self.model_parallel_size > 1:
@@ -439,14 +428,15 @@ def init_collective(
             else "init_collective"
         )
 
-        # Prepare rank
+        # Prepare rank. When using vLLM DP, workers in same vLLM instance but different vLLM DP
+        # ranks have unique ranks, so the rank prefix for same vLLM instance are made the same.
         total_workers = len(self.worker_group.workers)
-        if self.dp_size == 0:
-            raise RuntimeError(
-                "Data parallel size is zero, cannot initialize collective."
-            )
-        workers_per_group = total_workers // self.dp_size
-        rank_prefix_list = list(range(0, total_workers, workers_per_group))
+        workers_per_dp_rank = total_workers // self.dp_size
+        workers_per_vllm_instance = workers_per_dp_rank * self.vllm_dp_size
+        rank_prefix_list = []
+        for dp_rank in range(self.dp_size):
+            vllm_instance_idx = dp_rank // self.vllm_dp_size
+            rank_prefix_list.append(vllm_instance_idx * workers_per_vllm_instance)
 
         # Send world_size and rank for init collective to all workers
         futures = self.worker_group.run_all_workers_multiple_data(

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -454,6 +454,10 @@ def _patch_vllm_hermes_tool_parser_thread_safety():
 
             # Use Ray for distributed execution in parallel mode
             vllm_kwargs["distributed_executor_backend"] = "ray"
+        elif self.expert_parallel_size > 1:
+            # when there is data parallelism but no model-parallelism, we need to use
+            # the mp backend, otherwise it will default to ray and cause the worker to hang.
+            vllm_kwargs["distributed_executor_backend"] = "mp"
         else:
             # For non-parallel mode, explicitly set executor to None to avoid Ray issues
             vllm_kwargs["distributed_executor_backend"] = None
@@ -462,19 +466,34 @@ def _patch_vllm_hermes_tool_parser_thread_safety():
         os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
 
         # We should use vLLM DP if ep_size > tp_size since EP_SIZE = DP_SIZE * TP_SIZE in vLLM.
-        # See details in https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/data_parallel.py
         if self.expert_parallel_size > self.tensor_parallel_size:
-            # set vLLM DP rank
-            world_size = int(os.environ["VLLM_DP_SIZE"]) * model_parallel_size
-            rank = int(os.environ["RANK"]) % world_size
-            os.environ["VLLM_DP_RANK"] = str(rank // model_parallel_size)
-            os.environ["VLLM_DP_RANK_LOCAL"] = str((rank % 8) // model_parallel_size)
-            # set vLLM DP address and port
-            leader_rank = int(os.environ["RANK"]) // world_size * world_size
+            vllm_dp_size = self.expert_parallel_size // self.tensor_parallel_size
+            world_size_across_dp = vllm_dp_size * model_parallel_size
+            rank_in_world_across_dp = int(os.environ["RANK"]) % world_size_across_dp
+            vllm_dp_rank = rank_in_world_across_dp // model_parallel_size
+            leader_rank = (
+                int(os.environ["RANK"]) // world_size_across_dp * world_size_across_dp
+            )
             addr_list = eval(os.environ["AVAILABLE_ADDR_LIST"])
+            dp_addr = addr_list[leader_rank]
             port_list = eval(os.environ["AVAILABLE_PORT_LIST"])
-            os.environ["VLLM_DP_MASTER_IP"] = addr_list[leader_rank]
-            os.environ["VLLM_DP_MASTER_PORT"] = str(port_list[leader_rank])
+            dp_port = port_list[leader_rank]
+            use_sync_engine = not self.cfg["vllm_cfg"]["async_engine"]
+            if use_sync_engine:
+                # set vLLM DP flags according to https://github.com/vllm-project/vllm/blob/main/examples/features/data_parallel/data_parallel_offline.py
+                os.environ["VLLM_DP_SIZE"] = str(vllm_dp_size)
+                os.environ["VLLM_DP_RANK"] = str(vllm_dp_rank)
+                # Always set local rank to 0 because we only expose GPUs belong to this DP rank to the worker; if we set it to the actual local rank, it will cause the worker to hang.
+                os.environ["VLLM_DP_RANK_LOCAL"] = str(0)
+                os.environ["VLLM_DP_MASTER_IP"] = str(dp_addr)
+                os.environ["VLLM_DP_MASTER_PORT"] = str(dp_port)
+            else:
+                # set vLLM DP arguments according to https://docs.vllm.ai/en/latest/serving/data_parallel_deployment/#external-load-balancing
+                vllm_kwargs["data_parallel_size"] = vllm_dp_size
+                vllm_kwargs["data_parallel_rank"] = vllm_dp_rank
+                vllm_kwargs["data_parallel_external_lb"] = True
+                vllm_kwargs["data_parallel_address"] = str(dp_addr)
+                vllm_kwargs["data_parallel_rpc_port"] = dp_port
 
         load_format = self.cfg["vllm_cfg"]["load_format"]
         if ModelFlag.VLLM_LOAD_FORMAT_AUTO.matches(self.model_name):

tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-vllm-dp8.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=4
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=30
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+PYTHONPATH=$HF_HOME/modules:$PYTHONPATH uv run examples/run_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'median(data["train/token_mult_prob_error"]) < 1.1' \
+        'data["train/token_mult_prob_error"]["10"] < 1.1' \
+        'mean(data["train/reward"]) > 0.45' \
+        'mean(data["timing/train/total_step_time"], -11, -1) < 70'
+
+    # Clean up checkpoint directory after successful run to save space.
+    rm -rf "$CKPT_DIR"
+fi

tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-vllm-tp2ep16.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=4
+GPUS_PER_NODE=4
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=30
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+PYTHONPATH=$HF_HOME/modules:$PYTHONPATH uv run examples/run_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'median(data["train/token_mult_prob_error"]) < 1.1' \
+        'data["train/token_mult_prob_error"]["10"] < 1.1' \
+        'mean(data["train/reward"]) > 0.45' \
+        'mean(data["timing/train/total_step_time"], -11, -1) < 100'
+
+    # Clean up checkpoint directory after successful run to save space.
+    rm -rf "$CKPT_DIR"
+fi
+

tests/test_suites/nightly.txt

@@ -22,6 +22,10 @@ tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron.sh
 tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3.sh
 tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh
 
+# vLLM
+tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-vllm-dp8.sh
+tests/test_suites/llm/grpo-moonlight-16ba3b-4n8g-megatron-vllm-tp2ep16.sh
+
 # Functional 32b run
 tests/test_suites/llm/grpo-qwen2.5-32b-32n8g-fsdp2tp8-actckpt.v3.sh
 

tests/unit/test_recipes_and_test_suites.py

@@ -233,7 +233,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites(
     )
 
 
-def test_nightly_compute_stays_below_1360_hours(nightly_test_suite, tracker):
+def test_nightly_compute_stays_below_1380_hours(nightly_test_suite, tracker):
     command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}"
 
     print(f"Running command: {command}")
@@ -265,8 +265,8 @@ def test_nightly_compute_stays_below_1360_hours(nightly_test_suite, tracker):
         f"Last line of output was not as expected: '{last_line}'"
     )
     total_gpu_hours = float(last_line.split(":")[-1].strip())
-    assert total_gpu_hours <= 1360, (
-        f"Total GPU hours exceeded 1360: {last_line}. We should revisit the test suites to reduce the total GPU hours."
+    assert total_gpu_hours <= 1380, (
+        f"Total GPU hours exceeded 1380: {last_line}. We should revisit the test suites to reduce the total GPU hours."
     )
     tracker.track("total_nightly_gpu_hours", total_gpu_hours)