[Feature] support rollout DP+EP for vLLM - #125
Conversation
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>
There was a problem hiding this comment.
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.
| 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) |
There was a problem hiding this comment.
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.
| 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) |
| denom = pp * dp | ||
| if gpus_per_engine % denom != 0: |
There was a problem hiding this comment.
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.
| 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: |
| 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): |
There was a problem hiding this comment.
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.
| 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): |
There was a problem hiding this comment.
Why disable async-scheduling and enforce eager here?
…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>
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>
…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>
…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>
…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>
…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>
…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>
|
Closed because #108 has covered. |
Wire vLLM rollout data parallelism into TP sizing and use conservative DP sleep-mode defaults so Qwen3 DP+EP DeepEP validation can run reliably.