Skip to content

[Feature] support rollout DP+EP for vLLM - #125

Closed
princepride wants to merge 2 commits into
mainfrom
feature/rollout-dp-ep
Closed

[Feature] support rollout DP+EP for vLLM#125
princepride wants to merge 2 commits into
mainfrom
feature/rollout-dp-ep

Conversation

@princepride

Copy link
Copy Markdown
Collaborator

Wire vLLM rollout data parallelism into TP sizing and use conservative DP sleep-mode defaults so Qwen3 DP+EP DeepEP validation can run reliably.

Wire vLLM rollout data parallelism into TP sizing and use conservative DP sleep-mode defaults so Qwen3 DP+EP DeepEP validation can run reliably.
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for vLLM data parallelism (DP) alongside expert parallelism (EP) and the DeepEP communication kernel. It adds compatibility shims for legacy slime imports, updates parallel size resolution to account for DP, extends engine timeouts for DP configurations, and applies conservative defaults for DP engines in sleep mode. Additionally, comprehensive integration and unit tests are added. Feedback on the changes suggests enforcing positive integer values for pipeline and data parallel sizes to prevent potential division-by-zero errors, and using safe dictionary access for dp_size in server_args to avoid KeyError exceptions.

Comment on lines 347 to +349
pp_size = temp_args.vllm_pipeline_parallel_size
vllm_tp_size = temp_args.rollout_num_gpus_per_engine // pp_size
dp_size = temp_args.vllm_data_parallel_size
vllm_tp_size = temp_args.rollout_num_gpus_per_engine // (pp_size * dp_size)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If vllm-pipeline-parallel-size or vllm-data-parallel-size are passed as 0 or negative values, this division will raise a ZeroDivisionError or result in an invalid negative parallel size. We should enforce a minimum value of 1 for both sizes to ensure robustness.

Suggested change
pp_size = temp_args.vllm_pipeline_parallel_size
vllm_tp_size = temp_args.rollout_num_gpus_per_engine // pp_size
dp_size = temp_args.vllm_data_parallel_size
vllm_tp_size = temp_args.rollout_num_gpus_per_engine // (pp_size * dp_size)
pp_size = max(1, temp_args.vllm_pipeline_parallel_size)
dp_size = max(1, temp_args.vllm_data_parallel_size)
vllm_tp_size = temp_args.rollout_num_gpus_per_engine // (pp_size * dp_size)

Comment on lines +132 to +133
denom = pp * dp
if gpus_per_engine % denom != 0:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If denom is 0 or negative due to invalid or zero parallel sizes, this modulo operation will raise a ZeroDivisionError or result in an invalid parallel configuration. Adding a defensive check to ensure denom is positive prevents unexpected crashes.

Suggested change
denom = pp * dp
if gpus_per_engine % denom != 0:
denom = pp * dp
if denom <= 0:
raise ValueError(
f"Invalid parallel configuration: vllm_pipeline_parallel_size ({pp}) "
f"and vllm_data_parallel_size ({dp}) must be positive integers."
)
if gpus_per_engine % denom != 0:

Comment on lines +480 to +487
if (
server_args["dp_size"] > 1
and getattr(args, "vllm_enable_sleep_mode", False)
and not _user_overrode(args, "vllm_async_scheduling")
):
cmd += ["--no-async-scheduling"]

if server_args["dp_size"] > 1 and getattr(args, "vllm_enable_sleep_mode", False):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Accessing server_args["dp_size"] directly can raise a KeyError if the dictionary is partially populated (e.g., in certain unit tests or custom orchestration flows). Using .get("dp_size", 1) is safer and consistent with how dp_size is accessed in build_vllm_subprocess_env.

Suggested change
if (
server_args["dp_size"] > 1
and getattr(args, "vllm_enable_sleep_mode", False)
and not _user_overrode(args, "vllm_async_scheduling")
):
cmd += ["--no-async-scheduling"]
if server_args["dp_size"] > 1 and getattr(args, "vllm_enable_sleep_mode", False):
dp_size = server_args.get("dp_size", 1)
if (
dp_size > 1
and getattr(args, "vllm_enable_sleep_mode", False)
and not _user_overrode(args, "vllm_async_scheduling")
):
cmd += ["--no-async-scheduling"]
if dp_size > 1 and getattr(args, "vllm_enable_sleep_mode", False):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why disable async-scheduling and enforce eager here?

aoshen02 added a commit that referenced this pull request Jun 3, 2026
…M #24882 workaround)

The DP+EP (TP1/DP4/EP) colocate rollout path crashed with a CUDA illegal
memory access in update_weights, only with data parallelism (TP survived).
Root-caused via CUDA_LAUNCH_BLOCKING to the UNMERGED vLLM PR #24882: in DP
mode the coordinator issues execute_dummy_batch (_dummy_run forward) to an
idle replica to stay collective-synced; during vime's colocate weight sync
the engine is slept (weights freed) then woken weights-only, so that dummy
forward reads freed/half-written weight memory -> IMA.

