[Core] Sync DP state on the first step of a wave (#52957)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Sep 6, 2026 at 10:16 UTC
039ea82660621a9668c7d7e81c7144b6f2f75011
4 files changed
+65
-3
tests/v1/engine/test_engine_core.py
+54
-1
@@ -16,6 +16,7 @@ from vllm.config import (
16
ECTransferConfig,
17
KVTransferConfig,
18
ModelConfig,
19
+ ParallelConfig,
20
SchedulerConfig,
21
VllmConfig,
22
)
@@ -23,7 +24,7 @@ from vllm.engine.arg_utils import EngineArgs
24
from vllm.platforms import current_platform
25
from vllm.utils.torch_utils import set_default_torch_num_threads
26
from vllm.v1.engine import EngineCoreRequest
26
-from vllm.v1.engine.core import EngineCore, EngineCoreProc
27
+from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc
28
from vllm.v1.executor.abstract import Executor
29
from vllm.v1.executor.uniproc_executor import UniProcExecutor
30
from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -610,6 +611,58 @@ def test_encoder_instance_zero_kv_cache(
611
)
612
613
614
+def _cadenced_dp_engine_core(
615
+ monkeypatch, results: list[tuple[bool, bool]], dp_sync_interval: int = 32
616
+):
617
+ """A bare DPEngineCoreProc whose all-reduce is scripted by `results`;
618
+ returns the core and the step numbers at which the all-reduce ran."""
619
+ core = object.__new__(DPEngineCoreProc)
620
+ core.dp_group = object()
621
+ core.dp_sync_interval = dp_sync_interval
622
+ core.step_counter = 0
623
+ core.pending_pause = False
624
+ core.ignore_start_dp_wave = False
625
+ synced: list[int] = []
626
+
627
+ def scripted_sync(dp_group, has_unfinished, pending_pause):
628
+ synced.append(core.step_counter)
629
+ return results.pop(0)
630
+
631
+ monkeypatch.setattr(ParallelConfig, "sync_dp_state", staticmethod(scripted_sync))
632
+ return core, synced
633
+
634
+
635
+def test_dp_sync_interval_default_is_16():
636
+ """Regression pin: lowering the default narrows the mid-wave pause tail
637
+ for async RL, where the engine is rarely idle when pause lands."""
638
+ assert ParallelConfig.dp_sync_interval == 16
639
+
640
+
641
+@pytest.mark.parametrize("dp_sync_interval", [32, 16])
642
+def test_dp_sync_interval_normal_wave(monkeypatch, dp_sync_interval: int):
643
+ """First all-reduce of a wave after one step, then every N."""
644
+ core, synced = _cadenced_dp_engine_core(
645
+ monkeypatch, [(True, False)] * 3, dp_sync_interval=dp_sync_interval
646
+ )
647
+ for _ in range(dp_sync_interval * 2 + 6):
648
+ assert DPEngineCoreProc._has_global_unfinished_reqs(core, True)
649
+ assert synced == [1, dp_sync_interval, dp_sync_interval * 2]
650
+
651
+
652
+def test_dp_sync_interval_idle_pause_consensus_on_first_step(monkeypatch):
653
+ """A pause of an idle engine arms every rank before its kick-started
654
+ first step, so the step-1 sync reaches consensus after one dummy batch
655
+ regardless of the configured cadence."""
656
+ core, synced = _cadenced_dp_engine_core(
657
+ monkeypatch, [(False, True)], dp_sync_interval=32
658
+ )
659
+ core.pending_pause = True
660
+ assert DPEngineCoreProc._has_global_unfinished_reqs(core, False) is False
661
+ assert synced == [1]
662
+ assert core.ignore_start_dp_wave
663
+ assert not core.pending_pause
664
+
665
+
666
def _pausable_engine_core_proc() -> EngineCoreProc:
667
"""A bare EngineCoreProc holding just the state pause_scheduler touches."""
668
core = object.__new__(EngineCoreProc)
vllm/config/parallel.py
+2
@@ -142,6 +142,8 @@ class ParallelConfig:
142
"""IP of the data parallel master."""
143
data_parallel_rpc_port: int = Field(default=29550, ge=1, le=65535)
144
"""Fixed port for data parallel messaging, shared by all nodes."""
145
+ dp_sync_interval: int = Field(default=16, ge=1)
146
+ """Steps between DP finish-sync all-reduces; must match across DP ranks."""
147
data_parallel_master_port: int = 29500
148
"""Port of the data parallel master."""
149
data_parallel_backend: DataParallelBackend = "mp"
vllm/engine/arg_utils.py
+6
@@ -508,6 +508,7 @@ class EngineArgs:
508
enable_dbo: bool = ParallelConfig.enable_dbo
509
ubatch_size: int = ParallelConfig.ubatch_size
510
dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold
511
+ dp_sync_interval: int = ParallelConfig.dp_sync_interval
512
dbo_prefill_token_threshold: int = ParallelConfig.dbo_prefill_token_threshold
513
disable_nccl_for_dp_synchronization: bool | None = (
514
ParallelConfig.disable_nccl_for_dp_synchronization
@@ -1190,6 +1191,10 @@ class EngineArgs:
1191
"--dbo-decode-token-threshold",
1192
**parallel_kwargs["dbo_decode_token_threshold"],
1193
)
1194
+ parallel_group.add_argument(
1195
+ "--dp-sync-interval",
1196
+ **parallel_kwargs["dp_sync_interval"],
1197
+ )
1198
parallel_group.add_argument(
1199
"--dbo-prefill-token-threshold",
1200
**parallel_kwargs["dbo_prefill_token_threshold"],
@@ -2307,6 +2312,7 @@ class EngineArgs:
2312
enable_dbo=self.enable_dbo,
2313
ubatch_size=self.ubatch_size,
2314
dbo_decode_token_threshold=self.dbo_decode_token_threshold,
2315
+ dp_sync_interval=self.dp_sync_interval,
2316
dbo_prefill_token_threshold=self.dbo_prefill_token_threshold,
2317
disable_nccl_for_dp_synchronization=self.disable_nccl_for_dp_synchronization,
2318
enable_eplb=self.enable_eplb,
vllm/v1/engine/core.py
+3
-2
@@ -2032,6 +2032,7 @@ class DPEngineCoreProc(EngineCoreProc):
2032
2033
scheduler_config = vllm_config.scheduler_config
2034
self.prefill_schedule_interval = scheduler_config.prefill_schedule_interval
2035
+ self.dp_sync_interval = vllm_config.parallel_config.dp_sync_interval
2036
2037
# Counts forward-passes of the model so that we can synchronize
2038
# finished with DP peers every N steps.
@@ -2280,9 +2281,9 @@ class DPEngineCoreProc(EngineCoreProc):
2281
raise SystemExit
2282
2283
def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
2283
- # Optimization - only perform finish-sync all-reduce every 32 steps.
2284
+ # Sync step 1 too: an idle pause needs one dummy batch, not a full interval.
2285
self.step_counter += 1
2285
- if self.step_counter % 32 != 0:
2286
+ if self.step_counter != 1 and self.step_counter % self.dp_sync_interval != 0:
2287
return True
2288
2289
has_unfinished, pause_consensus = ParallelConfig.sync_dp_state(