[DP] Internal Load Balancing Per Node [one-pod-per-node]

[robertgshaw2-redhat]

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

• add capability of running with API server per node. This ensures we can use local communication between the EngineCore and AsyncLLM (i.e. do not have to go over the network), while allowing us to
• this setup allows us to load balance across nodes externally --- enabling to work in a one-pod-per-node configuration for llm-d to avoid the UCX issues we have with one-pod-per-rank balancing

NOTE:

• prometheus metrics are broken for pd. this Pr is compatible with the fix ([DP] Fix Prometheus Logging #21257)
• in this PR, we use the DP coordinator of the intra-node LB. We actually don't need to do this (we could use something local to LB). this change would require more complex surgery to vllm.

Resolves #21261

FOLLOW UPS:

• add ability to run with N Servers per node as well
• consider updating the --data-parallel-rank UX for the old external LB to a unified setup (cc @njhill)
• consider updating architecture such that the DPCoordinator to only send LB messages

Test Plan

MODEL:= "Qwen/Qwen3-30B-A3B-FP8"

dp_a:
  VLLM_LOGGING_LEVEL=DEBUG chg run --gpus 2 --  vllm serve {{MODEL}} \
    --port 8100 \
    --data-parallel-hybrid-lb \
    --data-parallel-size 4 \
    --data-parallel-size-local 2 \
    --data-parallel-start-rank 0 \
    --data-parallel-rpc-port 1234 \
    --enable-expert-parallel \
    --enforce-eager \
    --disable-log-requests

dp_b:
  VLLM_LOGGING_LEVEL=DEBUG chg run --gpus 2 -- vllm serve {{MODEL}} \
    --port 8200 \
    --data-parallel-hybrid-lb \
    --data-parallel-size 4 \
    --data-parallel-size-local 2 \
    --data-parallel-start-rank 2 \
    --data-parallel-rpc-port 1234 \
    --enable-expert-parallel \
    --enforce-eager \
    --disable-log-requests

eval PORT CONCURRENT LIMIT:
  lm_eval --model local-completions --tasks gsm8k \
    --model_args model={{MODEL}},base_url=http://127.0.0.1:{{PORT}}/v1/completions,num_concurrent={{CONCURRENT}},num_retries=0,tokenized_requests=False \
    --limit {{LIMIT}}

• launch

just dp_a

just dp_b

-run concurrently

just eval 8100 100 1000

just eval 8200 100 1000

Test Result

local-completions (model=Qwen/Qwen3-30B-A3B-FP8,base_url=http://127.0.0.1:8100/v1/completions,num_concurrent=10,num_retries=0,tokenized_requests=False), gen_kwargs: (None), limit: 100.0, num_fewshot: None, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.90|±  |0.0302|
|     |       |strict-match    |     5|exact_match|↑  | 0.96|±  |0.0197|

NOTE:

Also confirmed old ways works:

dp_a_internal_lb PORT:
  chg run --gpus 2 -- vllm serve {{MODEL}} \
    --port {{PORT}} \
    --data-parallel-size 4 \
    --data-parallel-size-local 2 \
    --data-parallel-rpc-port 1235 \
    --enable-expert-parallel \
    --enforce-eager \
    --disable-log-requests

dp_b_internal_lb:
  chg run --gpus 2 -- vllm serve {{MODEL}} \
    --headless \
    --data-parallel-size 4 \
    --data-parallel-size-local 2 \
    --data-parallel-start-rank 2 \
    --data-parallel-rpc-port 1235 \
    --enable-expert-parallel \
    --enforce-eager \
    --disable-log-requests

dp_a_external_lb PORT:
   chg run --gpus 1 -- vllm serve {{MODEL}} \
    --port {{PORT}} \
    --data-parallel-size 2 \
    --data-parallel-rank 0 \
    --data-parallel-rpc-port 1236 \
    --enable-expert-parallel \
    --enforce-eager \
    --disable-log-requests

dp_b_external_lb PORT:
  chg run --gpus 1 -- vllm serve {{MODEL}} \
    --port {{PORT}} \
    --data-parallel-size 2 \
    --data-parallel-rank 1 \
    --data-parallel-rpc-port 1236 \
    --enable-expert-parallel \
    --enforce-eager \
    --disable-log-requests

(Optional) Documentation Update

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request introduces changes to support internal load balancing for data-parallel setups, specifically in a "one-pod-per-node" configuration. The changes involve modifications to the engine argument parsing, distributed setup, and communication logic. The most critical issues are the presence of a "HACK" that hardcodes a key configuration variable and several instances of commented-out code and debugging artifacts (e.g., print and logger.info statements). These should be removed or replaced with proper, configurable implementations to ensure the code is clean, maintainable, and production-ready. Additionally, a todo comment indicates that some parts of the code may be incomplete or require further updates. Please address these points to improve the quality and clarity of the codebase.

[gemini-code-assist]

This validation logic has been commented out. If this check is no longer required, the commented code should be removed. If the check is still necessary, it should be re-enabled.

Suggested change

	# if args.data_parallel_start_rank:
	# raise ValueError(
	# "data_parallel_start_rank is only applicable "
	# "in headless mode. "
	# "Add --headless flag to enable headless mode.")
	if args.data_parallel_start_rank:
	raise ValueError(
	"data_parallel_start_rank is only applicable "
	"in headless mode. "
	"Add --headless flag to enable headless mode.")

[njhill]

@robertgshaw2-redhat I can spend some time on this tomorrow (Monday) if not before. This would be a third DP mode which is kind of a hybrid of the two existing ones. It will need some change to the coordinator and/or client load-balancing logic to constrain the set of engines considered to those associated with each API server.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @robertgshaw2-redhat.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[njhill]

The test failure is just due to too strict tolerance for the balancing. We can wait for the remaining tests to finish and I can then push a change to relax the tolerance.

[tlrmchlsmth]

This condition (self.data_parallel_size_local != self.data_parallel_size) makes it so that you can't set --data-parallel-hybrid-lb on a single node, which is really annoying since then you have to have different command line args for the single and multinode cases