_VLLMHijack now (always-on):
  - skips Worker.execute_dummy_batch while sleeping / _weight_update_active
    (sleep state tracked via Worker.sleep/wake_up; full-wake = kv_cache in
    tags). All vime engine replicas sleep in lockstep, so skipping does not
    desync the DP collective.
  - retries Worker.determine_available_memory on the colocate startup
    memory-profiling assertion (megatron frees GPU mem mid profile_run ->
    free goes up -> assert trips); settles and retries.

vllm_engine.py: conservative DP+EP sleep-mode defaults (port of PR #125) —
dp_size>1 + sleep => --enforce-eager, --no-async-scheduling,
--distributed-timeout-seconds 1800, VLLM_ENGINE_READY_TIMEOUT_S=1800; and
longer sleep/wake HTTP timeouts.

update_weight: packed single-buffer CUDA-IPC producer (one reduce_tensor
handle per chunk, held alive until the consumer copy completes).

model_provider: wire gradient_accumulation_fusion from args.

Validated on GB200 (2x4 arm64 cu13): glm4.7-30B DP+EP colocate runs the full
RL loop end-to-end (job succeeded, 2 iterations); qwen3.6-35B DP+EP clears
weight-sync + valid rollout (reward 0.31) + training step.

NOTE: also contains env-gated (default-off) diagnostic probes
(VIME_WT_TRACE, VIME_DISABLE_LAYERWISE_RELOAD) used during root-causing;
slated for removal after review.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request Jun 3, 2026
Port the 3 unit tests from vime PR #125 that exercise engine logic
already present in this branch (vllm_engine.py:448-450, 521-536) but
previously untested on the slime/ paths:

- build_vllm_subprocess_env extends VLLM_ENGINE_READY_TIMEOUT_S to 1800
  when dp_size>1.
- build_vllm_cmd_and_env forces --no-async-scheduling, --enforce-eager,
  and --distributed-timeout-seconds 1800 for colocate DP+EP sleep mode.
- resume_memory_occupation /wake_up uses the weight-transfer HTTP
  timeout (vllm_weight_transfer_timeout_sec) rather than requests' 30s
  default.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@aoshen02 aoshen02 closed this Jun 3, 2026
aoshen02 added a commit that referenced this pull request Jun 3, 2026
…outer

Port of PR #108 onto vime/main (post-#127 topology + #105 slime/->vime/
rename). Carries the full feature set #108 accumulated, including the
DP+EP work formerly tracked by the now-closed #125.

Engine (vllm_engine.py):
- DP+EP sizing/launch: tp = gpus_per_engine // (pp * dp); EP via
  --enable-expert-parallel; single-node DP uses --data-parallel-backend mp.
- Conservative DP+EP sleep-mode defaults (no-async-scheduling,
  enforce-eager, distributed-timeout-seconds=1800; VLLM_ENGINE_READY_TIMEOUT_S
  extended for DP) to avoid cuMemcpyDtoDAsync segfaults, honouring user
  overrides.
- PD NIXL wiring: --kv-transfer-config NixlConnector (kv_both),
  VLLM_NIXL_SIDE_CHANNEL_HOST/PORT, disaggregation_bootstrap_port in
  _compute_server_args + [vllm-topo] server_args trace.

Router/rollout (rollout.py):
- _launch_static_pd_router: SkyRL-style static prefill/decode URLs for PD.
- _start_router is now non-PD only; dropped the has_pd_disaggregation param.
  (Fixes a latent NameError #108 left: the dropped param was still
  referenced in the body. PD config lives entirely in the static path.)
- _sanitize_vllm_router_args / _vllm_router_args_from_cli helpers
  (negative-int CLI sanitisation; disable_health_check when supported).
- start_rollout_servers rewired for use_static_pd_router + engine_router_ip.

Weight sync (update_weight_from_tensor.py +140, _from_distributed.py +11):
- packed / tensor_sizes params paired with update_weights_from_ipc_handles.

Tests: DP / PD / DP+EP-r3 integration tests; test_vllm_engine updates.

Excluded vs upstream #108:
- vime_plugins/megatron_bridge/glm4v_moe.py changes dropped (per request).
- Removed 3 #125-labelled engine tests + the wake_up weight-transfer-timeout
  assertion (user's staged change). NOTE: the underlying feature code is
  still live, so those tests would have passed — flagged for review.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request Jun 3, 2026
…outer

Port of PR #108 onto vime/main (post-#127 topology + #105 slime/->vime/
rename). Carries the full feature set #108 accumulated, including the
DP+EP work formerly tracked by the now-closed #125.

Engine (vllm_engine.py):
- DP+EP sizing/launch: tp = gpus_per_engine // (pp * dp); EP via
  --enable-expert-parallel; single-node DP uses --data-parallel-backend mp.
- Conservative DP+EP sleep-mode defaults (no-async-scheduling,
  enforce-eager, distributed-timeout-seconds=1800; VLLM_ENGINE_READY_TIMEOUT_S
  extended for DP) to avoid cuMemcpyDtoDAsync segfaults, honouring user
  overrides.
- PD NIXL wiring: --kv-transfer-config NixlConnector (kv_both),
  VLLM_NIXL_SIDE_CHANNEL_HOST/PORT, disaggregation_bootstrap_port in
  _compute_server_args + [vllm-topo] server_args trace.

Router/rollout (rollout.py):
- _launch_static_pd_router: SkyRL-style static prefill/decode URLs for PD.
- _start_router is now non-PD only; dropped the has_pd_disaggregation param.
  (Fixes a latent NameError #108 left: the dropped param was still
  referenced in the body. PD config lives entirely in the static path.)
- _sanitize_vllm_router_args / _vllm_router_args_from_cli helpers
  (negative-int CLI sanitisation; disable_health_check when supported).
- start_rollout_servers rewired for use_static_pd_router + engine_router_ip.

Weight sync (update_weight_from_tensor.py +140, _from_distributed.py +11):
- packed / tensor_sizes params paired with update_weights_from_ipc_handles.

Tests: DP / PD / DP+EP-r3 integration tests; test_vllm_engine updates.

Excluded vs upstream #108:
- vime_plugins/megatron_bridge/glm4v_moe.py changes dropped (per request).
- Removed 3 #125-labelled engine tests + the wake_up weight-transfer-timeout
  assertion (user's staged change). NOTE: the underlying feature code is
  still live, so those tests would have passed — flagged for review.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request Jun 3, 2026
…outer

Port of PR #108 onto vime/main (post-#127 topology + #105 slime/->vime/
rename). Carries the full feature set #108 accumulated, including the
DP+EP work formerly tracked by the now-closed #125.

Engine (vllm_engine.py):
- DP+EP sizing/launch: tp = gpus_per_engine // (pp * dp); EP via
  --enable-expert-parallel; single-node DP uses --data-parallel-backend mp.
- Conservative DP+EP sleep-mode defaults (no-async-scheduling,
  enforce-eager, distributed-timeout-seconds=1800; VLLM_ENGINE_READY_TIMEOUT_S
  extended for DP) to avoid cuMemcpyDtoDAsync segfaults, honouring user
  overrides.
- PD NIXL wiring: --kv-transfer-config NixlConnector (kv_both),
  VLLM_NIXL_SIDE_CHANNEL_HOST/PORT, disaggregation_bootstrap_port in
  _compute_server_args + [vllm-topo] server_args trace.

Router/rollout (rollout.py):
- _launch_static_pd_router: SkyRL-style static prefill/decode URLs for PD.
- _start_router is now non-PD only; dropped the has_pd_disaggregation param.
  (Fixes a latent NameError #108 left: the dropped param was still
  referenced in the body. PD config lives entirely in the static path.)
- _sanitize_vllm_router_args / _vllm_router_args_from_cli helpers
  (negative-int CLI sanitisation; disable_health_check when supported).
- start_rollout_servers rewired for use_static_pd_router + engine_router_ip.

Weight sync (update_weight_from_tensor.py +140, _from_distributed.py +11):
- packed / tensor_sizes params paired with update_weights_from_ipc_handles.

Tests: DP / PD / DP+EP-r3 integration tests; test_vllm_engine updates.

Excluded vs upstream #108:
- vime_plugins/megatron_bridge/glm4v_moe.py changes dropped (per request).
- Removed 3 #125-labelled engine tests + the wake_up weight-transfer-timeout
  assertion (user's staged change). NOTE: the underlying feature code is
  still live, so those tests would have passed — flagged for review.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request Jun 4, 2026
…is the real guard)

The colocate DP+EP + sleep/wake crash at update_weights (cuMemcpyDtoDAsync
segfault / cublas CUBLAS_STATUS_EXECUTION_FAILED -> illegal memory access) is
fixed at the root by FIX A (execute_dummy_batch skip while the engine is asleep
/ only partially woken / mid weight-update), not by forcing eager + sync
scheduling. Drop the PR #125 --enforce-eager / --no-async-scheduling injection
so cudagraph + async scheduling stay enabled; keep the longer
--distributed-timeout-seconds for the weight-sync window.

Verified on gb200 2-node 35B PD: iter0 weight-sync + rollout + train pass with
cudagraph FULL + async scheduling enabled.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request Jun 4, 2026
…usion, VLLM_ENGINE_READY_TIMEOUT_S)

- model_provider.py: remove provider.gradient_accumulation_fusion assignment — main already
  sets it via #101 (model_provider.py:99); #108's copy is a duplicate from a stale base.
- vllm_engine.py: remove the dp_size>1 VLLM_ENGINE_READY_TIMEOUT_S=1800 env (PR #125 carryover);
  the longer --distributed-timeout-seconds already covers the colocate readiness window.

No behavior change on the validated path. Kept the router pd_disaggregation/disable_circuit_breaker
hasattr shims — those fields exist in the installed vllm-router and were active in the green run.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@princepride

Copy link
Copy Markdown
Collaborator Author

Closed because #108 has covered.

@princepride
princepride deleted the feature/rollout-dp-ep branch June 5, 2026 08:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants