A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 14:00 UTC
[1/N] Elastic EP Milestone 2 (#34861)
Signed-off-by: Yongji Wu <wuyongji317@gmail.com> Signed-off-by: Itay Alroy <ialroy@nvidia.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Ron Tourgeman <rtourgeman@nvidia.com> Co-authored-by: Yongji Wu <wuyongji317@gmail.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Ron Tourgeman <rtourgeman@nvidia.com>
Itay Alroy committed
Feb 28, 2026 at 06:46 UTC
dea268336fb51b5bc342dca29189f3d4440ca2a0
53 files changed
+3603
-1006
.buildkite/test_areas/expert_parallelism.yaml
+16
-1
@@ -20,4 +20,19 @@ steps:
20
- tests/distributed/test_eplb_execute.py
21
commands:
22
- pytest -v -s distributed/test_eplb_execute.py
23
- - pytest -v -s distributed/test_eplb_spec_decode.py
\ No newline at end of file
23
+ - pytest -v -s distributed/test_eplb_spec_decode.py
24
+
25
+- label: Elastic EP Scaling Test
26
+ timeout_in_minutes: 20
27
+ device: b200
28
+ optional: true
29
+ working_dir: "/vllm-workspace/tests"
30
+ num_devices: 4
31
+ source_file_dependencies:
32
+ - vllm/distributed/
33
+ - vllm/engine/
34
+ - vllm/executor/
35
+ - vllm/compilation/
36
+ - tests/distributed/
37
+ commands:
38
+ - pytest -v -s distributed/test_elastic_ep.py
tests/compile/passes/distributed/test_async_tp.py
+3
-5
@@ -316,7 +316,6 @@ def async_tp_pass_on_test_model(
316
317
# initialize distributed
318
init_distributed_environment()
319
- initialize_model_parallel(tensor_model_parallel_size=world_size)
319
320
# configure vllm config for SequenceParallelismPass
321
vllm_config = VllmConfig()
@@ -334,11 +333,10 @@ def async_tp_pass_on_test_model(
333
model=model_name, trust_remote_code=True, dtype=dtype, seed=42
334
)
335
337
- async_tp_pass = AsyncTPPass(vllm_config)
338
-
339
- # Set the global vllm_config for TestBackend which calls
340
- # get_current_vllm_config()
336
with set_current_vllm_config(vllm_config):
337
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
338
+
339
+ async_tp_pass = AsyncTPPass(vllm_config)
340
backend = TestBackend(async_tp_pass)
341
342
assert (
tests/compile/passes/distributed/test_fusion_all_reduce.py
+1
-1
@@ -278,7 +278,6 @@ def all_reduce_fusion_pass_on_test_model(
278
)
279
280
init_distributed_environment()
281
- initialize_model_parallel(tensor_model_parallel_size=world_size)
281
282
custom_ops = []
283
if enable_rms_norm_custom_op:
@@ -304,6 +303,7 @@ def all_reduce_fusion_pass_on_test_model(
303
model=model_name, trust_remote_code=True, dtype=dtype, seed=42
304
)
305
with set_current_vllm_config(vllm_config):
306
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
307
all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)
308
noop_pass = NoOpEliminationPass(vllm_config)
309
func_pass = FixFunctionalizationPass(vllm_config)
tests/compile/passes/distributed/test_sequence_parallelism.py
+1
-1
@@ -242,7 +242,6 @@ def sequence_parallelism_pass_on_test_model(
242
243
# initialize distributed
244
init_distributed_environment()
245
- initialize_model_parallel(tensor_model_parallel_size=world_size)
245
246
# configure vllm config for SequenceParallelismPass
247
custom_ops_list = custom_ops.split(",") if custom_ops else []
@@ -272,6 +271,7 @@ def sequence_parallelism_pass_on_test_model(
271
)
272
273
with set_current_vllm_config(vllm_config):
274
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
275
noop_pass = NoOpEliminationPass(vllm_config)
276
sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
277
cleanup_pass = PostCleanupPass(vllm_config)
tests/conftest.py
+13
-9
@@ -176,16 +176,20 @@ def init_test_http_connection():
176
177
@pytest.fixture
178
def dist_init():
179
+ from tests.utils import ensure_current_vllm_config
180
+
181
temp_file = tempfile.mkstemp()[1]
180
- init_distributed_environment(
181
- world_size=1,
182
- rank=0,
183
- distributed_init_method=f"file://{temp_file}",
184
- local_rank=0,
185
- backend="nccl",
186
- )
187
- initialize_model_parallel(1, 1)
188
- yield
182
+
183
+ with ensure_current_vllm_config():
184
+ init_distributed_environment(
185
+ world_size=1,
186
+ rank=0,
187
+ distributed_init_method=f"file://{temp_file}",
188
+ local_rank=0,
189
+ backend="nccl",
190
+ )
191
+ initialize_model_parallel(1, 1)
192
+ yield
193
cleanup_dist_env_and_memory()
194
195
tests/distributed/eplb_utils.py
+6
-1
@@ -7,6 +7,7 @@ import random
7
import torch
8
import torch.multiprocessing as mp
9
10
+from vllm.config import VllmConfig, set_current_vllm_config
11
from vllm.distributed.parallel_state import (
12
init_distributed_environment,
13
)
@@ -42,7 +43,11 @@ def set_env_vars_and_device(env: dict[str, str]) -> None:
43
local_rank = os.environ["LOCAL_RANK"]
44
device = torch.device(f"cuda:{local_rank}")
45
torch.cuda.set_device(device)
45
- init_distributed_environment()
46
+
47
+ # Create a minimal vllm config for init_distributed_environment
48
+ vllm_config = VllmConfig()
49
+ with set_current_vllm_config(vllm_config):
50
+ init_distributed_environment()
51
52
# Ensure each worker process has the same random seed
53
random.seed(42)
tests/distributed/test_elastic_ep.py
new
+202
@@ -0,0 +1,202 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import os
5
+import subprocess
6
+import time
7
+
8
+import pytest
9
+import requests
10
+
11
+from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k
12
+from ..utils import RemoteOpenAIServer, multi_gpu_test
13
+
14
+
15
+@pytest.fixture(autouse=True)
16
+def cleanup_ray_between_tests():
17
+ """Force-stop any lingering Ray processes between tests."""
18
+ subprocess.run(["ray", "stop", "--force"], timeout=30, capture_output=True)
19
+ time.sleep(5)
20
+ yield
21
+
22
+
23
+MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat"
24
+
25
+NUM_GSM8K_QUESTIONS = 256
26
+EXPECTED_ACCURACY = 0.58
27
+ACCURACY_TOL = 0.08
28
+MAX_NUM_SEQS = 32
29
+
30
+
31
+def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool:
32
+ url = server.url_for("scale_elastic_ep")
33
+ payload = {"new_data_parallel_size": new_dp_size}
34
+ headers = {"Content-Type": "application/json"}
35
+
36
+ try:
37
+ response = requests.post(url, json=payload, headers=headers, timeout=300)
38
+ return response.status_code == 200
39
+ except requests.exceptions.RequestException:
40
+ return False
41
+
42
+
43
+def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
44
+ assert server.port is not None
45
+ result = evaluate_gsm8k(
46
+ num_questions=NUM_GSM8K_QUESTIONS,
47
+ host=f"http://{server.host}",
48
+ port=server.port,
49
+ )
50
+ accuracy = result["accuracy"]
51
+ print(
52
+ f"[{stage}] GSM8K accuracy: {accuracy:.3f} "
53
+ f"({result['num_questions']} questions)"
54
+ )
55
+ assert accuracy >= EXPECTED_ACCURACY, (
56
+ f"[{stage}] GSM8K accuracy {accuracy:.3f} is below "
57
+ f"expected threshold {EXPECTED_ACCURACY}"
58
+ )
59
+ return accuracy
60
+
61
+
62
+@multi_gpu_test(num_gpus=4)
63
+def test_elastic_ep_scaling():
64
+ vllm_serve_args = [
65
+ "--trust-remote-code",
66
+ "--tensor-parallel-size",
67
+ "1",
68
+ "--gpu-memory-utilization",
69
+ "0.8",
70
+ "--max-model-len",
71
+ "4096",
72
+ "--max-num-seqs",
73
+ str(MAX_NUM_SEQS),
74
+ "--enable-expert-parallel",
75
+ "--all2all-backend",
76
+ "allgather_reducescatter",
77
+ "--enable-elastic-ep",
78
+ "--enable-eplb",
79
+ "--eplb-config.num_redundant_experts",
80
+ "0",
81
+ "--data-parallel-backend",
82
+ "ray",
83
+ "--data-parallel-size",
84
+ "2",
85
+ "--api-server-count",
86
+ "1",
87
+ ]
88
+
89
+ leader_address = os.environ.get("LEADER_ADDRESS")
90
+ if leader_address:
91
+ vllm_serve_args.extend(["--data-parallel-address", leader_address])
92
+
93
+ with RemoteOpenAIServer(
94
+ MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
95
+ ) as server:
96
+ initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")
97
+
98
+ assert _send_scale_command(server, 4)
99
+ time.sleep(10)
100
+ scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (4 GPUs)")
101
+
102
+ assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
103
+ f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "
104
+ f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
105
+ )
106
+
107
+ assert _send_scale_command(server, 2)
108
+ time.sleep(5)
109
+ scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
110
+
111
+ assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
112
+ f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "
113
+ f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
114
+ )
115
+
116
+ print("\nAccuracy Summary:")
117
+ print(f" Initial: {initial_accuracy:.3f}")
118
+ print(
119
+ f" Scale up: {scale_up_accuracy:.3f} "
120
+ f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})"
121
+ )
122
+ print(
123
+ f" Scale down: {scale_down_accuracy:.3f} "
124
+ f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"
125
+ )
126
+ print(f" Tolerance: {ACCURACY_TOL:.3f}")
127
+
128
+
129
+@multi_gpu_test(num_gpus=4)
130
+def test_elastic_ep_scaling_uneven():
131
+ """Test scale up with uneven worker distribution.
132
+
133
+ This tests the case where num_new_workers % old_dp_size != 0,
134
+ specifically 2 -> 3 where remainder = 1 % 2 = 1.
135
+ This exercises the remainder handling in sender-receiver pairing.
136
+ """
137
+ vllm_serve_args = [
138
+ "--trust-remote-code",
139
+ "--tensor-parallel-size",
140
+ "1",
141
+ "--gpu-memory-utilization",
142
+ "0.8",
143
+ "--max-model-len",
144
+ "4096",
145
+ "--max-num-seqs",
146
+ str(MAX_NUM_SEQS),
147
+ "--enable-expert-parallel",
148
+ "--all2all-backend",
149
+ "allgather_reducescatter",
150
+ "--enable-elastic-ep",
151
+ "--enable-eplb",
152
+ "--eplb-config.num_redundant_experts",
153
+ "0",
154
+ "--data-parallel-backend",
155
+ "ray",
156
+ "--data-parallel-size",
157
+ "2",
158
+ "--api-server-count",
159
+ "1",
160
+ ]
161
+
162
+ leader_address = os.environ.get("LEADER_ADDRESS")
163
+ if leader_address:
164
+ vllm_serve_args.extend(["--data-parallel-address", leader_address])
165
+
166
+ with RemoteOpenAIServer(
167
+ MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
168
+ ) as server:
169
+ initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")
170
+
171
+ # Scale 2 -> 3: This has remainder = 1 % 2 = 1
172
+ # Tests uneven sender-receiver pairing
173
+ assert _send_scale_command(server, 3)
174
+ time.sleep(10)
175
+ scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)")
176
+
177
+ assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
178
+ f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "
179
+ f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
180
+ )
181
+
182
+ # Scale back down to 2
183
+ assert _send_scale_command(server, 2)
184
+ time.sleep(5)
185
+ scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
186
+
187
+ assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
188
+ f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "
189
+ f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
190
+ )
191
+
192
+ print("\nAccuracy Summary (Uneven Scaling):")
193
+ print(f" Initial: {initial_accuracy:.3f}")
194
+ print(
195
+ f" Scale up: {scale_up_accuracy:.3f} "
196
+ f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})"
197
+ )
198
+ print(
199
+ f" Scale down: {scale_down_accuracy:.3f} "
200
+ f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"
201
+ )
202
+ print(f" Tolerance: {ACCURACY_TOL:.3f}")
tests/distributed/test_eplb_execute.py
+245
-224
@@ -8,6 +8,7 @@ import pytest
8
import torch
9
import torch.distributed
10
11
+from vllm.config import VllmConfig, set_current_vllm_config
12
from vllm.distributed.eplb.rebalance_execute import (
13
move_from_buffer,
14
rearrange_expert_weights_inplace,
@@ -244,91 +245,96 @@ def _test_async_transfer_layer_without_mtp_worker(
245
num_logical_experts: int,
246
) -> None:
247
set_env_vars_and_device(env)
247
- ensure_model_parallel_initialized(
248
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
249
- )
248
251
- tp_group = get_tp_group()
252
- ep_group = tp_group.device_group
253
- ep_rank = torch.distributed.get_rank()
254
- device = torch.device(f"cuda:{ep_rank}")
249
+ vllm_config = VllmConfig()
250
+ vllm_config.parallel_config.tensor_parallel_size = world_size
251
256
- total_physical_experts = world_size * num_local_experts
257
- hidden_sizes = [16, 32]
252
+ with set_current_vllm_config(vllm_config):
253
+ ensure_model_parallel_initialized(
254
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
255
+ )
256
259
- redundancy_config = create_redundancy_config(
260
- num_logical_experts,
261
- total_physical_experts,
262
- )
263
- old_indices = create_expert_indices_with_redundancy(
264
- num_layers,
265
- num_logical_experts,
266
- total_physical_experts,
267
- redundancy_config,
268
- )
257
+ tp_group = get_tp_group()
258
+ ep_group = tp_group.device_group
259
+ ep_rank = torch.distributed.get_rank()
260
+ device = torch.device(f"cuda:{ep_rank}")
261
270
- new_redundancy_config = create_redundancy_config(
271
- num_logical_experts,
272
- total_physical_experts,
273
- )
274
- new_indices = create_expert_indices_with_redundancy(
275
- num_layers,
276
- num_logical_experts,
277
- total_physical_experts,
278
- new_redundancy_config,
279
- )
262
+ total_physical_experts = world_size * num_local_experts
263
+ hidden_sizes = [16, 32]
264
281
- expert_weights = create_expert_weights(
282
- num_layers,
283
- num_local_experts,
284
- hidden_sizes,
285
- ep_rank,
286
- device,
287
- old_indices,
288
- )
289
- old_indices_cpu = old_indices.cpu()
290
- new_indices_cpu = new_indices.cpu()
265
+ redundancy_config = create_redundancy_config(
266
+ num_logical_experts,
267
+ total_physical_experts,
268
+ )
269
+ old_indices = create_expert_indices_with_redundancy(
270
+ num_layers,
271
+ num_logical_experts,
272
+ total_physical_experts,
273
+ redundancy_config,
274
+ )
275
292
- expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
293
- cuda_stream = torch.cuda.Stream(device=device)
276
+ new_redundancy_config = create_redundancy_config(
277
+ num_logical_experts,
278
+ total_physical_experts,
279
+ )
280
+ new_indices = create_expert_indices_with_redundancy(
281
+ num_layers,
282
+ num_logical_experts,
283
+ total_physical_experts,
284
+ new_redundancy_config,
285
+ )
286
295
- for layer_idx in range(num_layers):
296
- is_unchanged, is_received_locally, recv_metadata = asyncio.run(
297
- transfer_layer(
298
- old_layer_indices=old_indices_cpu[layer_idx],
299
- new_layer_indices=new_indices_cpu[layer_idx],
287
+ expert_weights = create_expert_weights(
288
+ num_layers,
289
+ num_local_experts,
290
+ hidden_sizes,
291
+ ep_rank,
292
+ device,
293
+ old_indices,
294
+ )
295
+ old_indices_cpu = old_indices.cpu()
296
+ new_indices_cpu = new_indices.cpu()
297
+
298
+ expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
299
+ cuda_stream = torch.cuda.Stream(device=device)
300
+
301
+ for layer_idx in range(num_layers):
302
+ is_unchanged, is_received_locally, recv_metadata = asyncio.run(
303
+ transfer_layer(
304
+ old_layer_indices=old_indices_cpu[layer_idx],
305
+ new_layer_indices=new_indices_cpu[layer_idx],
306
+ expert_weights=expert_weights[layer_idx],
307
+ expert_weights_buffer=expert_buffer,
308
+ ep_group=ep_group,
309
+ cuda_stream=cuda_stream,
310
+ )
311
+ )
312
+ cuda_stream.synchronize()
313
+ move_from_buffer(
314
expert_weights=expert_weights[layer_idx],
301
- expert_weights_buffer=expert_buffer,
302
- ep_group=ep_group,
303
- cuda_stream=cuda_stream,
315
+ expert_weights_buffers=expert_buffer,
316
+ is_unchanged=is_unchanged,
317
+ is_received_locally=is_received_locally,
318
+ recv_metadata=recv_metadata,
319
+ new_indices=new_indices_cpu[layer_idx].numpy(),
320
+ ep_rank=ep_rank,
321
)
322
+
323
+ verify_expert_weights_after_shuffle(
324
+ expert_weights,
325
+ new_indices,
326
+ hidden_sizes,
327
+ ep_rank,
328
+ num_local_experts,
329
)
306
- cuda_stream.synchronize()
307
- move_from_buffer(
308
- expert_weights=expert_weights[layer_idx],
309
- expert_weights_buffers=expert_buffer,
310
- is_unchanged=is_unchanged,
311
- is_received_locally=is_received_locally,
312
- recv_metadata=recv_metadata,
313
- new_indices=new_indices_cpu[layer_idx].numpy(),
314
- ep_rank=ep_rank,
330
+ verify_redundant_experts_have_same_weights(
331
+ expert_weights,
332
+ new_indices,
333
+ hidden_sizes,
334
+ world_size,
335
+ num_local_experts,
336
)
337
317
- verify_expert_weights_after_shuffle(
318
- expert_weights,
319
- new_indices,
320
- hidden_sizes,
321
- ep_rank,
322
- num_local_experts,
323
- )
324
- verify_redundant_experts_have_same_weights(
325
- expert_weights,
326
- new_indices,
327
- hidden_sizes,
328
- world_size,
329
- num_local_experts,
330
- )
331
-
338
339
def _test_rearrange_expert_weights_with_redundancy(
340
env, world_size, num_layers, num_local_experts, num_logical_experts
@@ -336,71 +342,76 @@ def _test_rearrange_expert_weights_with_redundancy(
342
# Initialize model parallel (using tensor parallel as an entrypoint
343
# to expert parallel)
344
set_env_vars_and_device(env)
339
- ensure_model_parallel_initialized(
340
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
341
- )
345
343
- ep_group = get_tp_group().cpu_group
344
- ep_rank = torch.distributed.get_rank()
345
- device = torch.device(f"cuda:{ep_rank}")
346
+ vllm_config = VllmConfig()
347
+ vllm_config.parallel_config.tensor_parallel_size = world_size
348
347
- # Test parameters
348
- total_physical_experts = world_size * num_local_experts
349
- hidden_sizes = [32, 64] # Two different weight matrices
349
+ with set_current_vllm_config(vllm_config):
350
+ ensure_model_parallel_initialized(
351
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
352
+ )
353
351
- # Create old expert indices (with redundancy)
352
- redundancy_config = create_redundancy_config(
353
- num_logical_experts, total_physical_experts
354
- )
354
+ ep_group = get_tp_group().cpu_group
355
+ ep_rank = torch.distributed.get_rank()
356
+ device = torch.device(f"cuda:{ep_rank}")
357
356
- old_indices = create_expert_indices_with_redundancy(
357
- num_layers,
358
- num_logical_experts,
359
- total_physical_experts,
360
- redundancy_config,
361
- )
358
+ # Test parameters
359
+ total_physical_experts = world_size * num_local_experts
360
+ hidden_sizes = [32, 64] # Two different weight matrices
361
363
- # Create new expert indices (with redundancy)
364
- new_redundancy_config = create_redundancy_config(
365
- num_logical_experts, total_physical_experts
366
- )
367
- new_indices = create_expert_indices_with_redundancy(
368
- num_layers,
369
- num_logical_experts,
370
- total_physical_experts,
371
- new_redundancy_config,
372
- )
362
+ # Create old expert indices (with redundancy)
363
+ redundancy_config = create_redundancy_config(
364
+ num_logical_experts, total_physical_experts
365
+ )
366
374
- # Create expert weights
375
- expert_weights = create_expert_weights(
376
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
377
- )
367
+ old_indices = create_expert_indices_with_redundancy(
368
+ num_layers,
369
+ num_logical_experts,
370
+ total_physical_experts,
371
+ redundancy_config,
372
+ )
373
379
- # Execute weight rearrangement
380
- rearrange_expert_weights_inplace(
381
- old_indices,
382
- new_indices,
383
- expert_weights,
384
- ep_group,
385
- is_profile=False,
386
- )
374
+ # Create new expert indices (with redundancy)
375
+ new_redundancy_config = create_redundancy_config(
376
+ num_logical_experts, total_physical_experts
377
+ )
378
+ new_indices = create_expert_indices_with_redundancy(
379
+ num_layers,
380
+ num_logical_experts,
381
+ total_physical_experts,
382
+ new_redundancy_config,
383
+ )
384
388
- # Verify the rearrangement result
389
- verify_expert_weights_after_shuffle(
390
- expert_weights,
391
- new_indices,
392
- hidden_sizes,
393
- ep_rank,
394
- num_local_experts,
395
- )
385
+ # Create expert weights
386
+ expert_weights = create_expert_weights(
387
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
388
+ )
389
397
- verify_redundant_experts_have_same_weights(
398
- expert_weights,
399
- new_indices,
400
- hidden_sizes,
401
- world_size,
402
- num_local_experts,
403
- )
390
+ # Execute weight rearrangement
391
+ rearrange_expert_weights_inplace(
392
+ old_indices,
393
+ new_indices,
394
+ expert_weights,
395
+ ep_group,
396
+ is_profile=False,
397
+ )
398
+
399
+ # Verify the rearrangement result
400
+ verify_expert_weights_after_shuffle(
401
+ expert_weights,
402
+ new_indices,
403
+ hidden_sizes,
404
+ ep_rank,
405
+ num_local_experts,
406
+ )
407
+
408
+ verify_redundant_experts_have_same_weights(
409
+ expert_weights,
410
+ new_indices,
411
+ hidden_sizes,
412
+ world_size,
413
+ num_local_experts,
414
+ )
415
416
417
@pytest.mark.parametrize(
@@ -444,58 +455,63 @@ def test_rearrange_expert_weights_with_redundancy(
455
456
def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
457
set_env_vars_and_device(env)
447
- ensure_model_parallel_initialized(
448
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
449
- )
458
451
- ep_group = get_tp_group().cpu_group
452
- ep_rank = torch.distributed.get_rank()
453
- device = torch.device(f"cuda:{ep_rank}")
459
+ vllm_config = VllmConfig()
460
+ vllm_config.parallel_config.tensor_parallel_size = world_size
461
455
- num_layers = 2
456
- num_local_experts = 2
457
- total_physical_experts = world_size * num_local_experts
458
- num_logical_experts = total_physical_experts // 2 # Some redundancy
459
- hidden_sizes = [32, 64]
462
+ with set_current_vllm_config(vllm_config):
463
+ ensure_model_parallel_initialized(
464
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
465
+ )
466
461
- # Create redundancy configuration
462
- redundancy_config = [2] * num_logical_experts
467
+ ep_group = get_tp_group().cpu_group
468
+ ep_rank = torch.distributed.get_rank()
469
+ device = torch.device(f"cuda:{ep_rank}")
470
464
- # Same indices - no change
465
- indices = create_expert_indices_with_redundancy(
466
- num_layers, num_logical_experts, total_physical_experts, redundancy_config
467
- )
471
+ num_layers = 2
472
+ num_local_experts = 2
473
+ total_physical_experts = world_size * num_local_experts
474
+ num_logical_experts = total_physical_experts // 2 # Some redundancy
475
+ hidden_sizes = [32, 64]
476
469
- expert_weights = create_expert_weights(
470
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
471
- )
477
+ # Create redundancy configuration
478
+ redundancy_config = [2] * num_logical_experts
479
473
- # Save original weights
474
- original_weights = []
475
- for layer_weights in expert_weights:
476
- layer_copy = []
477
- for weight in layer_weights:
478
- layer_copy.append(weight.clone())
479
- original_weights.append(layer_copy)
480
-
481
- # Execute rearrangement (should be no change)
482
- rearrange_expert_weights_inplace(
483
- indices,
484
- indices, # Same indices
485
- expert_weights,
486
- ep_group,
487
- is_profile=False,
488
- )
480
+ # Same indices - no change
481
+ indices = create_expert_indices_with_redundancy(
482
+ num_layers, num_logical_experts, total_physical_experts, redundancy_config
483
+ )
484
490
- # Verify that the weights have not changed
491
- for layer in range(num_layers):
492
- for weight_idx in range(len(hidden_sizes)):
493
- torch.testing.assert_close(
494
- expert_weights[layer][weight_idx],
495
- original_weights[layer][weight_idx],
496
- msg=f"""Layer {layer}, weight {weight_idx}
485
+ expert_weights = create_expert_weights(
486
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
487
+ )
488
+
489
+ # Save original weights
490
+ original_weights = []
491
+ for layer_weights in expert_weights:
492
+ layer_copy = []
493
+ for weight in layer_weights:
494
+ layer_copy.append(weight.clone())
495
+ original_weights.append(layer_copy)
496
+
497
+ # Execute rearrangement (should be no change)
498
+ rearrange_expert_weights_inplace(
499
+ indices,
500
+ indices, # Same indices
501
+ expert_weights,
502
+ ep_group,
503
+ is_profile=False,
504
+ )
505
+
506
+ # Verify that the weights have not changed
507
+ for layer in range(num_layers):
508
+ for weight_idx in range(len(hidden_sizes)):
509
+ torch.testing.assert_close(
510
+ expert_weights[layer][weight_idx],
511
+ original_weights[layer][weight_idx],
512
+ msg=f"""Layer {layer}, weight {weight_idx}
513
should remain unchanged""",
498
- )
514
+ )
515
516
517
@pytest.mark.parametrize(
@@ -538,64 +554,69 @@ def test_rearrange_expert_weights_no_change(world_size):
554
555
def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
556
set_env_vars_and_device(env)
541
- ensure_model_parallel_initialized(
542
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
543
- )
557
545
- ep_group = get_tp_group().cpu_group
546
- ep_rank = torch.distributed.get_rank()
547
- device = torch.device(f"cuda:{ep_rank}")
558
+ vllm_config = VllmConfig()
559
+ vllm_config.parallel_config.tensor_parallel_size = world_size
560
549
- num_layers = 1
550
- num_local_experts = 2
551
- total_physical_experts = world_size * num_local_experts
552
- num_logical_experts = total_physical_experts // 2
553
- hidden_sizes = [32]
561
+ with set_current_vllm_config(vllm_config):
562
+ ensure_model_parallel_initialized(
563
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
564
+ )
565
555
- # Create different index distributions
556
- old_redundancy = create_redundancy_config(
557
- num_logical_experts, total_physical_experts
558
- )
559
- new_redundancy = create_redundancy_config(
560
- num_logical_experts, total_physical_experts
561
- )
566
+ ep_group = get_tp_group().cpu_group
567
+ ep_rank = torch.distributed.get_rank()
568
+ device = torch.device(f"cuda:{ep_rank}")
569
563
- old_indices = create_expert_indices_with_redundancy(
564
- num_layers, num_logical_experts, total_physical_experts, old_redundancy
565
- )
566
- new_indices = create_expert_indices_with_redundancy(
567
- num_layers, num_logical_experts, total_physical_experts, new_redundancy
568
- )
570
+ num_layers = 1
571
+ num_local_experts = 2
572
+ total_physical_experts = world_size * num_local_experts
573
+ num_logical_experts = total_physical_experts // 2
574
+ hidden_sizes = [32]
575
570
- expert_weights = create_expert_weights(
571
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
572
- )
576
+ # Create different index distributions
577
+ old_redundancy = create_redundancy_config(
578
+ num_logical_experts, total_physical_experts
579
+ )
580
+ new_redundancy = create_redundancy_config(
581
+ num_logical_experts, total_physical_experts
582
+ )
583
574
- # Save original weights
575
- original_weights = []
576
- for layer_weights in expert_weights:
577
- layer_copy = []
578
- for weight in layer_weights:
579
- layer_copy.append(weight.clone())
580
- original_weights.append(layer_copy)
581
-
582
- # Execute profile mode rearrangement
583
- rearrange_expert_weights_inplace(
584
- old_indices,
585
- new_indices,
586
- expert_weights,
587
- ep_group,
588
- is_profile=True, # Profile mode
589
- )
584
+ old_indices = create_expert_indices_with_redundancy(
585
+ num_layers, num_logical_experts, total_physical_experts, old_redundancy
586
+ )
587
+ new_indices = create_expert_indices_with_redundancy(
588
+ num_layers, num_logical_experts, total_physical_experts, new_redundancy
589
+ )
590
591
- # In profile mode, the weights should remain unchanged
592
- for layer in range(num_layers):
593
- for weight_idx in range(len(hidden_sizes)):
594
- torch.testing.assert_close(
595
- expert_weights[layer][weight_idx],
596
- original_weights[layer][weight_idx],
597
- msg="In profile mode, the weights should remain unchanged",
598
- )
591
+ expert_weights = create_expert_weights(
592
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
593
+ )
594
+
595
+ # Save original weights
596
+ original_weights = []
597
+ for layer_weights in expert_weights:
598
+ layer_copy = []
599
+ for weight in layer_weights:
600
+ layer_copy.append(weight.clone())
601
+ original_weights.append(layer_copy)
602
+
603
+ # Execute profile mode rearrangement
604
+ rearrange_expert_weights_inplace(
605
+ old_indices,
606
+ new_indices,
607
+ expert_weights,
608
+ ep_group,
609
+ is_profile=True, # Profile mode
610
+ )
611
+
612
+ # In profile mode, the weights should remain unchanged
613
+ for layer in range(num_layers):
614
+ for weight_idx in range(len(hidden_sizes)):
615
+ torch.testing.assert_close(
616
+ expert_weights[layer][weight_idx],
617
+ original_weights[layer][weight_idx],
618
+ msg="In profile mode, the weights should remain unchanged",
619
+ )
620
621
622
@pytest.mark.parametrize("world_size", [2, 4])
tests/distributed/test_nccl_symm_mem_allreduce.py
+3
-1
@@ -10,6 +10,7 @@ import torch.distributed as dist
10
import torch.multiprocessing as mp
11
12
import vllm.envs as envs
13
+from tests.utils import ensure_current_vllm_config
14
from vllm.distributed import cleanup_dist_env_and_memory
15
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
16
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
@@ -51,7 +52,8 @@ def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
52
)
53
54
init_distributed_environment()
54
- initialize_model_parallel(tensor_model_parallel_size=world_size)
55
+ with ensure_current_vllm_config():
56
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
57
58
cuda_communicator = typing.cast(
59
CudaCommunicator, get_tp_group().device_communicator
tests/distributed/test_pynccl.py
+3
-1
@@ -9,6 +9,7 @@ import pytest
9
import torch
10
import torch.distributed
11
12
+from tests.utils import ensure_current_vllm_config
13
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
14
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
15
from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
@@ -112,7 +113,8 @@ def test_pynccl_multiple_allreduce():
113
@worker_fn_wrapper
114
def multiple_allreduce_with_vllm_worker_fn():
115
device = torch.device(f"cuda:{torch.distributed.get_rank()}")
115
- ensure_model_parallel_initialized(2, 2)
116
+ with ensure_current_vllm_config():
117
+ ensure_model_parallel_initialized(2, 2)
118
tensor = torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
119
with graph_capture(device=device):
120
# two tp groups can communicate independently
tests/kernels/mamba/test_mamba_mixer2.py
+3
-2
@@ -6,7 +6,7 @@ import unittest
6
import pytest
7
import torch
8
9
-from tests.utils import multi_gpu_test
9
+from tests.utils import ensure_current_vllm_config, multi_gpu_test
10
from vllm.distributed.parallel_state import (
11
init_distributed_environment,
12
initialize_model_parallel,
@@ -87,7 +87,8 @@ def mixer2_gated_norm_tensor_parallel(
87
88
# initialize distributed
89
init_distributed_environment()
90
- initialize_model_parallel(tensor_model_parallel_size=world_size)
90
+ with ensure_current_vllm_config():
91
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
92
93
# create random weights an inputs
94
weight = torch.rand((hidden_size,), dtype=dtype, device=device)
tests/lora/conftest.py
+12
-9
@@ -45,21 +45,24 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):
45
46
@pytest.fixture
47
def dist_init():
48
+ from tests.utils import ensure_current_vllm_config
49
+
50
temp_file = tempfile.mkstemp()[1]
51
52
backend = "nccl"
53
if current_platform.is_cpu() or current_platform.is_tpu():
54
backend = "gloo"
55
54
- init_distributed_environment(
55
- world_size=1,
56
- rank=0,
57
- distributed_init_method=f"file://{temp_file}",
58
- local_rank=0,
59
- backend=backend,
60
- )
61
- initialize_model_parallel(1, 1)
62
- yield
56
+ with ensure_current_vllm_config():
57
+ init_distributed_environment(
58
+ world_size=1,
59
+ rank=0,
60
+ distributed_init_method=f"file://{temp_file}",
61
+ local_rank=0,
62
+ backend=backend,
63
+ )
64
+ initialize_model_parallel(1, 1)
65
+ yield
66
cleanup_dist_env_and_memory(shutdown_ray=True)
67
68
tests/lora/test_fused_moe_lora_kernel.py
+3
-2
@@ -6,7 +6,7 @@ import random
6
import pytest
7
import torch
8
9
-from tests.utils import multi_gpu_test
9
+from tests.utils import ensure_current_vllm_config, multi_gpu_test
10
from vllm import _custom_ops as ops
11
from vllm.distributed import (
12
init_distributed_environment,
@@ -631,7 +631,8 @@ def use_fused_moe_lora_kernel_tensor_parallel(
631
local_rank=local_rank,
632
distributed_init_method=init_method,
633
)
634
- initialize_model_parallel(world_size, 1)
634
+ with ensure_current_vllm_config():
635
+ initialize_model_parallel(world_size, 1)
636
tp_size = get_tensor_model_parallel_world_size()
637
638
input_dim = K if column_parallel else N
tests/lora/test_worker.py
+4
-2
@@ -13,6 +13,7 @@ from vllm.config import (
13
ParallelConfig,
14
SchedulerConfig,
15
VllmConfig,
16
+ set_current_vllm_config,
17
)
18
from vllm.config.load import LoadConfig
19
from vllm.config.lora import LoRAConfig
@@ -77,8 +78,9 @@ def test_worker_apply_lora(qwen3_lora_files):
78
distributed_init_method=f"file://{tempfile.mkstemp()[1]}",
79
)
80
80
- worker.init_device()
81
- worker.load_model()
81
+ with set_current_vllm_config(vllm_config):
82
+ worker.init_device()
83
+ worker.load_model()
84
85
set_active_loras(worker, [])
86
assert worker.list_loras() == set()
tests/models/test_vision.py
+9
-5
@@ -6,7 +6,7 @@ import pytest
6
import torch
7
import torch.multiprocessing as mp
8
9
-from tests.utils import multi_gpu_test
9
+from tests.utils import ensure_current_vllm_config, multi_gpu_test
10
from vllm.distributed import get_tensor_model_parallel_world_size
11
from vllm.distributed.parallel_state import (
12
init_distributed_environment,
@@ -117,7 +117,8 @@ def run_dp_sharded_vision_model_vs_direct(
117
118
# initialize distributed
119
init_distributed_environment()
120
- initialize_model_parallel(tensor_model_parallel_size=world_size)
120
+ with ensure_current_vllm_config():
121
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
122
123
# Create a test input tensor
124
image_input = torch.randn(batch_size, 3, 224, 224)
@@ -302,7 +303,8 @@ def run_dp_sharded_mrope_vision_model_vs_direct(
303
304
# initialize distributed
305
init_distributed_environment()
305
- initialize_model_parallel(tensor_model_parallel_size=world_size)
306
+ with ensure_current_vllm_config():
307
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
308
309
# Create test data
310
grid_thw_list = []
@@ -377,7 +379,8 @@ def run_dp_sharded_mrope_vision_model_empty_input_worker(
379
)
380
381
init_distributed_environment()
380
- initialize_model_parallel(tensor_model_parallel_size=world_size)
382
+ with ensure_current_vllm_config():
383
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
384
385
# Create empty inputs
386
pixel_values = torch.empty((0, 768))
@@ -425,7 +428,8 @@ def run_dp_sharded_mrope_vision_model_uneven_load_worker(
428
)
429
430
init_distributed_environment()
428
- initialize_model_parallel(tensor_model_parallel_size=world_size)
431
+ with ensure_current_vllm_config():
432
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
433
434
# Create images with very different sizes
435
grid_thw_list = [
tests/utils.py
+32
-1
@@ -895,6 +895,36 @@ def compare_all_settings(
895
)
896
897
898
+@contextmanager
899
+def ensure_current_vllm_config():
900
+ """Ensures a vllm config is set for the duration of the context.
901
+
902
+ If a config is already set, this is a no-op. Otherwise, it creates a default
903
+ VllmConfig and sets it for the duration of the context.
904
+
905
+ Used for tests that call functions which require a vllm config but don't
906
+ need a specific config.
907
+
908
+ Example:
909
+ with ensure_current_vllm_config():
910
+ init_distributed_environment(...)
911
+ ensure_model_parallel_initialized(...)
912
+ """
913
+ from vllm.config import (
914
+ VllmConfig,
915
+ get_current_vllm_config_or_none,
916
+ set_current_vllm_config,
917
+ )
918
+
919
+ if get_current_vllm_config_or_none() is not None:
920
+ # Config already set, just yield
921
+ yield
922
+ else:
923
+ # No config set, create a default one for the duration
924
+ with set_current_vllm_config(VllmConfig()):
925
+ yield
926
+
927
+
928
def init_test_distributed_environment(
929
tp_size: int,
930
pp_size: int,
@@ -921,6 +951,7 @@ def init_test_distributed_environment(
951
distributed_init_method=distributed_init_method,
952
local_rank=local_rank,
953
)
954
+ ensure_model_parallel_initialized(tp_size, pp_size)
955
else:
956
# No config set, create a default one for the test
957
with set_current_vllm_config(VllmConfig()):
@@ -930,7 +961,7 @@ def init_test_distributed_environment(
961
distributed_init_method=distributed_init_method,
962
local_rank=local_rank,
963
)
933
- ensure_model_parallel_initialized(tp_size, pp_size)
964
+ ensure_model_parallel_initialized(tp_size, pp_size)
965
966
967
def multi_process_parallel(
tests/v1/worker/test_gpu_model_runner.py
+5
-2
@@ -789,8 +789,11 @@ def test_hybrid_attention_mamba_tensor_shapes():
789
"MASTER_PORT": "12345",
790
}
791
)
792
- init_distributed_environment()
793
- initialize_model_parallel(tensor_model_parallel_size=1)
792
+ from tests.utils import ensure_current_vllm_config
793
+
794
+ with ensure_current_vllm_config():
795
+ init_distributed_environment()
796
+ initialize_model_parallel(tensor_model_parallel_size=1)
797
torch.set_default_dtype(torch.float16)
798
799
model_config = ModelConfig(
tests/v1/worker/test_worker_memory_snapshot.py
+7
-1
@@ -10,6 +10,7 @@ from unittest.mock import patch
10
import pytest
11
import torch
12
13
+from vllm.config import set_current_vllm_config
14
from vllm.engine.arg_utils import EngineArgs
15
from vllm.utils.mem_utils import MemorySnapshot
16
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
@@ -95,7 +96,12 @@ def worker_process(
96
side_effect=make_operation_tracker("nccl_all_reduce", original_all_reduce),
97
)
98
98
- with init_patch, memory_patch, all_reduce_patch:
99
+ with (
100
+ init_patch,
101
+ memory_patch,
102
+ all_reduce_patch,
103
+ set_current_vllm_config(vllm_config),
104
+ ):
105
# Initialize device (this is where we test the order)
106
worker.init_device()
107
vllm/compilation/wrapper.py
+49
@@ -319,3 +319,52 @@ class TorchCompileWithNoGuardsWrapper:
319
yield
320
finally:
321
self.__class__.forward.__code__ = original
322
+
323
+
324
+def reset_compile_wrapper(model: torch.nn.Module) -> None:
325
+ """
326
+ Clean up compiled model and captured CUDA graphs for elastic EP.
327
+ """
328
+ if not isinstance(model, TorchCompileWithNoGuardsWrapper) and hasattr(
329
+ model, "model"
330
+ ):
331
+ model = model.model
332
+ if not isinstance(model, TorchCompileWithNoGuardsWrapper):
333
+ return
334
+ # model.do_not_compile is set by the @support_torch_compile decorator
335
+ if hasattr(model, "do_not_compile") and model.do_not_compile:
336
+ return
337
+ from vllm.compilation.counter import compilation_counter
338
+
339
+ # reset the compilation counter
340
+ compilation_counter.num_models_seen = 0
341
+ compilation_counter.num_graphs_seen = 0
342
+ compilation_counter.num_piecewise_graphs_seen = 0
343
+ compilation_counter.num_piecewise_capturable_graphs_seen = 0
344
+ compilation_counter.num_backend_compilations = 0
345
+ compilation_counter.num_gpu_runner_capture_triggers = 0
346
+ compilation_counter.num_cudagraph_captured = 0
347
+ compilation_counter.num_inductor_compiles = 0
348
+ compilation_counter.num_eager_compiles = 0
349
+ compilation_counter.num_cache_entries_updated = 0
350
+ compilation_counter.num_compiled_artifacts_saved = 0
351
+ compilation_counter.stock_torch_compile_count = 0
352
+
353
+ # Clear the AOT compiled function so the model is forced to
354
+ # recompile on the next call. Without this, decorators.py
355
+ # __call__ uses the stale aot_compiled_fn whose torchinductor
356
+ # kernels have old parameters (expert_map size for example)
357
+ # baked in as compile-time constants.
358
+ if hasattr(model, "aot_compiled_fn"):
359
+ model.aot_compiled_fn = None
360
+ if hasattr(model, "was_aot_compile_fn_loaded_from_disk"):
361
+ model.was_aot_compile_fn_loaded_from_disk = False
362
+
363
+ # Reset the cache_dir so VllmBackend recomputes the hash
364
+ # (data_parallel_size changed, so the config hash differs).
365
+ compilation_config = model.vllm_config.compilation_config
366
+ compilation_config.cache_dir = ""
367
+ compilation_config.local_cache_dir = ""
368
+
369
+ model.__class__.forward.__code__ = model.original_code_object()
370
+ TorchCompileWithNoGuardsWrapper.__init__(model)
vllm/config/parallel.py
+116
-6
@@ -165,6 +165,9 @@ class ParallelConfig:
165
disable_custom_all_reduce: bool = False
166
"""Disable the custom all-reduce kernel and fall back to NCCL."""
167
168
+ enable_elastic_ep: bool = False
169
+ """Enable elastic expert parallelism with stateless NCCL groups for DP/EP."""
170
+
171
enable_dbo: bool = False
172
"""Enable dual batch overlap for the model executor."""
173
ubatch_size: int = 0
@@ -244,6 +247,34 @@ class ParallelConfig:
247
Set to be private as it's not intended to be configured by users.
248
"""
249
250
+ _stateless_dp_group_port_list: list[list[int]] = Field(default_factory=list)
251
+ """List of open ports for stateless DP groups when enable_elastic_ep is True.
252
+ Set to be private as it's not intended to be configured by users.
253
+ It is a list of list[int], with each inner list contains a set of 3 ports
254
+ to be used for setting up the stateless CPU/device/TCPStore groups
255
+ in StatelessGroupCoordinator. The number of inner lists is equal to
256
+ the number of DP groups,
257
+ i.e., len(self._stateless_dp_group_port_list) == world_size_across_dp // dp_size,
258
+ and len(self._stateless_dp_group_port_list[i]) == 3 for all i.
259
+ """
260
+
261
+ _stateless_ep_group_port_list: list[list[int]] = Field(default_factory=list)
262
+ """List of open ports for stateless EP groups when enable_elastic_ep is True.
263
+ Set to be private as it's not intended to be configured by users.
264
+ len(self._stateless_ep_group_port_list) == world_size_across_dp // ep_size,
265
+ """
266
+
267
+ _stateless_eplb_group_port_list: list[list[int]] = Field(default_factory=list)
268
+ """List of open ports for stateless EPLB groups when enable_elastic_ep is True.
269
+ Same topology as EP but separate NCCL communicator to avoid deadlocks.
270
+ """
271
+
272
+ _stateless_world_group_port_list: list[list[int]] = Field(default_factory=list)
273
+ """List of open ports for stateless world group when enable_elastic_ep is True.
274
+ Set to be private as it's not intended to be configured by users.
275
+ len(self._stateless_world_group_port_list) == 1,
276
+ """
277
+
278
decode_context_parallel_size: int = 1
279
"""Number of decode context parallel groups, because the world size does
280
not change by dcp, it simply reuse the GPUs of TP group, and tp_size
@@ -402,7 +433,67 @@ class ParallelConfig:
433
434
return answer
435
405
- def stateless_init_dp_group(self) -> ProcessGroup:
436
+ def allocate_elastic_ep_ports(self) -> None:
437
+ """Allocate all ports for elastic EP (stateless groups + DP master).
438
+
439
+ Must be called AFTER ray.init() so that ports claimed by Ray's
440
+ idle worker pool are already in use and won't be returned by
441
+ get_open_ports_list().
442
+ """
443
+ if not self.enable_elastic_ep:
444
+ return
445
+ if self._stateless_world_group_port_list:
446
+ return
447
+
448
+ num_world_groups = 1
449
+ dp_size = self.data_parallel_size
450
+ ep_size = self.data_parallel_size * self.world_size_across_dp
451
+ num_dp_groups = max(1, self.world_size_across_dp // dp_size)
452
+ num_ep_groups = max(1, self.world_size_across_dp // ep_size)
453
+ num_eplb_groups = num_ep_groups
454
+ total_stateless_ports = (
455
+ num_world_groups + num_dp_groups + num_ep_groups + num_eplb_groups
456
+ ) * 3
457
+ num_dp_master_ports = 5
458
+
459
+ all_ports = get_open_ports_list(total_stateless_ports + num_dp_master_ports)
460
+
461
+ self._data_parallel_master_port_list = all_ports[-num_dp_master_ports:]
462
+ self.data_parallel_master_port = self._data_parallel_master_port_list.pop()
463
+ all_ports = all_ports[:-num_dp_master_ports]
464
+
465
+ self._stateless_world_group_port_list = [
466
+ all_ports[i : i + 3] for i in range(0, num_world_groups * 3, 3)
467
+ ]
468
+ start_idx = num_world_groups * 3
469
+ self._stateless_dp_group_port_list = [
470
+ all_ports[i : i + 3]
471
+ for i in range(start_idx, start_idx + num_dp_groups * 3, 3)
472
+ ]
473
+ start_idx += num_dp_groups * 3
474
+ self._stateless_ep_group_port_list = [
475
+ all_ports[i : i + 3]
476
+ for i in range(start_idx, start_idx + num_ep_groups * 3, 3)
477
+ ]
478
+ start_idx += num_ep_groups * 3
479
+ self._stateless_eplb_group_port_list = [
480
+ all_ports[i : i + 3]
481
+ for i in range(start_idx, start_idx + num_eplb_groups * 3, 3)
482
+ ]
483
+
484
+ def get_next_stateless_world_group_port(self) -> list[int]:
485
+ return self._stateless_world_group_port_list.pop()
486
+
487
+ def get_next_stateless_dp_group_port(self) -> list[int]:
488
+ return self._stateless_dp_group_port_list.pop()
489
+
490
+ def get_next_stateless_ep_group_port(self) -> list[int]:
491
+ return self._stateless_ep_group_port_list.pop()
492
+
493
+ def get_next_stateless_eplb_group_port(self) -> list[int]:
494
+ return self._stateless_eplb_group_port_list.pop()
495
+
496
+ def stateless_init_dp_group(self, return_store: bool = False) -> ProcessGroup:
497
# NOTE: In high-concurrency scenarios multiple processes
498
# can pick the same (currently free) port through a race
499
# condition when calling `get_open_port()`. When the first
@@ -426,7 +517,8 @@ class ParallelConfig:
517
self.get_next_dp_init_port(),
518
self.data_parallel_rank,
519
self.data_parallel_size,
429
- backend=current_platform.dist_backend,
520
+ backend="gloo",
521
+ return_store=return_store,
522
)
523
except DistNetworkError as e:
524
# We only want to retry when the root cause is EADDRINUSE.
@@ -561,6 +653,21 @@ class ParallelConfig:
653
logger.info("Using external launcher for distributed inference.")
654
self.world_size *= self.data_parallel_size
655
656
+ if self.enable_elastic_ep:
657
+ if not self.enable_eplb:
658
+ raise ValueError("Elastic EP is only supported with enable_eplb=True.")
659
+ if self.pipeline_parallel_size > 1:
660
+ raise ValueError(
661
+ "Elastic EP is not supported with pipeline parallelism "
662
+ f"(pipeline_parallel_size={self.pipeline_parallel_size})."
663
+ )
664
+ if self.data_parallel_external_lb or self.data_parallel_hybrid_lb:
665
+ raise NotImplementedError(
666
+ "Elastic EP is not compatible with data_parallel_external_lb "
667
+ "or data_parallel_hybrid_lb. Elastic EP relies on a single API "
668
+ "server and core client to coordinate scale up/down."
669
+ )
670
+
671
if self.data_parallel_size > 1 or self.data_parallel_size_local == 0:
672
# Data parallel was specified in the engine args.
673
if self.distributed_executor_backend == "external_launcher":
@@ -573,9 +680,12 @@ class ParallelConfig:
680
"Set data_parallel_rank to %d automatically.",
681
self.data_parallel_rank,
682
)
576
- if not self._data_parallel_master_port_list:
577
- self._data_parallel_master_port_list = get_open_ports_list(5)
578
- self.data_parallel_master_port = self._data_parallel_master_port_list.pop()
683
+ if not self.enable_elastic_ep:
684
+ if not self._data_parallel_master_port_list:
685
+ self._data_parallel_master_port_list = get_open_ports_list(5)
686
+ self.data_parallel_master_port = (
687
+ self._data_parallel_master_port_list.pop()
688
+ )
689
690
if not (0 <= self.data_parallel_rank < self.data_parallel_size):
691
raise ValueError(
@@ -602,7 +712,7 @@ class ParallelConfig:
712
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
713
logger.info("Disabling V1 multiprocessing for external launcher.")
714
605
- if self.distributed_executor_backend is None and self.world_size > 1:
715
+ if self.distributed_executor_backend is None and self.world_size_across_dp > 1:
716
# We use multiprocessing by default if world_size fits on the
717
# current node and we aren't in a ray placement group.
718
vllm/distributed/device_communicators/all2all.py
+18
-13
@@ -31,8 +31,8 @@ class NaiveAll2AllManager(All2AllManagerBase):
31
debugging.
32
"""
33
34
- def __init__(self, cpu_group):
35
- super().__init__(cpu_group)
34
+ def __init__(self, cpu_group, tcp_store_group=None):
35
+ super().__init__(cpu_group, tcp_store_group)
36
37
def naive_multicast(
38
self,
@@ -138,8 +138,8 @@ class AgRsAll2AllManager(All2AllManagerBase):
138
all-gather (dispatch) and reduce-scatter (combine).
139
"""
140
141
- def __init__(self, cpu_group):
142
- super().__init__(cpu_group)
141
+ def __init__(self, cpu_group, tcp_store_group=None):
142
+ super().__init__(cpu_group, tcp_store_group)
143
144
def dispatch_router_logits(
145
self,
@@ -239,12 +239,12 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):
239
All2All communication based on DeepEP High-Throughput kernels.
240
"""
241
242
- def __init__(self, cpu_group):
242
+ def __init__(self, cpu_group, tcp_store_group=None):
243
assert has_deep_ep(), (
244
"DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md"
245
" to install DeepEP kernels."
246
) # noqa
247
- super().__init__(cpu_group)
247
+ super().__init__(cpu_group, tcp_store_group)
248
self.handle_cache = Cache()
249
250
# This is the DeepEP default. Stick to it till we can establish
@@ -282,7 +282,10 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):
282
raise NotImplementedError
283
284
def destroy(self):
285
- pass
285
+ with self.handle_cache._lock:
286
+ for _, handle in self.handle_cache._cache.items():
287
+ handle.destroy()
288
+ self.handle_cache._cache.clear()
289
290
291
class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
@@ -290,8 +293,8 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
293
All2All communication based on DeepEP High-Throughput kernels.
294
"""
295
293
- def __init__(self, cpu_group):
294
- super().__init__(cpu_group)
296
+ def __init__(self, cpu_group, tcp_store_group=None):
297
+ super().__init__(cpu_group, tcp_store_group)
298
299
def _make_all2all_kwargs(self) -> dict[Any, Any]:
300
# Defaults for internode and intranode are taken from DeepEP tests.
@@ -314,6 +317,7 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
317
num_rdma_bytes=num_rdma_bytes,
318
low_latency_mode=False,
319
num_qps_per_rank=num_qps_per_rank,
320
+ explicitly_destroy=True,
321
)
322
323
def get_handle(self, kwargs):
@@ -347,8 +351,8 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
351
All2All communication based on DeepEP Low-Latency kernels.
352
"""
353
350
- def __init__(self, cpu_group):
351
- super().__init__(cpu_group)
354
+ def __init__(self, cpu_group, tcp_store_group=None):
355
+ super().__init__(cpu_group, tcp_store_group)
356
357
def _make_all2all_kwargs(
358
self,
@@ -387,6 +391,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
391
num_qps_per_rank=num_qps_per_rank,
392
allow_nvlink_for_low_latency_mode=True,
393
allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
394
+ explicitly_destroy=True,
395
)
396
397
def get_handle(self, kwargs):
@@ -418,11 +423,11 @@ class FlashInferAllToAllManager(All2AllManagerBase):
423
rank: int
424
world_size: int
425
421
- def __init__(self, cpu_group):
426
+ def __init__(self, cpu_group, tcp_store_group=None):
427
assert has_flashinfer_all2all(), (
428
"flashinfer all2all module not found. Please install/check flashinfer"
429
) # noqa
425
- super().__init__(cpu_group)
430
+ super().__init__(cpu_group, tcp_store_group)
431
logger.debug(
432
"Initialize for flashinfer All2All rank=%d, world size=%d",
433
self.rank,
vllm/distributed/device_communicators/base_device_communicator.py
+46
-11
@@ -29,8 +29,9 @@ class All2AllManagerBase:
29
rank: int
30
world_size: int
31
32
- def __init__(self, cpu_group):
32
+ def __init__(self, cpu_group, tcp_store_group=None):
33
self.cpu_group = cpu_group
34
+ self.tcp_store_group = tcp_store_group
35
36
# compute some common properties
37
from vllm.distributed.parallel_state import (
@@ -47,12 +48,17 @@ class All2AllManagerBase:
48
# when we create this object
49
self.dp_rank = self.dp_group.rank_in_group
50
self.dp_world_size = self.dp_group.world_size
50
- self.rank = dist.get_rank(cpu_group)
51
- self.world_size = dist.get_world_size(cpu_group)
51
+ self.rank = cpu_group.rank()
52
+ self.world_size = cpu_group.size()
53
54
# all2all communication often has separate implementations for
55
# intra-node and inter-node communication
55
- self.internode = not all(in_the_same_node_as(cpu_group, source_rank=0))
56
+ if tcp_store_group is None:
57
+ self.internode = not all(in_the_same_node_as(cpu_group, source_rank=0))
58
+ else:
59
+ self.internode = not all(
60
+ in_the_same_node_as(tcp_store_group, source_rank=0)
61
+ )
62
63
def get_handle(self, kwargs):
64
# get a handle for the all2all communication,
@@ -121,17 +127,36 @@ class DeviceCommunicatorBase:
127
device: torch.device | None = None,
128
device_group: ProcessGroup | None = None,
129
unique_name: str = "",
130
+ global_ranks: list[int] | None = None,
131
+ global_world_size: int | None = None,
132
):
133
self.device = device or torch.device("cpu")
134
self.cpu_group = cpu_group
135
self.device_group = device_group
136
self.unique_name = unique_name
129
- self.rank = dist.get_rank(cpu_group)
130
- self.world_size = dist.get_world_size(cpu_group)
131
- self.ranks = dist.get_process_group_ranks(cpu_group)
132
- self.global_rank = dist.get_rank()
133
- self.global_world_size = dist.get_world_size()
134
- self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
137
+
138
+ # Check if this is a stateless process group
139
+ from torch.distributed.distributed_c10d import _world
140
+
141
+ is_stateless = _world.pg_map.get(cpu_group, None) is None
142
+
143
+ if is_stateless:
144
+ # For stateless groups, we can't use torch.distributed methods
145
+ self.rank = cpu_group.rank()
146
+ self.world_size = cpu_group.size()
147
+ assert global_ranks is not None
148
+ assert global_world_size is not None
149
+ self.ranks = global_ranks
150
+ self.global_rank = self.ranks[self.rank]
151
+ self.global_world_size = global_world_size
152
+ self.rank_in_group = self.rank
153
+ else:
154
+ self.rank = dist.get_rank(cpu_group)
155
+ self.world_size = dist.get_world_size(cpu_group)
156
+ self.ranks = dist.get_process_group_ranks(cpu_group)
157
+ self.global_rank = dist.get_rank()
158
+ self.global_world_size = dist.get_world_size()
159
+ self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
160
161
use_ep = False
162
all2all_backend = None
@@ -145,7 +170,7 @@ class DeviceCommunicatorBase:
170
use_ep = config.parallel_config.data_parallel_size > 1
171
all2all_backend = config.parallel_config.all2all_backend
172
148
- self.is_ep_communicator = "ep" in unique_name
173
+ self.is_ep_communicator = unique_name.split(":")[0] == "ep"
174
self.use_all2all = self.is_ep_communicator and use_ep
175
self.all2all_backend = all2all_backend
176
self.all2all_manager: All2AllManagerBase | None = None
@@ -275,6 +300,13 @@ class DeviceCommunicatorBase:
300
torch.distributed.recv(tensor, self.ranks[src], self.device_group)
301
return tensor
302
303
+ def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
304
+ """Broadcast a tensor from source rank to all ranks."""
305
+ if self.world_size == 1:
306
+ return tensor
307
+ torch.distributed.broadcast(tensor, self.ranks[src], self.device_group)
308
+ return tensor
309
+
310
def destroy(self):
311
pass
312
@@ -343,3 +375,6 @@ class DeviceCommunicatorBase:
375
This is a no-op in the base class.
376
"""
377
return hidden_states
378
+
379
+ def batch_isend_irecv(self, p2p_ops: list):
380
+ raise NotImplementedError
vllm/distributed/device_communicators/cuda_communicator.py
+47
-7
@@ -16,6 +16,7 @@ from vllm.distributed.device_communicators.pynccl_allocator import (
16
from vllm.logger import init_logger
17
from vllm.platforms import current_platform
18
19
+from ..utils import StatelessProcessGroup
20
from .base_device_communicator import DeviceCommunicatorBase
21
22
logger = init_logger(__name__)
@@ -28,8 +29,18 @@ class CudaCommunicator(DeviceCommunicatorBase):
29
device: torch.device | None = None,
30
device_group: ProcessGroup | None = None,
31
unique_name: str = "",
32
+ global_ranks: list[int] | None = None,
33
+ global_world_size: int | None = None,
34
+ tcp_store_group: StatelessProcessGroup | None = None,
35
):
32
- super().__init__(cpu_group, device, device_group, unique_name)
36
+ super().__init__(
37
+ cpu_group,
38
+ device,
39
+ device_group,
40
+ unique_name,
41
+ global_ranks,
42
+ global_world_size,
43
+ )
44
if "tp" not in unique_name:
45
# custom allreduce or torch symm mem can be used only by tp
46
use_custom_allreduce = False
@@ -62,7 +73,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
73
self.pynccl_comm: PyNcclCommunicator | None = None
74
if self.world_size > 1:
75
self.pynccl_comm = PyNcclCommunicator(
65
- group=self.cpu_group,
76
+ group=self.cpu_group if tcp_store_group is None else tcp_store_group,
77
device=self.device,
78
)
79
if is_symmetric_memory_enabled():
@@ -107,19 +118,27 @@ class CudaCommunicator(DeviceCommunicatorBase):
118
if self.all2all_backend == "naive":
119
from .all2all import NaiveAll2AllManager
120
110
- self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
121
+ self.all2all_manager = NaiveAll2AllManager(
122
+ self.cpu_group, tcp_store_group
123
+ )
124
elif self.all2all_backend == "allgather_reducescatter":
125
from .all2all import AgRsAll2AllManager
126
114
- self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
127
+ self.all2all_manager = AgRsAll2AllManager(
128
+ self.cpu_group, tcp_store_group
129
+ )
130
elif self.all2all_backend == "deepep_high_throughput":
131
from .all2all import DeepEPHTAll2AllManager
132
118
- self.all2all_manager = DeepEPHTAll2AllManager(self.cpu_group)
133
+ self.all2all_manager = DeepEPHTAll2AllManager(
134
+ self.cpu_group, tcp_store_group
135
+ )
136
elif self.all2all_backend == "deepep_low_latency":
137
from .all2all import DeepEPLLAll2AllManager
138
122
- self.all2all_manager = DeepEPLLAll2AllManager(self.cpu_group)
139
+ self.all2all_manager = DeepEPLLAll2AllManager(
140
+ self.cpu_group, tcp_store_group
141
+ )
142
elif self.all2all_backend == "mori":
143
from .all2all import MoriAll2AllManager
144
@@ -127,7 +146,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
146
elif self.all2all_backend == "flashinfer_all2allv":
147
from .all2all import FlashInferAllToAllManager
148
130
- self.all2all_manager = FlashInferAllToAllManager(self.cpu_group)
149
+ self.all2all_manager = FlashInferAllToAllManager(
150
+ self.cpu_group, tcp_store_group
151
+ )
152
else:
153
raise ValueError(f"Unknown all2all backend: {self.all2all_backend}")
154
@@ -284,6 +305,18 @@ class CudaCommunicator(DeviceCommunicatorBase):
305
torch.distributed.recv(tensor, self.ranks[src], self.device_group)
306
return tensor
307
308
+ def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
309
+ """Broadcast a tensor from source rank to all ranks."""
310
+ if self.world_size == 1:
311
+ return tensor
312
+
313
+ pynccl_comm = self.pynccl_comm
314
+ if pynccl_comm is not None and not pynccl_comm.disabled:
315
+ pynccl_comm.broadcast(tensor, src)
316
+ return tensor
317
+ else:
318
+ raise ValueError("No PyNCCL communicator found")
319
+
320
def destroy(self):
321
if self.pynccl_comm is not None:
322
self.pynccl_comm = None
@@ -403,3 +436,10 @@ class CudaCommunicator(DeviceCommunicatorBase):
436
hidden_states,
437
is_sequence_parallel,
438
)
439
+
440
+ def batch_isend_irecv(self, p2p_ops: list):
441
+ pynccl_comm = self.pynccl_comm
442
+ if pynccl_comm is not None and not pynccl_comm.disabled:
443
+ pynccl_comm.batch_isend_irecv(p2p_ops)
444
+ else:
445
+ raise ValueError("No PyNCCL communicator found")
vllm/distributed/device_communicators/pynccl.py
+34
-2
@@ -312,10 +312,19 @@ class PyNcclCommunicator:
312
)
313
if stream is None:
314
stream = current_stream()
315
+ if tensor.dtype in [
316
+ torch.float8_e5m2,
317
+ torch.float8_e4m3fn,
318
+ torch.float8_e4m3fnuz,
319
+ torch.float8_e5m2fnuz,
320
+ ]:
321
+ nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8)
322
+ else:
323
+ nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype)
324
self.nccl.ncclSend(
325
buffer_type(tensor.data_ptr()),
326
tensor.numel(),
318
- ncclDataTypeEnum.from_torch(tensor.dtype),
327
+ nccl_dtype,
328
dst,
329
self.comm,
330
cudaStream_t(stream.cuda_stream),
@@ -330,10 +339,19 @@ class PyNcclCommunicator:
339
)
340
if stream is None:
341
stream = current_stream()
342
+ if tensor.dtype in [
343
+ torch.float8_e5m2,
344
+ torch.float8_e4m3fn,
345
+ torch.float8_e4m3fnuz,
346
+ torch.float8_e5m2fnuz,
347
+ ]:
348
+ nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8)
349
+ else:
350
+ nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype)
351
self.nccl.ncclRecv(
352
buffer_type(tensor.data_ptr()),
353
tensor.numel(),
336
- ncclDataTypeEnum.from_torch(tensor.dtype),
354
+ nccl_dtype,
355
src,
356
self.comm,
357
cudaStream_t(stream.cuda_stream),
@@ -384,3 +402,17 @@ class PyNcclCommunicator:
402
403
def deregister_comm_window(self, window):
404
return self.nccl.ncclCommWindowDeregister(self.comm, window)
405
+
406
+ def batch_isend_irecv(self, p2p_ops: list, stream=None):
407
+ if self.disabled:
408
+ return
409
+ if stream is None:
410
+ stream = current_stream()
411
+ self.group_start()
412
+ for op in p2p_ops:
413
+ if op.op is torch.distributed.isend:
414
+ self.send(op.tensor, op.group_peer, stream)
415
+ elif op.op is torch.distributed.irecv:
416
+ self.recv(op.tensor, op.group_peer, stream)
417
+
418
+ self.group_end()
vllm/distributed/elastic_ep/__init__.py
vllm/distributed/elastic_ep/elastic_execute.py
new
+529
@@ -0,0 +1,529 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
+import gc
5
+import weakref
6
+from collections.abc import Iterable, Sequence
7
+
8
+import torch
9
+import torch.nn as nn
10
+import torch.nn.functional as F
11
+from torch.distributed import P2POp
12
+
13
+from vllm.compilation.counter import compilation_counter
14
+from vllm.compilation.cuda_graph import CUDAGraphWrapper
15
+from vllm.compilation.wrapper import reset_compile_wrapper
16
+from vllm.config import (
17
+ CompilationMode,
18
+ set_current_vllm_config,
19
+)
20
+from vllm.distributed import (
21
+ get_dp_group,
22
+ get_ep_group,
23
+ get_pcp_group,
24
+ get_tp_group,
25
+)
26
+from vllm.distributed.elastic_ep.standby_state import (
27
+ create_standby_groups,
28
+ get_standby_dp_group,
29
+ get_standby_ep_group,
30
+ pop_standby_groups,
31
+)
32
+from vllm.distributed.parallel_state import (
33
+ _replace_active_groups,
34
+ prepare_communication_buffer_for_model,
35
+)
36
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
37
+from vllm.logger import init_logger
38
+from vllm.model_executor.layers.fused_moe.layer import FusedMoEParallelConfig
39
+from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
40
+from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
41
+from vllm.v1.worker.workspace import lock_workspace, unlock_workspace
42
+
43
+logger = init_logger(__name__)
44
+
45
+
46
+def batch_transfer_weights(
47
+ model: nn.Module,
48
+ is_sender: bool,
49
+ peer_rank: int,
50
+ dp_group: StatelessGroupCoordinator,
51
+ expert_weights: Sequence[Iterable[torch.Tensor]],
52
+) -> None:
53
+ device_comm = dp_group.device_communicator
54
+ if device_comm is None:
55
+ raise ValueError("No device communicator found")
56
+
57
+ expert_weights_set = set()
58
+ for weight_group in expert_weights:
59
+ for weight in weight_group:
60
+ expert_weights_set.add(weight.data_ptr())
61
+
62
+ state_dict = model.state_dict()
63
+ all_params = []
64
+
65
+ for name, param in state_dict.items():
66
+ if name.endswith("expert_map"):
67
+ continue
68
+ if param.data_ptr() not in expert_weights_set:
69
+ all_params.append(param.data)
70
+
71
+ assert len(all_params) > 0
72
+ p2p_ops = []
73
+ for param in all_params:
74
+ op = object.__new__(P2POp)
75
+ if is_sender:
76
+ op.op = torch.distributed.isend
77
+ op.tensor = param
78
+ else:
79
+ op.op = torch.distributed.irecv
80
+ op.tensor = param
81
+ op.group_peer = peer_rank
82
+ p2p_ops.append(op)
83
+ device_comm.batch_isend_irecv(p2p_ops)
84
+
85
+
86
+def broadcast_expert_mapping(
87
+ physical_to_logical: torch.Tensor | None,
88
+ num_local_physical_experts: int | None,
89
+ num_logical_experts: int | None,
90
+ dp_group: StatelessGroupCoordinator,
91
+ device: torch.device,
92
+ src_rank: int = 0,
93
+) -> tuple[torch.Tensor, int, int]:
94
+ if dp_group.rank_in_group == src_rank:
95
+ assert physical_to_logical is not None
96
+ assert num_local_physical_experts is not None
97
+ assert num_logical_experts is not None
98
+ assert physical_to_logical.dtype == torch.int64
99
+ shape_tensor = torch.tensor(
100
+ list(physical_to_logical.shape), dtype=torch.int64, device="cpu"
101
+ )
102
+ metadata_tensor = torch.tensor(
103
+ [num_local_physical_experts, num_logical_experts],
104
+ dtype=torch.int64,
105
+ device="cpu",
106
+ )
107
+ else:
108
+ shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
109
+ metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
110
+
111
+ shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank)
112
+ metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_tensor, src_rank)
113
+
114
+ if dp_group.rank_in_group != src_rank:
115
+ assert device is not None
116
+ physical_to_logical = torch.empty(
117
+ tuple(shape_tensor.tolist()),
118
+ dtype=torch.int64,
119
+ device=device,
120
+ )
121
+
122
+ assert physical_to_logical is not None
123
+ physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank)
124
+ num_local_physical_experts = int(metadata_tensor[0].item())
125
+ num_logical_experts = int(metadata_tensor[1].item())
126
+
127
+ return physical_to_logical, num_local_physical_experts, num_logical_experts
128
+
129
+
130
+class ElasticEPScalingExecutor:
131
+ def __init__(self, worker):
132
+ self.worker_ref = weakref.ref(worker)
133
+ self.reconfig_request = None
134
+
135
+ @property
136
+ def worker(self):
137
+ worker = self.worker_ref()
138
+ if worker is None:
139
+ raise RuntimeError("Worker has been garbage collected")
140
+ return worker
141
+
142
+ def execute(self, execute_method: str, *args, **kwargs):
143
+ method = getattr(self, execute_method, None)
144
+ if method is None:
145
+ raise ValueError(f"Unknown execute method: {execute_method}")
146
+ return method(*args, **kwargs)
147
+
148
+ def create_standby_groups(
149
+ self, reconfig_request: ReconfigureDistributedRequest
150
+ ) -> None:
151
+ self.reconfig_request = reconfig_request
152
+ new_dp_size = reconfig_request.new_data_parallel_size
153
+ world_size = self.worker.vllm_config.parallel_config.world_size
154
+ new_world_size_across_dp = world_size * new_dp_size
155
+ updated_config = copy.copy(self.worker.vllm_config)
156
+ updated_config.parallel_config = copy.deepcopy(
157
+ self.worker.vllm_config.parallel_config
158
+ )
159
+ updated_config.parallel_config.data_parallel_size = new_dp_size
160
+ with set_current_vllm_config(updated_config):
161
+ create_standby_groups(
162
+ new_dp_size=new_dp_size,
163
+ new_world_size_across_dp=new_world_size_across_dp,
164
+ master_ip=reconfig_request.new_data_parallel_master_ip,
165
+ world_group_ports=reconfig_request.new_stateless_world_group_port_list,
166
+ dp_group_ports=reconfig_request.new_stateless_dp_group_port_list,
167
+ ep_group_ports=reconfig_request.new_stateless_ep_group_port_list,
168
+ eplb_group_ports=reconfig_request.new_stateless_eplb_group_port_list,
169
+ )
170
+ self.worker.model_runner.eep_eplb_suppressed = True
171
+ standby_ep_group = get_standby_ep_group()
172
+ assert standby_ep_group is not None
173
+ if standby_ep_group.rank == 0:
174
+ logger.info("[Elastic EP] EPLB disabled during elastic scaling transition")
175
+
176
+ def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None:
177
+ standby_dp_group = get_standby_dp_group()
178
+ assert standby_dp_group is not None
179
+ # Broadcast old_dp_size to all workers in standby group
180
+ if standby_dp_group.rank_in_group < old_dp_size:
181
+ old_dp_size_tensor = torch.tensor(
182
+ [old_dp_size], dtype=torch.int64, device="cpu"
183
+ )
184
+ else:
185
+ old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu")
186
+ old_dp_size_tensor = standby_dp_group.tcp_store_group.broadcast(
187
+ old_dp_size_tensor, 0
188
+ )
189
+
190
+ num_new_workers = new_dp_size - old_dp_size
191
+ dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
192
+
193
+ # Sender-receiver pairing: the first new_workers % old_dp_size
194
+ # senders get (k+1) contiguous receivers, the rest get k
195
+ # receivers.
196
+ num_dst_per_sender = num_new_workers // old_dp_size
197
+ remainder = num_new_workers % old_dp_size
198
+
199
+ if dp_rank < remainder:
200
+ recv_begin = dp_rank * (num_dst_per_sender + 1)
201
+ recv_end = recv_begin + num_dst_per_sender + 1
202
+ else:
203
+ recv_begin = (
204
+ remainder * (num_dst_per_sender + 1)
205
+ + (dp_rank - remainder) * num_dst_per_sender
206
+ )
207
+ recv_end = recv_begin + num_dst_per_sender
208
+
209
+ ranks_to_send = list(range(old_dp_size + recv_begin, old_dp_size + recv_end))
210
+
211
+ model = self.worker.model_runner.get_model()
212
+ for new_worker_rank in sorted(ranks_to_send):
213
+ batch_transfer_weights(
214
+ model=model,
215
+ is_sender=True,
216
+ peer_rank=new_worker_rank,
217
+ dp_group=standby_dp_group,
218
+ expert_weights=model.expert_weights,
219
+ )
220
+ torch.cuda.synchronize()
221
+
222
+ def broadcast_expert_mapping(self) -> None:
223
+ standby_dp_group = get_standby_dp_group()
224
+ assert standby_dp_group is not None
225
+ model_config = self.worker.model_runner.model_config
226
+ eplb_state = self.worker.model_runner.eplb_state
227
+ assert eplb_state is not None
228
+ eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
229
+ physical_to_logical = eplb_model_state.physical_to_logical_map
230
+ num_physical_experts = physical_to_logical.shape[1]
231
+ num_local_physical_experts = num_physical_experts // get_ep_group().world_size
232
+ num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
233
+ broadcast_expert_mapping(
234
+ physical_to_logical=physical_to_logical,
235
+ num_local_physical_experts=num_local_physical_experts,
236
+ num_logical_experts=num_logical_experts,
237
+ dp_group=standby_dp_group,
238
+ src_rank=0,
239
+ device=self.worker.device,
240
+ )
241
+
242
+ def switch_and_remove(self) -> None:
243
+ _replace_active_groups(world=None, dp=None, ep=None, eplb=None, node_count=None)
244
+
245
+ def switch_and_prepare(self) -> None:
246
+ old_dp_size = get_dp_group().world_size
247
+ old_ep_size = get_ep_group().world_size
248
+
249
+ _replace_active_groups(**pop_standby_groups())
250
+
251
+ parallel_config = self.worker.vllm_config.parallel_config
252
+ reconfig_request = self.reconfig_request
253
+ assert reconfig_request is not None
254
+ new_dp_size = reconfig_request.new_data_parallel_size
255
+ new_ep_size = get_ep_group().world_size
256
+
257
+ parallel_config.data_parallel_size = new_dp_size
258
+ if (
259
+ reconfig_request.new_data_parallel_rank
260
+ != ReconfigureRankType.KEEP_CURRENT_RANK
261
+ ):
262
+ parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
263
+ if (
264
+ reconfig_request.new_data_parallel_rank_local
265
+ != ReconfigureRankType.KEEP_CURRENT_RANK
266
+ ):
267
+ parallel_config.data_parallel_rank_local = (
268
+ reconfig_request.new_data_parallel_rank_local
269
+ )
270
+ parallel_config.data_parallel_master_ip = (
271
+ reconfig_request.new_data_parallel_master_ip
272
+ )
273
+ parallel_config.data_parallel_master_port = (
274
+ reconfig_request.new_data_parallel_master_port
275
+ )
276
+
277
+ # Reconfigure MoE modules with new EP size
278
+ moe_modules = [
279
+ module
280
+ for module in self.worker.model_runner.model.modules()
281
+ if (
282
+ module.__class__.__name__ == "FusedMoE"
283
+ or module.__class__.__name__ == "SharedFusedMoE"
284
+ )
285
+ ]
286
+ num_local_experts = moe_modules[0].moe_config.num_local_experts
287
+ assert all(
288
+ module.moe_config.num_local_experts == num_local_experts
289
+ for module in moe_modules
290
+ ), "All MoE modules must have the same number of experts"
291
+ for module in moe_modules:
292
+ module.moe_config.num_experts = num_local_experts * new_ep_size
293
+ module.global_num_experts = module.moe_config.num_experts
294
+ tp_size = get_tp_group().world_size
295
+ is_sequence_parallel = parallel_config.use_sequence_parallel_moe
296
+ sp_size = tp_size if is_sequence_parallel else 1
297
+ module.moe_parallel_config = FusedMoEParallelConfig.make(
298
+ tp_size_=tp_size,
299
+ pcp_size_=get_pcp_group().world_size,
300
+ dp_size_=get_dp_group().world_size,
301
+ sp_size_=sp_size,
302
+ vllm_parallel_config=parallel_config,
303
+ )
304
+ module.moe_config.moe_parallel_config = module.moe_parallel_config
305
+
306
+ # Update EPLB state
307
+ eplb_state = self.worker.model_runner.eplb_state
308
+ assert eplb_state is not None
309
+ model_config = self.worker.model_runner.model_config
310
+ eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
311
+
312
+ num_physical_experts = num_local_experts * new_ep_size
313
+ num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
314
+ parallel_config.eplb_config.num_redundant_experts = (
315
+ num_physical_experts - num_logical_experts
316
+ )
317
+ old_physical_to_logical = eplb_model_state.physical_to_logical_map
318
+ num_moe_layers = old_physical_to_logical.shape[0]
319
+ num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size
320
+ if new_dp_size > old_dp_size:
321
+ expanded_physical_to_logical = torch.full(
322
+ (num_moe_layers, num_local_experts * new_ep_size),
323
+ -1,
324
+ dtype=old_physical_to_logical.dtype,
325
+ device=old_physical_to_logical.device,
326
+ )
327
+ expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = (
328
+ old_physical_to_logical
329
+ )
330
+ eplb_model_state.physical_to_logical_map = expanded_physical_to_logical
331
+
332
+ old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1]
333
+ pad_size = num_physical_experts - old_num_physical_experts
334
+ if new_dp_size > old_dp_size:
335
+ assert pad_size > 0
336
+ expanded_expert_load_pass = F.pad(
337
+ eplb_model_state.expert_load_pass, (0, pad_size), value=0
338
+ )
339
+ expanded_expert_load_window = F.pad(
340
+ eplb_model_state.expert_load_window, (0, pad_size), value=0
341
+ )
342
+ eplb_model_state.expert_load_pass = expanded_expert_load_pass
343
+ eplb_model_state.expert_load_window = expanded_expert_load_window
344
+ eplb_state.num_valid_physical_experts = old_num_physical_experts
345
+ else:
346
+ assert pad_size < 0
347
+ eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[
348
+ :, :num_physical_experts
349
+ ]
350
+ eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[
351
+ :, :, :num_physical_experts
352
+ ]
353
+ eplb_state.num_valid_physical_experts = num_physical_experts
354
+
355
+ model = self.worker.model_runner.get_model()
356
+ model.expert_weights = []
357
+ with set_current_vllm_config(self.worker.vllm_config):
358
+ model.set_eplb_state(
359
+ eplb_model_state.expert_load_pass,
360
+ eplb_model_state.logical_to_physical_map,
361
+ eplb_model_state.logical_replica_count,
362
+ )
363
+ model.update_physical_experts_metadata(
364
+ num_physical_experts=num_physical_experts,
365
+ num_local_physical_experts=num_local_experts,
366
+ )
367
+ # Force re-creation of the modular kernel (and all2all manager)
368
+ # for the new EP size by resetting quant_method to base
369
+ for module in moe_modules:
370
+ if hasattr(module.quant_method, "old_quant_method"):
371
+ module.quant_method = module.quant_method.old_quant_method
372
+ module.runner = module._init_runner()
373
+ prepare_communication_buffer_for_model(self.worker.model_runner.model)
374
+ if (
375
+ self.worker.vllm_config.compilation_config.mode
376
+ == CompilationMode.STOCK_TORCH_COMPILE
377
+ ):
378
+ # NOTE(yongji): when using stock torch.compile,
379
+ # torch.compile is triggered during GPUModelRunner's load_model()
380
+ # TODO(yongji):check do we need to re-trigger torch.compile here?
381
+ # any changes to the tensor shapes in execution should already
382
+ # be handled internally by torch.compile.
383
+ backend = self.worker.vllm_config.compilation_config.init_backend(
384
+ self.worker.vllm_config
385
+ )
386
+ compilation_counter.stock_torch_compile_count += 1
387
+ self.worker.model_runner.model.compile(fullgraph=True, backend=backend)
388
+
389
+ # release all previously captured CUDA graphs
390
+ if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
391
+ wrapper = self.worker.model_runner.model
392
+ wrapper.concrete_cudagraph_entries = {}
393
+ elif isinstance(self.worker.model_runner.model, UBatchWrapper):
394
+ raise RuntimeError("DBO is not yet supported in elastic EP")
395
+
396
+ multi_block_table = self.worker.model_runner.input_batch.block_table
397
+ saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
398
+ for bt in multi_block_table.block_tables:
399
+ saved_block_tables.append(
400
+ (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
401
+ )
402
+ multi_block_table.clear()
403
+
404
+ # reset the compile wrapper
405
+ torch.compiler.reset()
406
+ with set_current_vllm_config(self.worker.vllm_config):
407
+ reset_compile_wrapper(self.worker.model_runner.get_model())
408
+
409
+ gc.collect()
410
+ torch.cuda.synchronize()
411
+ torch.cuda.empty_cache()
412
+ unlock_workspace()
413
+ self.worker.compile_or_warm_up_model()
414
+ lock_workspace()
415
+
416
+ for bt, (saved_gpu, saved_cpu) in zip(
417
+ multi_block_table.block_tables, saved_block_tables
418
+ ):
419
+ bt.block_table.gpu.copy_(saved_gpu)
420
+ bt.block_table.cpu.copy_(saved_cpu)
421
+
422
+ def perform_eplb_reshuffle(self, new_dp_size: int | None = None) -> None:
423
+ if get_ep_group().rank == 0:
424
+ logger.info("[Elastic EP] Starting expert resharding...")
425
+
426
+ eplb_state = self.worker.model_runner.eplb_state
427
+ assert eplb_state is not None
428
+
429
+ model_config = self.worker.model_runner.model_config
430
+ eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
431
+ is_async_enabled = eplb_state.is_async
432
+ eplb_state.is_async = False
433
+ if new_dp_size is None:
434
+ eplb_state.rearrange()
435
+ else:
436
+ # scale down
437
+ parallel_config = self.worker.vllm_config.parallel_config
438
+ tp_size = parallel_config.tensor_parallel_size
439
+ old_ep_size = parallel_config.data_parallel_size * tp_size
440
+ new_ep_size = new_dp_size * tp_size
441
+
442
+ rank_mapping = {
443
+ old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
444
+ for old_ep_rank in range(old_ep_size)
445
+ }
446
+
447
+ eplb_state.rearrange(rank_mapping=rank_mapping)
448
+ # NOTE(yongji): check whether we need to synchronize here
449
+ torch.cuda.synchronize()
450
+ # reset expert_rearrangement_step to ensure all ranks are synchronized
451
+ eplb_state.expert_rearrangement_step = 0
452
+ eplb_state.num_valid_physical_experts = (
453
+ eplb_model_state.physical_to_logical_map.shape[1]
454
+ )
455
+ eplb_state.is_async = is_async_enabled
456
+ self.worker.model_runner.eep_eplb_suppressed = False
457
+ if get_ep_group().rank == 0:
458
+ logger.info("[Elastic EP] Expert resharding completed")
459
+
460
+ def receive_weights(self) -> None:
461
+ dp_group = get_dp_group()
462
+ assert isinstance(dp_group, StatelessGroupCoordinator)
463
+ new_dp_size = dp_group.world_size
464
+ dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
465
+
466
+ # Receive old_dp_size broadcasted during transfer_weights
467
+ old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu")
468
+ old_dp_size_tensor = dp_group.tcp_store_group.broadcast(old_dp_size_tensor, 0)
469
+ old_dp_size = int(old_dp_size_tensor[0].item())
470
+
471
+ # Calculate which existing worker will send to this new worker
472
+ num_new_workers = new_dp_size - old_dp_size
473
+ new_worker_idx = dp_rank - old_dp_size
474
+ num_dst_per_sender = num_new_workers // old_dp_size
475
+ remainder = num_new_workers % old_dp_size
476
+
477
+ if new_worker_idx < remainder * (num_dst_per_sender + 1):
478
+ sender_rank = new_worker_idx // (num_dst_per_sender + 1)
479
+ else:
480
+ sender_rank = (
481
+ remainder
482
+ + (new_worker_idx - remainder * (num_dst_per_sender + 1))
483
+ // num_dst_per_sender
484
+ )
485
+
486
+ model = self.worker.model_runner.get_model()
487
+ batch_transfer_weights(
488
+ model=model,
489
+ is_sender=False,
490
+ peer_rank=sender_rank,
491
+ dp_group=dp_group,
492
+ expert_weights=model.expert_weights,
493
+ )
494
+ torch.cuda.synchronize()
495
+
496
+ def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]:
497
+ dp_group = get_dp_group()
498
+ assert isinstance(dp_group, StatelessGroupCoordinator)
499
+ physical_to_logical, num_local_physical_experts, num_logical_experts = (
500
+ broadcast_expert_mapping(
501
+ physical_to_logical=None,
502
+ num_local_physical_experts=None,
503
+ num_logical_experts=None,
504
+ dp_group=dp_group,
505
+ src_rank=0,
506
+ device=self.worker.device,
507
+ )
508
+ )
509
+ num_moe_layers = physical_to_logical.shape[0]
510
+ new_dp_size = get_dp_group().world_size
511
+ tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size
512
+ new_ep_size = new_dp_size * tp_size
513
+ expanded_physical_to_logical = torch.full(
514
+ (num_moe_layers, num_local_physical_experts * new_ep_size),
515
+ -1,
516
+ dtype=physical_to_logical.dtype,
517
+ device=physical_to_logical.device,
518
+ )
519
+ old_num_physical_experts = physical_to_logical.shape[1]
520
+ expanded_physical_to_logical[:, :old_num_physical_experts] = physical_to_logical
521
+ return (
522
+ expanded_physical_to_logical,
523
+ num_logical_experts,
524
+ old_num_physical_experts,
525
+ )
526
+
527
+ def prepare_new_worker(self) -> None:
528
+ with set_current_vllm_config(self.worker.vllm_config):
529
+ prepare_communication_buffer_for_model(self.worker.model_runner.get_model())
vllm/distributed/elastic_ep/elastic_state.py
new
+563
@@ -0,0 +1,563 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import enum
4
+import time
5
+import weakref
6
+from datetime import timedelta
7
+from typing import TYPE_CHECKING, Literal
8
+
9
+import torch.distributed
10
+
11
+from vllm.config import ParallelConfig
12
+from vllm.distributed import (
13
+ sched_yield,
14
+ stateless_destroy_torch_distributed_process_group,
15
+)
16
+from vllm.logger import init_logger
17
+from vllm.v1.engine import (
18
+ EEPNotificationType,
19
+ ReconfigureDistributedRequest,
20
+ ReconfigureRankType,
21
+)
22
+from vllm.v1.engine.core import DPEngineCoreProc
23
+
24
+if TYPE_CHECKING:
25
+ from vllm.config import VllmConfig
26
+ from vllm.v1.executor.abstract import Executor
27
+
28
+logger = init_logger(__name__)
29
+
30
+WorkerType = Literal["existing", "new", "removing"]
31
+
32
+
33
+class ScaleUpExistingEngineState(enum.IntEnum):
34
+ WAIT_NEW_CORE_ENGINES_INIT = 0
35
+ CREATE_STANDBY_GROUPS = 1
36
+ TRANSFER_EXPERT_MAPPING = 2
37
+ WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT = 3
38
+ TRANSFER_WEIGHTS = 4
39
+ SYNC_KV_CACHE_MEMORY_SIZE = 5
40
+ SWITCH_AND_PREPARE = 6
41
+ EPLB_RESHUFFLE = 7
42
+ COMPLETE = 8
43
+
44
+
45
+class ScaleUpNewEngineState(enum.IntEnum):
46
+ PREPARE = 0
47
+ EPLB_RESHUFFLE = 1
48
+ COMPLETE = 2
49
+
50
+
51
+class ScaleDownRemainingEngineState(enum.IntEnum):
52
+ PREPARE = 0
53
+ EPLB_RESHUFFLE = 1
54
+ SWITCH_AND_PREPARE = 2
55
+ COMPLETE = 3
56
+
57
+
58
+class ScaleDownRemovingEngineState(enum.IntEnum):
59
+ PREPARE = 0
60
+ EPLB_RESHUFFLE = 1
61
+ COMPLETE = 2
62
+
63
+
64
+class _BarrierTimeoutError(RuntimeError):
65
+ """
66
+ Exception raised for timeout
67
+ in the first stage of our two-staged
68
+ TCPStore based barrier to synchronize the
69
+ execution of all engines in the DP group.
70
+ """
71
+
72
+
73
+class ElasticEPScalingState:
74
+ def __init__(
75
+ self,
76
+ model_executor: "Executor",
77
+ engine_core: "DPEngineCoreProc",
78
+ vllm_config: "VllmConfig",
79
+ new_parallel_config: ParallelConfig,
80
+ worker_type: WorkerType,
81
+ scale_type: Literal["scale_up", "scale_down"],
82
+ reconfig_request: ReconfigureDistributedRequest | None = None,
83
+ ):
84
+ self.model_executor_ref = weakref.ref(model_executor)
85
+ self.engine_core_ref = weakref.ref(engine_core)
86
+ self.vllm_config = vllm_config
87
+ self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None
88
+ self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None
89
+ self.new_parallel_config: ParallelConfig = new_parallel_config
90
+ self.new_dp_group: torch.distributed.ProcessGroup | None = (
91
+ self.engine_core.dp_group if worker_type == "new" else None
92
+ )
93
+ self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None
94
+ self.worker_type = worker_type
95
+ self.scale_type = scale_type
96
+ self.reconfig_request = reconfig_request
97
+
98
+ if scale_type == "scale_up":
99
+ self.state = (
100
+ ScaleUpNewEngineState.PREPARE
101
+ if worker_type == "new"
102
+ else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
103
+ )
104
+ else:
105
+ self.state = (
106
+ ScaleDownRemovingEngineState.PREPARE
107
+ if worker_type == "removing"
108
+ else ScaleDownRemainingEngineState.PREPARE
109
+ )
110
+
111
+ @property
112
+ def model_executor(self) -> "Executor":
113
+ model_executor = self.model_executor_ref()
114
+ if model_executor is None:
115
+ raise RuntimeError("Model executor has been garbage collected")
116
+ return model_executor
117
+
118
+ @property
119
+ def engine_core(self) -> "DPEngineCoreProc":
120
+ engine_core = self.engine_core_ref()
121
+ if engine_core is None:
122
+ raise RuntimeError("Engine core has been garbage collected")
123
+ return engine_core
124
+
125
+ def progress(self) -> bool:
126
+ if self.scale_type == "scale_up":
127
+ return (
128
+ self._progress_new_engine()
129
+ if self.worker_type == "new"
130
+ else self._progress_existing_engine()
131
+ )
132
+ return (
133
+ self._progress_removing_engine()
134
+ if self.worker_type == "removing"
135
+ else self._progress_remaining_engine()
136
+ )
137
+
138
+ def _execute_tcp_store_barrier(
139
+ self, dp_store, group_rank, group_size, barrier_id, timeout=None
140
+ ):
141
+ arrival_key = f"arrival_{barrier_id}_{group_rank}"
142
+ dp_store.set(arrival_key, b"1")
143
+
144
+ start_time = time.time()
145
+ processes_arrived: set[int] = set()
146
+
147
+ while len(processes_arrived) < group_size:
148
+ if (
149
+ timeout is not None
150
+ and time.time() - start_time > timeout.total_seconds()
151
+ ):
152
+ raise _BarrierTimeoutError(
153
+ f"Barrier timed out after {timeout.total_seconds()} seconds"
154
+ )
155
+
156
+ for i in range(group_size):
157
+ if i in processes_arrived:
158
+ continue
159
+
160
+ key = f"arrival_{barrier_id}_{i}"
161
+ present = dp_store.check([key])
162
+ if present:
163
+ processes_arrived.add(i)
164
+
165
+ if len(processes_arrived) < group_size:
166
+ sched_yield()
167
+
168
+ def _staged_barrier(self, use_new_group: bool, barrier_name: str) -> bool:
169
+ """
170
+ Execute a two-staged barrier to synchronize all engines in the DP group.
171
+
172
+ Some DP EngineCores may receive the reconfiguration notifications
173
+ later than others, and already proceed to engine step (model forward)
174
+ in the busy loop.
175
+ In this case, EngineCores that already proceed to reconfiguration
176
+ should skip reconfiguration and execute model forward for one more
177
+ step, so in the next step, all EngineCores will be synchronized.
178
+ We use a two-staged barrier to achieve this. The first time each
179
+ EngineCore executes the barrier, if a timeout is reached before the
180
+ barrier completes, that means some EngineCores have already entered
181
+ engine step. The EngineCores that timed out will then proceed to
182
+ engine step, and will synchronize with the other EngineCores in the
183
+ next step with a barrier without timeout.
184
+ """
185
+ dp_store = self.new_dp_store if use_new_group else self.old_dp_store
186
+ dp_group = self.new_dp_group if use_new_group else self.old_dp_group
187
+ assert dp_group is not None
188
+
189
+ group_rank = dp_group.rank()
190
+ group_size = dp_group.size()
191
+ barrier_id = f"eep_barrier_{barrier_name}"
192
+ sync_key = f"{barrier_id}_sync"
193
+
194
+ # TODO(yongji): figure out appropriate timeout for the barrier
195
+ timeout = None if dp_store.check([sync_key]) else timedelta(seconds=5)
196
+
197
+ try:
198
+ self._execute_tcp_store_barrier(
199
+ dp_store, group_rank, group_size, barrier_id, timeout=timeout
200
+ )
201
+ torch.distributed.barrier(dp_group)
202
+ if group_rank == 0:
203
+ dp_store.delete_key(sync_key)
204
+ for i in range(group_size):
205
+ dp_store.delete_key(f"arrival_{barrier_id}_{i}")
206
+ return True
207
+ except _BarrierTimeoutError as e:
208
+ if timeout is None:
209
+ raise RuntimeError("Unexpected timeout encountered") from e
210
+ dp_store.compare_set(sync_key, "", b"1")
211
+ return False
212
+
213
+ def _progress_existing_engine(self) -> bool:
214
+ state = self.state
215
+
216
+ if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT:
217
+ return False
218
+
219
+ elif state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS:
220
+ # NOTE(yongji): wait for all existing workers to receive the request
221
+ if (
222
+ int(self.old_dp_store.get("eep_barrier_engine_count"))
223
+ < self.old_dp_group.size()
224
+ ):
225
+ return False
226
+ if not self._staged_barrier(
227
+ use_new_group=False, barrier_name="create_standby_groups"
228
+ ):
229
+ return False
230
+ if self.old_dp_group.rank() == 0:
231
+ self.old_dp_store.delete_key("eep_barrier_engine_count")
232
+ self._create_standby_groups()
233
+ self.state = ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING
234
+ return True
235
+
236
+ elif state == ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING:
237
+ self._transfer_expert_mapping()
238
+ self.state = ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT
239
+ return True
240
+
241
+ elif state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT:
242
+ return False
243
+
244
+ elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS:
245
+ if (
246
+ int(self.old_dp_store.get("eep_barrier_engine_count"))
247
+ < self.old_dp_group.size()
248
+ ):
249
+ return False
250
+ if not self._staged_barrier(
251
+ use_new_group=False, barrier_name="transfer_weights"
252
+ ):
253
+ return False
254
+ if self.old_dp_group.rank() == 0:
255
+ self.old_dp_store.delete_key("eep_barrier_engine_count")
256
+ self._transfer_weights()
257
+ self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE
258
+ return True
259
+
260
+ elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE:
261
+ self._sync_kv_cache_memory_size()
262
+ self.state = ScaleUpExistingEngineState.SWITCH_AND_PREPARE
263
+ return True
264
+
265
+ elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE:
266
+ self._switch_and_prepare()
267
+ self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE
268
+ self.new_dp_store.add("eep_barrier_engine_count", 1)
269
+ return True
270
+
271
+ elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE:
272
+ assert self.new_dp_group is not None
273
+ if (
274
+ int(self.new_dp_store.get("eep_barrier_engine_count"))
275
+ < self.new_dp_group.size()
276
+ ):
277
+ return False
278
+ if not self._staged_barrier(
279
+ use_new_group=True, barrier_name="eplb_reshuffle"
280
+ ):
281
+ return False
282
+ if self.new_dp_group.rank() == 0:
283
+ self.new_dp_store.delete_key("eep_barrier_engine_count")
284
+ self._eplb_reshuffle()
285
+ self.state = ScaleUpExistingEngineState.COMPLETE
286
+ self._update_parallel_config()
287
+ return True
288
+
289
+ else:
290
+ assert self.state == ScaleUpExistingEngineState.COMPLETE
291
+ return True
292
+
293
+ def _progress_new_engine(self) -> bool:
294
+ state = self.state
295
+ assert self.new_dp_group is not None
296
+
297
+ if state == ScaleUpNewEngineState.PREPARE:
298
+ tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")
299
+ torch.distributed.all_reduce(
300
+ tensor,
301
+ op=torch.distributed.ReduceOp.MAX,
302
+ group=self.new_dp_group,
303
+ )
304
+ data = tensor.tolist()
305
+ self.engine_core.engines_running = bool(data[0])
306
+ self.engine_core.current_wave = int(data[1])
307
+ self.engine_core.step_counter = int(data[2])
308
+ self.state = ScaleUpNewEngineState.EPLB_RESHUFFLE
309
+ self.new_dp_store.add("eep_barrier_engine_count", 1)
310
+ return True
311
+
312
+ elif state == ScaleUpNewEngineState.EPLB_RESHUFFLE:
313
+ if (
314
+ int(self.new_dp_store.get("eep_barrier_engine_count"))
315
+ < self.new_dp_group.size()
316
+ ):
317
+ return False
318
+ if not self._staged_barrier(
319
+ use_new_group=True, barrier_name="eplb_reshuffle"
320
+ ):
321
+ return False
322
+ assert self.new_dp_group.rank() > 0
323
+ self._eplb_reshuffle()
324
+ self.state = ScaleUpNewEngineState.COMPLETE
325
+ return True
326
+
327
+ else:
328
+ assert self.state == ScaleUpNewEngineState.COMPLETE
329
+ return True
330
+
331
+ def _progress_remaining_engine(self) -> bool:
332
+ state = self.state
333
+
334
+ if state == ScaleDownRemainingEngineState.PREPARE:
335
+ self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE
336
+ self.old_dp_store.add("eep_barrier_engine_count", 1)
337
+ return True
338
+
339
+ elif state == ScaleDownRemainingEngineState.EPLB_RESHUFFLE:
340
+ if (
341
+ int(self.old_dp_store.get("eep_barrier_engine_count"))
342
+ < self.old_dp_group.size()
343
+ ):
344
+ return False
345
+ if not self._staged_barrier(
346
+ use_new_group=False, barrier_name="eplb_reshuffle"
347
+ ):
348
+ return False
349
+ if self.old_dp_group.rank() == 0:
350
+ self.old_dp_store.delete_key("eep_barrier_engine_count")
351
+ self._eplb_reshuffle_before_scale_down()
352
+ self.state = ScaleDownRemainingEngineState.SWITCH_AND_PREPARE
353
+ # NOTE(yongji): currently, after EPLB reshuffle
354
+ # that redistributes experts to remaining workers, workers
355
+ # to be removed will immediately initiate shutdown;
356
+ # existing workers can no longer execute forward steps using
357
+ # the old setup. In the future, we may keep
358
+ # the removing workers alive a bit longer,
359
+ # e.g., to drain in-batch requests.
360
+ self._create_standby_groups()
361
+ self._switch_and_prepare()
362
+ self._update_parallel_config()
363
+ self.state = ScaleDownRemainingEngineState.COMPLETE
364
+ return True
365
+
366
+ else:
367
+ assert self.state == ScaleDownRemainingEngineState.COMPLETE
368
+ return True
369
+
370
+ def _progress_removing_engine(self) -> bool:
371
+ state = self.state
372
+
373
+ if state == ScaleDownRemovingEngineState.PREPARE:
374
+ self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE
375
+ self.old_dp_store.add("eep_barrier_engine_count", 1)
376
+ return True
377
+
378
+ if state == ScaleDownRemovingEngineState.EPLB_RESHUFFLE:
379
+ if (
380
+ int(self.old_dp_store.get("eep_barrier_engine_count"))
381
+ < self.old_dp_group.size()
382
+ ):
383
+ return False
384
+ if not self._staged_barrier(
385
+ use_new_group=False, barrier_name="eplb_reshuffle"
386
+ ):
387
+ return False
388
+ assert self.old_dp_group.rank() > 0
389
+ self._eplb_reshuffle_before_scale_down()
390
+ self._switch_and_remove()
391
+ self.state = ScaleDownRemovingEngineState.COMPLETE
392
+ self.engine_core._eep_send_engine_core_notification(
393
+ EEPNotificationType.SHUTDOWN_COMPLETE
394
+ )
395
+ self.engine_core.shutdown()
396
+ return True
397
+
398
+ else:
399
+ assert self.state == ScaleDownRemovingEngineState.COMPLETE
400
+ return True
401
+
402
+ def handle_notification(self, notification_type: EEPNotificationType):
403
+ assert self.worker_type != "new"
404
+ if (
405
+ notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY
406
+ and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
407
+ ):
408
+ self.old_dp_store.add("eep_barrier_engine_count", 1)
409
+ self.state = ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS
410
+ elif (
411
+ notification_type == EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY
412
+ and self.state
413
+ == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT
414
+ ):
415
+ self.old_dp_store.add("eep_barrier_engine_count", 1)
416
+ self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS
417
+
418
+ def is_complete(self) -> bool:
419
+ if self.scale_type == "scale_up":
420
+ return (
421
+ self.state == ScaleUpNewEngineState.COMPLETE
422
+ if self.worker_type == "new"
423
+ else self.state == ScaleUpExistingEngineState.COMPLETE
424
+ )
425
+ return (
426
+ self.state == ScaleDownRemovingEngineState.COMPLETE
427
+ if self.worker_type == "removing"
428
+ else self.state == ScaleDownRemainingEngineState.COMPLETE
429
+ )
430
+
431
+ def _create_standby_groups(self):
432
+ self.new_dp_group, self.new_dp_store = (
433
+ self.new_parallel_config.stateless_init_dp_group(return_store=True)
434
+ )
435
+ self.model_executor.collective_rpc(
436
+ "elastic_ep_execute", args=("create_standby_groups", self.reconfig_request)
437
+ )
438
+ if self.old_dp_group.rank() == 0:
439
+ logger.info("[Elastic EP] Created standby communication groups")
440
+
441
+ def _transfer_weights(self):
442
+ assert self.reconfig_request is not None
443
+ old_dp_size = self.old_dp_group.size()
444
+ new_dp_size = self.reconfig_request.new_data_parallel_size
445
+
446
+ self.model_executor.collective_rpc(
447
+ "elastic_ep_execute", args=("transfer_weights", old_dp_size, new_dp_size)
448
+ )
449
+ if self.old_dp_group.rank() == 0:
450
+ logger.info("[Elastic EP] Transferred weights to new workers")
451
+
452
+ def _transfer_expert_mapping(self):
453
+ self.model_executor.collective_rpc(
454
+ "elastic_ep_execute", args=("broadcast_expert_mapping",)
455
+ )
456
+ if self.old_dp_group.rank() == 0:
457
+ logger.info("[Elastic EP] Broadcasted expert mapping to new workers")
458
+
459
+ def _sync_kv_cache_memory_size(self):
460
+ assert self.engine_core.available_gpu_memory_for_kv_cache > 0
461
+ assert self.new_dp_group is not None
462
+ ParallelConfig.sync_kv_cache_memory_size(
463
+ self.new_dp_group,
464
+ self.engine_core.available_gpu_memory_for_kv_cache,
465
+ )
466
+ if self.old_dp_group.rank() == 0:
467
+ logger.info("[Elastic EP] Synced KV cache memory size to new workers")
468
+
469
+ def _switch_and_prepare(self):
470
+ self.model_executor.collective_rpc(
471
+ "elastic_ep_execute", args=("switch_and_prepare",)
472
+ )
473
+ old_dp_group = self.old_dp_group
474
+ stateless_destroy_torch_distributed_process_group(old_dp_group)
475
+ assert self.new_dp_group is not None
476
+ new_dp_group = self.new_dp_group
477
+ self.engine_core.dp_group = new_dp_group
478
+ self.engine_core.dp_rank = new_dp_group.rank()
479
+ self.engine_core.dp_store = self.new_dp_store
480
+ engines_running = int(self.engine_core.engines_running)
481
+ current_wave = self.engine_core.current_wave
482
+ step_counter = self.engine_core.step_counter
483
+ tensor = torch.tensor(
484
+ [engines_running, current_wave, step_counter],
485
+ dtype=torch.int32,
486
+ device="cpu",
487
+ )
488
+ torch.distributed.all_reduce(
489
+ tensor, op=torch.distributed.ReduceOp.MAX, group=new_dp_group
490
+ )
491
+ data = tensor.tolist()
492
+ self.engine_core.engines_running = bool(data[0])
493
+ self.engine_core.current_wave = int(data[1])
494
+ self.engine_core.step_counter = int(data[2])
495
+ if new_dp_group.rank() == 0:
496
+ self.engine_core._eep_send_engine_core_notification(
497
+ EEPNotificationType.RECONFIGURE_FINISHED
498
+ )
499
+ logger.info("[Elastic EP] Switched to new setup")
500
+
501
+ def _eplb_reshuffle(self):
502
+ self.model_executor.collective_rpc(
503
+ "elastic_ep_execute", args=("perform_eplb_reshuffle",)
504
+ )
505
+ assert self.new_dp_group is not None
506
+ if self.new_dp_group.rank() == 0:
507
+ logger.info("[Elastic EP] EPLB reshuffle completed")
508
+
509
+ def _eplb_reshuffle_before_scale_down(self):
510
+ assert self.reconfig_request is not None
511
+ self.model_executor.collective_rpc(
512
+ "elastic_ep_execute",
513
+ args=(
514
+ "perform_eplb_reshuffle",
515
+ self.reconfig_request.new_data_parallel_size,
516
+ ),
517
+ )
518
+ if self.old_dp_group.rank() == 0:
519
+ logger.info("[Elastic EP] EPLB reshuffle completed")
520
+
521
+ def _switch_and_remove(self):
522
+ self.model_executor.collective_rpc(
523
+ "elastic_ep_execute", args=("switch_and_remove",)
524
+ )
525
+
526
+ def _update_parallel_config(self):
527
+ assert self.reconfig_request is not None
528
+ reconfig_request = self.reconfig_request
529
+ parallel_config = self.vllm_config.parallel_config
530
+ parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size
531
+ if (
532
+ reconfig_request.new_data_parallel_rank
533
+ != ReconfigureRankType.KEEP_CURRENT_RANK
534
+ ):
535
+ parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
536
+ if (
537
+ reconfig_request.new_data_parallel_rank_local
538
+ != ReconfigureRankType.KEEP_CURRENT_RANK
539
+ ):
540
+ parallel_config.data_parallel_rank_local = (
541
+ reconfig_request.new_data_parallel_rank_local
542
+ )
543
+ parallel_config.data_parallel_master_ip = (
544
+ reconfig_request.new_data_parallel_master_ip
545
+ )
546
+ parallel_config.data_parallel_master_port = (
547
+ reconfig_request.new_data_parallel_master_port
548
+ )
549
+ parallel_config._data_parallel_master_port_list = (
550
+ reconfig_request.new_data_parallel_master_port_list
551
+ )
552
+ parallel_config._stateless_world_group_port_list = (
553
+ reconfig_request.new_stateless_world_group_port_list
554
+ )
555
+ parallel_config._stateless_dp_group_port_list = (
556
+ reconfig_request.new_stateless_dp_group_port_list
557
+ )
558
+ parallel_config._stateless_ep_group_port_list = (
559
+ reconfig_request.new_stateless_ep_group_port_list
560
+ )
561
+ parallel_config._stateless_eplb_group_port_list = (
562
+ reconfig_request.new_stateless_eplb_group_port_list
563
+ )
vllm/distributed/elastic_ep/standby_state.py
new
+117
@@ -0,0 +1,117 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.distributed.parallel_state import (
6
+ _init_stateless_group,
7
+ _node_count,
8
+ get_pp_group,
9
+ get_tp_group,
10
+ get_world_group,
11
+)
12
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
13
+
14
+_STANDBY_WORLD: StatelessGroupCoordinator | None = None
15
+_STANDBY_WORLD_NODE_COUNT: int | None = None
16
+_STANDBY_DP: StatelessGroupCoordinator | None = None
17
+_STANDBY_EP: StatelessGroupCoordinator | None = None
18
+_STANDBY_EPLB: StatelessGroupCoordinator | None = None
19
+
20
+
21
+def get_standby_dp_group() -> StatelessGroupCoordinator | None:
22
+ return _STANDBY_DP
23
+
24
+
25
+def get_standby_ep_group() -> StatelessGroupCoordinator | None:
26
+ return _STANDBY_EP
27
+
28
+
29
+def get_standby_eplb_group() -> StatelessGroupCoordinator | None:
30
+ return _STANDBY_EPLB
31
+
32
+
33
+def get_standby_world_group() -> StatelessGroupCoordinator | None:
34
+ return _STANDBY_WORLD
35
+
36
+
37
+def create_standby_groups(
38
+ new_dp_size: int,
39
+ new_world_size_across_dp: int,
40
+ master_ip: str,
41
+ world_group_ports: list[list[int]],
42
+ dp_group_ports: list[list[int]],
43
+ ep_group_ports: list[list[int]],
44
+ eplb_group_ports: list[list[int]] | None = None,
45
+ backend: str | None = None,
46
+) -> None:
47
+ global \
48
+ _STANDBY_WORLD, \
49
+ _STANDBY_WORLD_NODE_COUNT, \
50
+ _STANDBY_DP, \
51
+ _STANDBY_EP, \
52
+ _STANDBY_EPLB
53
+
54
+ assert new_world_size_across_dp == torch.distributed.get_world_size() * new_dp_size
55
+ world_group = get_world_group()
56
+ assert isinstance(world_group, StatelessGroupCoordinator)
57
+ backend = backend or world_group.backend
58
+
59
+ standby_world_ranks = [list(range(new_world_size_across_dp))]
60
+ _STANDBY_WORLD = _init_stateless_group(
61
+ standby_world_ranks,
62
+ "world",
63
+ world_group_ports,
64
+ master_ip,
65
+ backend,
66
+ use_device_communicator=False,
67
+ )
68
+ _STANDBY_WORLD_NODE_COUNT = _node_count(_STANDBY_WORLD.tcp_store_group)
69
+
70
+ tp_size = get_tp_group().world_size
71
+ pp_size = get_pp_group().world_size
72
+
73
+ all_ranks = torch.arange(new_world_size_across_dp).reshape(
74
+ -1, new_dp_size, pp_size, tp_size
75
+ )
76
+ standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)
77
+ standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
78
+ _STANDBY_DP = _init_stateless_group(
79
+ standby_dp_ranks, "dp", dp_group_ports, master_ip, backend
80
+ )
81
+
82
+ standby_ep_ranks = (
83
+ all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
84
+ )
85
+ standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
86
+ _STANDBY_EP = _init_stateless_group(
87
+ standby_ep_ranks, "ep", ep_group_ports, master_ip, backend
88
+ )
89
+
90
+ if eplb_group_ports is not None:
91
+ _STANDBY_EPLB = _init_stateless_group(
92
+ standby_ep_ranks, "eplb", eplb_group_ports, master_ip, backend
93
+ )
94
+
95
+
96
+def pop_standby_groups() -> dict:
97
+ """Return all standby groups and clear the standby state."""
98
+ global \
99
+ _STANDBY_WORLD, \
100
+ _STANDBY_WORLD_NODE_COUNT, \
101
+ _STANDBY_DP, \
102
+ _STANDBY_EP, \
103
+ _STANDBY_EPLB
104
+
105
+ result = dict(
106
+ world=_STANDBY_WORLD,
107
+ dp=_STANDBY_DP,
108
+ ep=_STANDBY_EP,
109
+ eplb=_STANDBY_EPLB,
110
+ node_count=_STANDBY_WORLD_NODE_COUNT,
111
+ )
112
+ _STANDBY_WORLD = None
113
+ _STANDBY_WORLD_NODE_COUNT = None
114
+ _STANDBY_DP = None
115
+ _STANDBY_EP = None
116
+ _STANDBY_EPLB = None
117
+ return result
vllm/distributed/eplb/async_worker.py
-4
@@ -24,7 +24,6 @@ logger = init_logger(__name__)
24
25
def start_async_worker(
26
state: "EplbState",
27
- rank_mapping: dict[int, int] | None = None,
27
is_profile: bool = False,
28
) -> threading.Thread:
29
eplb_group = get_eplb_group().device_group
@@ -45,7 +44,6 @@ def start_async_worker(
44
eplb_group=eplb_group,
45
cuda_stream=cuda_stream,
46
is_profile=is_profile,
48
- rank_mapping=rank_mapping,
47
)
48
)
49
except Exception as exc: # pragma: no cover - diagnostic path
@@ -107,7 +105,6 @@ async def transfer_run_periodically(
105
eplb_group: ProcessGroup,
106
cuda_stream: torch.cuda.Stream,
107
is_profile: bool = False,
110
- rank_mapping: dict[int, int] | None = None,
108
) -> None:
109
while True:
110
await asyncio.to_thread(state.rearrange_event.wait)
@@ -176,7 +173,6 @@ async def transfer_run_periodically(
173
ep_group=eplb_group,
174
is_profile=is_profile,
175
cuda_stream=cuda_stream,
179
- rank_mapping=rank_mapping,
176
)
177
event = torch.cuda.Event(blocking=False)
178
cuda_stream.record_event(event)
vllm/distributed/eplb/eplb_state.py
+98
-214
@@ -40,6 +40,7 @@ from vllm.distributed.parallel_state import (
40
get_node_count,
41
in_the_same_node_as,
42
)
43
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
44
from vllm.distributed.utils import StatelessProcessGroup
45
from vllm.logger import init_logger
46
from vllm.model_executor.models.interfaces import MixtureOfExperts
@@ -302,6 +303,14 @@ class EplbState:
303
"""
304
CUDA device index for the async EPLB worker thread.
305
"""
306
+ self.num_valid_physical_experts: int = 0
307
+ """
308
+ Number of valid physical experts.
309
+ This is the number of physical experts that are
310
+ actually mapped to logical experts. In elastic EP,
311
+ newly started EP ranks may not have physical experts
312
+ mapped yet.
313
+ """
314
if self.device.type == "cuda":
315
self.cuda_device_index = self.device.index
316
if self.cuda_device_index is None and torch.cuda.is_available():
@@ -367,9 +376,6 @@ class EplbState:
376
self,
377
model: MixtureOfExperts,
378
model_config: ModelConfig,
370
- global_expert_load: torch.Tensor | None = None,
371
- old_global_expert_indices: torch.Tensor | None = None,
372
- rank_mapping: dict[int, int] | None = None,
379
):
380
"""
381
Build the initial EPLB state.
@@ -462,75 +468,15 @@ class EplbState:
468
)
469
self.expert_rearrangement_step_interval = eplb_step_interval
470
465
- # Set the policy based on the selected eplb algorithm type.
471
policy_type = self.parallel_config.eplb_config.policy
472
self.policy = EPLB_POLICIES[policy_type]
473
logger.debug("Selected EPLB policy: %s", policy_type)
469
- if global_expert_load is not None:
470
- ep_group = get_ep_group().device_group
471
- assert global_expert_load.shape == (
472
- model.num_moe_layers,
473
- model.num_logical_experts,
474
- )
475
- assert global_expert_load.dtype == torch.int64
476
-
477
- num_replicas = model.num_physical_experts
478
- num_groups = model.num_expert_groups
479
- num_nodes = get_node_count()
480
- num_gpus = ep_group.size()
481
-
482
- if num_gpus % num_nodes != 0:
483
- num_nodes = 1
484
- logger.warning_once(
485
- f"num_gpus % num_nodes != 0, "
486
- "not using hierarchical rearrangement algorithm.\n"
487
- f"{num_gpus=}, {num_nodes=}"
488
- )
489
-
490
- # Get new expert mappings
491
- (
492
- new_physical_to_logical_map,
493
- new_logical_to_physical_map,
494
- new_logical_replica_count,
495
- ) = self.policy.rebalance_experts(
496
- global_expert_load,
497
- num_replicas,
498
- num_groups,
499
- num_nodes,
500
- num_gpus,
501
- )
502
-
503
- max_physical_slots = new_logical_to_physical_map.shape[-1]
504
- assert max_physical_slots <= logical_to_physical_map.shape[-1]
505
- new_logical_to_physical_map = torch.nn.functional.pad(
506
- new_logical_to_physical_map,
507
- (0, logical_to_physical_map.shape[-1] - max_physical_slots),
508
- value=-1,
509
- )
510
- physical_to_logical_map = new_physical_to_logical_map.to(self.device)
511
- logical_to_physical_map.copy_(new_logical_to_physical_map)
512
- logical_replica_count.copy_(new_logical_replica_count)
513
- else:
514
- new_physical_to_logical_map = None
515
-
516
- new_logical_to_physical_map = None
474
518
- new_logical_replica_count = None
475
model.set_eplb_state(
476
expert_load_pass,
477
logical_to_physical_map,
478
logical_replica_count,
479
)
524
- if global_expert_load is not None:
525
- rearrange_expert_weights_inplace(
526
- old_global_expert_indices,
527
- new_physical_to_logical_map,
528
- model.expert_weights,
529
- ep_group,
530
- False,
531
- rank_mapping,
532
- )
533
- self.expert_rearrangement_step = 0
480
481
expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]
482
@@ -561,11 +507,12 @@ class EplbState:
507
recv_dst_rows=np.array([]),
508
),
509
cuda_device_index=self.cuda_device_index,
564
- new_physical_to_logical_map=new_physical_to_logical_map,
565
- new_logical_to_physical_map=new_logical_to_physical_map,
566
- new_logical_replica_count=new_logical_replica_count,
510
+ new_physical_to_logical_map=None,
511
+ new_logical_to_physical_map=None,
512
+ new_logical_replica_count=None,
513
)
514
self.model_states[model_config.compute_hash()] = model_state
515
+ self.num_valid_physical_experts = model.num_physical_experts
516
517
def step(
518
self,
@@ -696,8 +643,6 @@ class EplbState:
643
def rearrange(
644
self,
645
is_profile: bool = False,
699
- execute_shuffle: bool = True,
700
- global_expert_loads: list[torch.Tensor] | None = None,
646
rank_mapping: dict[int, int] | None = None,
647
) -> torch.Tensor | None:
648
"""
@@ -707,12 +652,6 @@ class EplbState:
652
is_profile (bool): If `True`, perform a dummy rearrangement.
653
This is used in `profile_run` to reserve enough memory,
654
no memory movement will be performed. Default is False.
710
- execute_shuffle (bool): If `True`, execute the shuffle
711
- in elastic expert parallel (EEP). Default is True.
712
- global_expert_loads (list[torch.Tensor] | None): The global expert
713
- loads when scaling is done in EEP.
714
- List of expert loads for the main and drafter
715
- (when spec decode is used) models.
655
rank_mapping (dict[int, int] | None): The rank mapping
656
when scaling is done in EEP.
657
"""
@@ -734,67 +673,34 @@ class EplbState:
673
"(profile)" if is_profile else "",
674
)
675
737
- if global_expert_loads is None:
738
- # Map the physical expert load to global logical experts
739
- global_expert_load_windows = []
740
- if not execute_shuffle:
741
- num_models = torch.tensor(
742
- [len(self.model_states)], dtype=torch.int32, device="cpu"
743
- )
744
- torch.distributed.broadcast(
745
- num_models, group=get_ep_group().cpu_group, group_src=0
746
- )
747
-
748
- for eplb_model_state in self.model_states.values():
749
- logical_expert_load_window = torch.zeros(
750
- self.expert_load_window_size,
751
- eplb_model_state.model.num_moe_layers,
752
- eplb_model_state.model.num_logical_experts,
753
- dtype=eplb_model_state.expert_load_window.dtype,
754
- device=eplb_model_state.expert_load_window.device,
755
- )
756
- logical_expert_load_window.scatter_add_(
757
- dim=-1,
758
- index=eplb_model_state.physical_to_logical_map.unsqueeze(0)
759
- .expand_as(eplb_model_state.expert_load_window)
760
- .long(),
761
- src=eplb_model_state.expert_load_window,
762
- )
763
-
764
- if not execute_shuffle:
765
- metadata = torch.tensor(
766
- [
767
- eplb_model_state.model.num_moe_layers,
768
- eplb_model_state.model.num_logical_experts,
769
- eplb_model_state.physical_to_logical_map.shape[1],
770
- ],
771
- dtype=torch.int32,
772
- device="cpu",
773
- )
774
- torch.distributed.broadcast(
775
- metadata, group=get_ep_group().cpu_group, group_src=0
776
- )
777
-
778
- global_expert_load_window = logical_expert_load_window.sum(dim=0)
779
- global_expert_load_windows.append(global_expert_load_window)
780
- # Perform all-reduce to get the expert load across all ranks for each model
781
- global_expert_load_windows = self._allreduce_list(
782
- global_expert_load_windows
676
+ # Map the physical expert load to global logical experts
677
+ global_expert_load_windows = []
678
+ for eplb_model_state in self.model_states.values():
679
+ expert_load_window = eplb_model_state.expert_load_window[
680
+ :, :, : self.num_valid_physical_experts
681
+ ]
682
+ logical_expert_load_window = torch.zeros(
683
+ self.expert_load_window_size,
684
+ eplb_model_state.model.num_moe_layers,
685
+ eplb_model_state.model.num_logical_experts,
686
+ dtype=eplb_model_state.expert_load_window.dtype,
687
+ device=eplb_model_state.expert_load_window.device,
688
+ )
689
+ logical_expert_load_window.scatter_add_(
690
+ dim=-1,
691
+ index=eplb_model_state.physical_to_logical_map[
692
+ :, : self.num_valid_physical_experts
693
+ ]
694
+ .unsqueeze(0)
695
+ .expand_as(expert_load_window)
696
+ .long(),
697
+ src=expert_load_window,
698
)
784
- if not execute_shuffle:
785
- for eplb_model_state, global_expert_load_window in zip(
786
- self.model_states.values(), global_expert_load_windows
787
- ):
788
- # (num_moe_layers, old_num_physical_experts)
789
- old_global_expert_indices = eplb_model_state.physical_to_logical_map
790
- torch.distributed.broadcast(
791
- old_global_expert_indices, group=ep_group, group_src=0
792
- )
793
- if not execute_shuffle:
794
- return global_expert_load_windows
795
- else:
796
- assert execute_shuffle
797
- global_expert_load_windows = global_expert_loads
699
+
700
+ global_expert_load_window = logical_expert_load_window.sum(dim=0)
701
+ global_expert_load_windows.append(global_expert_load_window)
702
+ # Perform all-reduce to get the expert load across all ranks for each model
703
+ global_expert_load_windows = self._allreduce_list(global_expert_load_windows)
704
705
# TODO(bowen): Treat differently for prefill and decode nodes
706
eplb_model_state = next(iter(self.model_states.values()))
@@ -806,8 +712,10 @@ class EplbState:
712
# NOTE(yongji): scale down, we need to rebalance the experts on
713
# remaining GPUs, transfer the experts while we haven't shutdown
714
# the GPUs to be released.
809
- cpu_group = get_ep_group().cpu_group
810
- num_nodes = _node_count_with_rank_mapping(cpu_group, rank_mapping)
715
+ coordinator = get_ep_group()
716
+ assert isinstance(coordinator, StatelessGroupCoordinator)
717
+ tcp_store_group = coordinator.tcp_store_group
718
+ num_nodes = _node_count_with_rank_mapping(tcp_store_group, rank_mapping)
719
num_gpus = sum(new_rank != -1 for new_rank in rank_mapping.values())
720
num_replicas = (
721
num_replicas // ep_group.size() * num_gpus
@@ -933,7 +841,6 @@ class EplbState:
841
if self.async_worker is None:
842
self.async_worker = start_async_worker(
843
self,
936
- rank_mapping=rank_mapping,
844
is_profile=is_profile,
845
)
846
@@ -1089,83 +996,6 @@ class EplbState:
996
model_state.new_logical_to_physical_map = None
997
model_state.new_logical_replica_count = None
998
1092
- @staticmethod
1093
- def recv_state() -> tuple[list[torch.Tensor], list[torch.Tensor]]:
1094
- """
1095
- Receive the expert load and old placement from the master rank.
1096
- """
1097
- ep_group = get_ep_group()
1098
- num_models = torch.empty(1, dtype=torch.int32, device="cpu")
1099
- torch.distributed.broadcast(num_models, group=ep_group.cpu_group, group_src=0)
1100
- num_models = num_models.item()
1101
- global_expert_loads = []
1102
- old_global_expert_indices_per_model = []
1103
- for _ in range(num_models):
1104
- metadata = torch.empty(3, dtype=torch.int32, device="cpu")
1105
- torch.distributed.broadcast(metadata, group=ep_group.cpu_group, group_src=0)
1106
- num_moe_layers, num_logical_experts, num_old_physical_experts = (
1107
- metadata.tolist()
1108
- )
1109
- global_expert_load = torch.zeros(
1110
- (num_moe_layers, num_logical_experts),
1111
- dtype=torch.int64,
1112
- device=ep_group.device,
1113
- )
1114
- all_reduce(global_expert_load, group=ep_group.device_group)
1115
- old_global_expert_indices = torch.empty(
1116
- (num_moe_layers, num_old_physical_experts),
1117
- dtype=torch.int64,
1118
- device=ep_group.device,
1119
- )
1120
- torch.distributed.broadcast(
1121
- old_global_expert_indices,
1122
- group=ep_group.device_group,
1123
- group_src=0,
1124
- )
1125
- global_expert_loads.append(global_expert_load)
1126
- old_global_expert_indices_per_model.append(old_global_expert_indices)
1127
- return global_expert_loads, old_global_expert_indices_per_model
1128
-
1129
- @classmethod
1130
- def get_eep_state(
1131
- cls, parallel_config: ParallelConfig
1132
- ) -> tuple[
1133
- list[torch.Tensor] | None,
1134
- list[torch.Tensor] | None,
1135
- dict[int, int] | None,
1136
- ]:
1137
- num_local_physical_experts = torch.empty(1, dtype=torch.int32, device="cpu")
1138
- torch.distributed.broadcast(
1139
- num_local_physical_experts,
1140
- group=get_ep_group().cpu_group,
1141
- group_src=0,
1142
- )
1143
- num_local_physical_experts = int(num_local_physical_experts.item())
1144
- new_ep_size = get_ep_group().world_size
1145
- global_expert_loads, old_global_expert_indices_per_model = (
1146
- EplbState.recv_state()
1147
- )
1148
-
1149
- # EP configuration for all models has to be the same so as eplb config
1150
- num_logical_experts = global_expert_loads[0].shape[1]
1151
- parallel_config.eplb_config.num_redundant_experts = (
1152
- num_local_physical_experts * new_ep_size - num_logical_experts
1153
- )
1154
- assert (
1155
- old_global_expert_indices_per_model[0].shape[1] % num_local_physical_experts
1156
- == 0
1157
- )
1158
- old_ep_size = (
1159
- old_global_expert_indices_per_model[0].shape[1]
1160
- // num_local_physical_experts
1161
- )
1162
- rank_mapping = {old_ep_rank: old_ep_rank for old_ep_rank in range(old_ep_size)}
1163
- return (
1164
- global_expert_loads,
1165
- old_global_expert_indices_per_model,
1166
- rank_mapping,
1167
- )
1168
-
999
def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
1000
"""
1001
All-reduce a list of tensors.
@@ -1203,6 +1033,60 @@ class EplbState:
1033
load_pass_list.append(eplb_model_state.expert_load_pass.clone())
1034
return self._allreduce_list(load_pass_list)
1035
1036
+ @classmethod
1037
+ def from_mapping(
1038
+ cls,
1039
+ model: MixtureOfExperts,
1040
+ model_config: ModelConfig,
1041
+ device: torch.device,
1042
+ parallel_config: ParallelConfig,
1043
+ expanded_physical_to_logical: torch.Tensor,
1044
+ num_valid_physical_experts: int,
1045
+ ) -> "EplbState":
1046
+ eplb_state = cls(
1047
+ parallel_config=parallel_config,
1048
+ device=device,
1049
+ )
1050
+ eplb_state.add_model(
1051
+ model=model,
1052
+ model_config=model_config,
1053
+ )
1054
+ eplb_state.num_valid_physical_experts = num_valid_physical_experts
1055
+ num_moe_layers = expanded_physical_to_logical.shape[0]
1056
+ num_physical_experts = expanded_physical_to_logical.shape[1]
1057
+ eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
1058
+ eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)
1059
+
1060
+ logical_to_physical_map = torch.full(
1061
+ (
1062
+ num_moe_layers,
1063
+ model.num_logical_experts,
1064
+ eplb_model_state.logical_to_physical_map.shape[2],
1065
+ ),
1066
+ -1,
1067
+ dtype=torch.int64,
1068
+ )
1069
+ logical_replica_count = torch.zeros(
1070
+ (num_moe_layers, model.num_logical_experts),
1071
+ dtype=torch.int64,
1072
+ )
1073
+ expanded_physical_to_logical_numpy = expanded_physical_to_logical.cpu().numpy()
1074
+ for layer_idx in range(num_moe_layers):
1075
+ for phys_idx in range(num_physical_experts):
1076
+ logical_idx = expanded_physical_to_logical_numpy[layer_idx, phys_idx]
1077
+ if logical_idx >= 0:
1078
+ replica_idx = logical_replica_count[layer_idx, logical_idx]
1079
+ logical_to_physical_map[layer_idx, logical_idx, replica_idx] = (
1080
+ phys_idx
1081
+ )
1082
+ logical_replica_count[layer_idx, logical_idx] += 1
1083
+
1084
+ logical_to_physical_map = logical_to_physical_map.to(device)
1085
+ logical_replica_count = logical_replica_count.to(device)
1086
+ eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)
1087
+ eplb_model_state.logical_replica_count.copy_(logical_replica_count)
1088
+ return eplb_state
1089
+
1090
1091
@dataclass
1092
class EplbLayerState:
vllm/distributed/eplb/rebalance_execute.py
+56
-24
@@ -19,6 +19,8 @@ from torch.distributed import (
19
get_global_rank,
20
)
21
22
+from vllm.distributed.parallel_state import get_ep_group
23
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
24
from vllm.logger import init_logger
25
26
logger = init_logger(__name__)
@@ -249,10 +251,18 @@ def move_to_buffer(
251
b[dst].copy_(w[src_local], non_blocking=True)
252
253
p2p_ops: list[P2POp] = []
254
+ if isinstance(get_ep_group(), StatelessGroupCoordinator):
255
+ ep_group = get_ep_group()
256
+ is_stateless = True
257
+ else:
258
+ is_stateless = False
259
253
- # Pre-compute global ranks mapping
260
+ # Pre-compute global ranks mapping (only needed for non-stateless groups)
261
ep_size = ep_group.size()
255
- rank_to_global = {rank: get_global_rank(ep_group, rank) for rank in range(ep_size)}
262
+ if not is_stateless:
263
+ rank_to_global = {
264
+ rank: get_global_rank(ep_group, rank) for rank in range(ep_size)
265
+ }
266
267
# 2. Post sends
268
if send_count > 0:
@@ -284,15 +294,23 @@ def move_to_buffer(
294
if recver_pos < len(ranks_to_recv):
295
recv_ranks.append(ranks_to_recv[recver_pos])
296
for dst in recv_ranks:
287
- dst_global = rank_to_global[dst]
288
- p2p_ops += [
289
- P2POp(
290
- torch.distributed.isend,
291
- w[src],
292
- dst_global,
293
- )
294
- for w in expert_weights
295
- ]
297
+ if is_stateless:
298
+ for w in expert_weights:
299
+ op = object.__new__(P2POp)
300
+ op.op = torch.distributed.isend
301
+ op.tensor = w[src]
302
+ op.group_peer = dst
303
+ p2p_ops.append(op)
304
+ else:
305
+ dst_global = rank_to_global[dst]
306
+ p2p_ops += [
307
+ P2POp(
308
+ torch.distributed.isend,
309
+ w[src],
310
+ dst_global,
311
+ )
312
+ for w in expert_weights
313
+ ]
314
315
# 3. Post recvs
316
if recv_count > 0:
@@ -321,26 +339,40 @@ def move_to_buffer(
339
src = ranks_to_send[recver_pos // num_dst_per_sender]
340
else:
341
src = ranks_to_send[recver_pos - remainder_start]
324
- src_global = rank_to_global[src]
325
- p2p_ops += [
326
- P2POp(
327
- torch.distributed.irecv,
328
- b[dst],
329
- src_global,
330
- )
331
- for b in expert_weights_buffers
332
- ]
342
+ if is_stateless:
343
+ for b in expert_weights_buffers:
344
+ op = object.__new__(P2POp)
345
+ op.op = torch.distributed.irecv
346
+ op.tensor = b[dst]
347
+ op.group_peer = src
348
+ p2p_ops.append(op)
349
+ else:
350
+ src_global = rank_to_global[src]
351
+ p2p_ops += [
352
+ P2POp(
353
+ torch.distributed.irecv,
354
+ b[dst],
355
+ src_global,
356
+ )
357
+ for b in expert_weights_buffers
358
+ ]
359
360
# 4. Execute the P2P operations. The real communication happens here.
361
if p2p_ops and cuda_stream is not None:
362
with torch.cuda.stream(cuda_stream):
363
+ if is_stateless:
364
+ ep_group.device_communicator.batch_isend_irecv(p2p_ops)
365
+ else:
366
+ reqs = batch_isend_irecv(p2p_ops)
367
+ for req in reqs:
368
+ req.wait()
369
+ elif p2p_ops:
370
+ if is_stateless:
371
+ ep_group.device_communicator.batch_isend_irecv(p2p_ops)
372
+ else:
373
reqs = batch_isend_irecv(p2p_ops)
374
for req in reqs:
375
req.wait()
340
- elif p2p_ops:
341
- reqs = batch_isend_irecv(p2p_ops)
342
- for req in reqs:
343
- req.wait()
376
# wait for the communication to finish
377
return (
378
is_unchanged,
vllm/distributed/parallel_state.py
+208
-23
@@ -33,7 +33,7 @@ from contextlib import contextmanager, nullcontext
33
from dataclasses import dataclass
34
from datetime import timedelta
35
from multiprocessing import shared_memory
36
-from typing import Any, Protocol
36
+from typing import TYPE_CHECKING, Any, Protocol
37
from unittest.mock import patch
38
39
import torch
@@ -55,6 +55,9 @@ from vllm.utils.torch_utils import (
55
direct_register_custom_op,
56
)
57
58
+if TYPE_CHECKING:
59
+ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
60
+
61
62
@dataclass
63
class GraphCaptureContext:
@@ -1157,6 +1160,55 @@ def init_model_parallel_group(
1160
)
1161
1162
1163
+def _init_stateless_group(
1164
+ group_ranks: list[list[int]],
1165
+ group_name: str,
1166
+ group_ports: list[list[int]],
1167
+ host: str,
1168
+ backend: str,
1169
+ use_device_communicator: bool = True,
1170
+) -> "StatelessGroupCoordinator":
1171
+ """Create a StatelessGroupCoordinator with the given parameters."""
1172
+ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
1173
+
1174
+ world = get_world_group()
1175
+ return StatelessGroupCoordinator(
1176
+ group_ranks=group_ranks,
1177
+ local_rank=world.local_rank,
1178
+ torch_distributed_backend=backend,
1179
+ use_device_communicator=use_device_communicator,
1180
+ group_name=group_name,
1181
+ host=host,
1182
+ group_ports=group_ports,
1183
+ global_rank=world.rank,
1184
+ global_world_size=world.world_size,
1185
+ )
1186
+
1187
+
1188
+def _replace_active_groups(
1189
+ *,
1190
+ world: GroupCoordinator | None,
1191
+ dp: GroupCoordinator | None,
1192
+ ep: GroupCoordinator | None,
1193
+ eplb: GroupCoordinator | None,
1194
+ node_count: int | None,
1195
+) -> None:
1196
+ """Destroy the current DP/EP/WORLD/EPLB groups and replace them.
1197
+
1198
+ Destruction is collective — all ranks in the old groups must call this
1199
+ function together. Pass all-``None`` to tear down without replacement.
1200
+ """
1201
+ global _WORLD, _DP, _EP, _EPLB, _NODE_COUNT
1202
+ for group in (_DP, _EP, _WORLD, _EPLB):
1203
+ if group is not None:
1204
+ group.destroy()
1205
+ _WORLD = world
1206
+ _DP = dp
1207
+ _EP = ep
1208
+ _EPLB = eplb
1209
+ _NODE_COUNT = node_count
1210
+
1211
+
1212
_TP: GroupCoordinator | None = None
1213
1214
@@ -1254,6 +1306,39 @@ def set_custom_all_reduce(enable: bool):
1306
_ENABLE_CUSTOM_ALL_REDUCE = enable
1307
1308
1309
+def _init_elastic_ep_world(
1310
+ config, local_rank: int, backend: str, rank: int, world_size: int
1311
+) -> None:
1312
+ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
1313
+
1314
+ global _WORLD, _NODE_COUNT
1315
+ assert _WORLD is None, "world group already initialized"
1316
+ parallel_config = config.parallel_config
1317
+ global_rank = parallel_config.data_parallel_rank * world_size + rank
1318
+ global_world_size = parallel_config.world_size_across_dp
1319
+ all_ranks = list(range(global_world_size))
1320
+ group_ranks = [all_ranks[i : i + 1] for i in range(global_world_size)]
1321
+ if global_rank in all_ranks:
1322
+ group_ranks = [all_ranks]
1323
+ group_ports = [parallel_config.get_next_stateless_world_group_port()]
1324
+ world = StatelessGroupCoordinator(
1325
+ group_ranks=group_ranks,
1326
+ local_rank=local_rank,
1327
+ torch_distributed_backend=backend,
1328
+ use_device_communicator=False,
1329
+ group_name="world",
1330
+ host=parallel_config.data_parallel_master_ip,
1331
+ group_ports=group_ports,
1332
+ global_rank=global_rank,
1333
+ global_world_size=global_world_size,
1334
+ )
1335
+ assert parallel_config.nnodes_within_dp == 1, (
1336
+ "Elastic EP is not supported with multi-node TP/PP"
1337
+ )
1338
+ _NODE_COUNT = _node_count(world.tcp_store_group)
1339
+ _WORLD = world
1340
+
1341
+
1342
def init_distributed_environment(
1343
world_size: int = -1,
1344
rank: int = -1,
@@ -1273,6 +1358,7 @@ def init_distributed_environment(
1358
from vllm.config import get_current_vllm_config_or_none
1359
1360
config = get_current_vllm_config_or_none()
1361
+ enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep
1362
if (
1363
config is not None
1364
and config.parallel_config.distributed_executor_backend != "external_launcher"
@@ -1280,6 +1366,7 @@ def init_distributed_environment(
1366
config.parallel_config.nnodes > 1
1367
or config.parallel_config.data_parallel_size > 1
1368
)
1369
+ and not enable_elastic_ep
1370
):
1371
parallel_config = config.parallel_config
1372
# adjust to take into account data parallelism
@@ -1333,6 +1420,18 @@ def init_distributed_environment(
1420
rank=rank,
1421
timeout=timeout,
1422
)
1423
+ if enable_elastic_ep:
1424
+ tp_pp_cpu_group = torch.distributed.new_group(
1425
+ backend="gloo", timeout=timeout
1426
+ )
1427
+ if _node_count(tp_pp_cpu_group) > 1:
1428
+ # NOTE(yongji): StatelessGroupCoordinator uses data_parallel_master_ip
1429
+ # to initialize all DP/EP groups, hence all ranks within TP/PP group
1430
+ # must reside on the same node
1431
+ raise RuntimeError(
1432
+ "Elastic EP is not yet supported with multi-node TP/PP"
1433
+ )
1434
+
1435
# set the local rank
1436
# local_rank is not available in torch ProcessGroup,
1437
# see https://github.com/pytorch/pytorch/issues/122816
@@ -1341,6 +1440,9 @@ def init_distributed_environment(
1440
# setting, where we can use rank as local rank
1441
local_rank = envs.LOCAL_RANK if distributed_init_method == "env://" else rank
1442
global _WORLD, _NODE_COUNT, _INNER_DP_WORLD
1443
+ if enable_elastic_ep:
1444
+ _init_elastic_ep_world(config, local_rank, backend, rank, world_size)
1445
+ return
1446
if _WORLD is None:
1447
ranks = list(range(torch.distributed.get_world_size()))
1448
_WORLD = init_world_group(ranks, local_rank, backend)
@@ -1404,16 +1506,33 @@ def initialize_model_parallel(
1506
"""
1507
# Get world size and rank. Ensure some consistencies.
1508
assert torch.distributed.is_initialized()
1407
- world_size: int = torch.distributed.get_world_size()
1408
- rank = torch.distributed.get_rank()
1409
- backend = backend or torch.distributed.get_backend(get_world_group().device_group)
1410
-
1411
- data_parallel_size = 1
1412
- from vllm.config import get_current_vllm_config_or_none
1509
1414
- config = get_current_vllm_config_or_none()
1415
- if config is not None:
1416
- data_parallel_size = config.parallel_config.data_parallel_size
1510
+ from vllm.config import get_current_vllm_config
1511
+
1512
+ config = get_current_vllm_config()
1513
+ data_parallel_size = config.parallel_config.data_parallel_size
1514
+ enable_elastic_ep = config.parallel_config.enable_elastic_ep
1515
+ if enable_elastic_ep:
1516
+ # Use stateless world group for global information
1517
+ world_size = get_world_group().world_size
1518
+ rank = get_world_group().rank
1519
+ backend = backend or "nccl"
1520
+ tp_pp_pcp_size = (
1521
+ tensor_model_parallel_size
1522
+ * pipeline_model_parallel_size
1523
+ * prefill_context_model_parallel_size
1524
+ )
1525
+ local_all_ranks = torch.arange(tp_pp_pcp_size).reshape(
1526
+ pipeline_model_parallel_size,
1527
+ prefill_context_model_parallel_size,
1528
+ tensor_model_parallel_size,
1529
+ )
1530
+ else:
1531
+ world_size = torch.distributed.get_world_size()
1532
+ rank = torch.distributed.get_rank()
1533
+ backend = backend or torch.distributed.get_backend(
1534
+ get_world_group().device_group
1535
+ )
1536
1537
# the layout order is: ExternalDP x DP x PP x TP
1538
# ExternalDP is the data parallel group that is not part of the model,
@@ -1437,7 +1556,9 @@ def initialize_model_parallel(
1556
assert _TP is None, "tensor model parallel group is already initialized"
1557
group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
1558
group_ranks = [x.tolist() for x in group_ranks]
1440
-
1559
+ if enable_elastic_ep:
1560
+ group_ranks = local_all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
1561
+ group_ranks = [x.tolist() for x in group_ranks]
1562
# message queue broadcaster is only used in tensor model parallel group
1563
_TP = init_model_parallel_group(
1564
group_ranks,
@@ -1456,6 +1577,11 @@ def initialize_model_parallel(
1577
# TP group into tp_size//dcp_size DCP groups.
1578
group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0)
1579
group_ranks = [x.tolist() for x in group_ranks]
1580
+ if enable_elastic_ep:
1581
+ group_ranks = local_all_ranks.reshape(
1582
+ -1, decode_context_model_parallel_size
1583
+ ).unbind(0)
1584
+ group_ranks = [x.tolist() for x in group_ranks]
1585
_DCP = init_model_parallel_group(
1586
group_ranks,
1587
get_world_group().local_rank,
@@ -1472,6 +1598,13 @@ def initialize_model_parallel(
1598
.unbind(0)
1599
)
1600
group_ranks = [x.tolist() for x in group_ranks]
1601
+ if enable_elastic_ep:
1602
+ group_ranks = (
1603
+ local_all_ranks.transpose(1, 2)
1604
+ .reshape(-1, prefill_context_model_parallel_size)
1605
+ .unbind(0)
1606
+ )
1607
+ group_ranks = [x.tolist() for x in group_ranks]
1608
_PCP = init_model_parallel_group(
1609
group_ranks, get_world_group().local_rank, backend, group_name="pcp"
1610
)
@@ -1483,6 +1616,13 @@ def initialize_model_parallel(
1616
all_ranks.transpose(2, 4).reshape(-1, pipeline_model_parallel_size).unbind(0)
1617
)
1618
group_ranks = [x.tolist() for x in group_ranks]
1619
+ if enable_elastic_ep:
1620
+ group_ranks = (
1621
+ local_all_ranks.transpose(0, 2)
1622
+ .reshape(-1, pipeline_model_parallel_size)
1623
+ .unbind(0)
1624
+ )
1625
+ group_ranks = [x.tolist() for x in group_ranks]
1626
_PP = init_model_parallel_group(
1627
group_ranks, get_world_group().local_rank, backend, group_name="pp"
1628
)
@@ -1491,14 +1631,27 @@ def initialize_model_parallel(
1631
assert _DP is None, "data parallel group is already initialized"
1632
group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
1633
group_ranks = [x.tolist() for x in group_ranks]
1494
- _DP = init_model_parallel_group(
1495
- group_ranks, get_world_group().local_rank, backend, group_name="dp"
1496
- )
1634
+ if enable_elastic_ep:
1635
+ parallel_config = config.parallel_config
1636
+ dp_ports = [
1637
+ parallel_config.get_next_stateless_dp_group_port() for _ in group_ranks
1638
+ ]
1639
+ _DP = _init_stateless_group(
1640
+ group_ranks,
1641
+ "dp",
1642
+ dp_ports,
1643
+ parallel_config.data_parallel_master_ip,
1644
+ backend,
1645
+ )
1646
+ else:
1647
+ _DP = init_model_parallel_group(
1648
+ group_ranks, get_world_group().local_rank, backend, group_name="dp"
1649
+ )
1650
1651
global _EP
1652
assert _EP is None, "expert parallel group is already initialized"
1653
# Don't create EP group for dense models.
1501
- if config is None or config.model_config is None or config.model_config.is_moe:
1654
+ if config.model_config is None or config.model_config.is_moe:
1655
group_ranks = (
1656
all_ranks.transpose(1, 2)
1657
.reshape(
@@ -1510,9 +1663,22 @@ def initialize_model_parallel(
1663
.unbind(0)
1664
)
1665
group_ranks = [x.tolist() for x in group_ranks]
1513
- _EP = init_model_parallel_group(
1514
- group_ranks, get_world_group().local_rank, backend, group_name="ep"
1515
- )
1666
+ if enable_elastic_ep:
1667
+ parallel_config = config.parallel_config
1668
+ ep_ports = [
1669
+ parallel_config.get_next_stateless_ep_group_port() for _ in group_ranks
1670
+ ]
1671
+ _EP = _init_stateless_group(
1672
+ group_ranks,
1673
+ "ep",
1674
+ ep_ports,
1675
+ parallel_config.data_parallel_master_ip,
1676
+ backend,
1677
+ )
1678
+ else:
1679
+ _EP = init_model_parallel_group(
1680
+ group_ranks, get_world_group().local_rank, backend, group_name="ep"
1681
+ )
1682
1683
# Create EPLB group with the same ranks as EP if EPLB is enabled.
1684
# This is a separate process group to isolate EPLB communications
@@ -1525,10 +1691,25 @@ def initialize_model_parallel(
1691
and config.parallel_config is not None
1692
and config.parallel_config.enable_eplb
1693
):
1528
- # Reuse the same group_ranks from EP
1529
- _EPLB = init_model_parallel_group(
1530
- group_ranks, get_world_group().local_rank, backend, group_name="eplb"
1531
- )
1694
+ if enable_elastic_ep:
1695
+ eplb_ports = [
1696
+ parallel_config.get_next_stateless_eplb_group_port()
1697
+ for _ in group_ranks
1698
+ ]
1699
+ _EPLB = _init_stateless_group(
1700
+ group_ranks,
1701
+ "eplb",
1702
+ eplb_ports,
1703
+ parallel_config.data_parallel_master_ip,
1704
+ backend,
1705
+ )
1706
+ else:
1707
+ _EPLB = init_model_parallel_group(
1708
+ group_ranks,
1709
+ get_world_group().local_rank,
1710
+ backend,
1711
+ group_name="eplb",
1712
+ )
1713
# If no EP group needed, _EP remains None
1714
# If no EPLB group needed, _EPLB remains None
1715
@@ -1558,7 +1739,11 @@ def ensure_model_parallel_initialized(
1739
or ensure tensor-parallel and pipeline-parallel sizes are equal to expected
1740
values if the model parallel groups are initialized.
1741
"""
1561
- backend = backend or torch.distributed.get_backend(get_world_group().device_group)
1742
+ world_group = get_world_group()
1743
+ if hasattr(world_group, "backend"):
1744
+ backend = backend or world_group.backend
1745
+ else:
1746
+ backend = backend or torch.distributed.get_backend(world_group.device_group)
1747
if not model_parallel_is_initialized():
1748
initialize_model_parallel(
1749
tensor_model_parallel_size,
vllm/distributed/stateless_coordinator.py
new
+322
@@ -0,0 +1,322 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any, Optional
4
+
5
+import torch
6
+from torch.distributed import Backend, ProcessGroup
7
+
8
+from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
9
+from vllm.distributed.parallel_state import (
10
+ GroupCoordinator,
11
+ TensorMetadata,
12
+ _get_unique_name,
13
+ _register_group,
14
+ _split_tensor_dict,
15
+)
16
+from vllm.distributed.utils import (
17
+ StatelessProcessGroup,
18
+ stateless_destroy_torch_distributed_process_group,
19
+ stateless_init_torch_distributed_process_group,
20
+)
21
+from vllm.logger import init_logger
22
+from vllm.utils.import_utils import resolve_obj_by_qualname
23
+
24
+logger = init_logger(__name__)
25
+
26
+
27
+class StatelessGroupCoordinator(GroupCoordinator):
28
+ """
29
+ A stateless version of the GroupCoordinator class in parallel_state,
30
+ It will create CPU, device and TCPStore based communication groups
31
+ that are independent of PyTorch's WORLD group. Hence,
32
+ communication groups with a different set of participants GPUs
33
+ can be created without destroying the existing ones.
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ group_ranks: list[list[int]],
39
+ local_rank: int,
40
+ torch_distributed_backend: str | Backend,
41
+ use_device_communicator: bool,
42
+ use_message_queue_broadcaster: bool = False,
43
+ group_name: str | None = None,
44
+ host: str = "127.0.0.1",
45
+ group_ports: list[list[int]] | None = None,
46
+ global_rank: int = 0,
47
+ global_world_size: int = 1,
48
+ ):
49
+ group_name = group_name or "anonymous"
50
+ self.unique_name = _get_unique_name(group_name)
51
+ _register_group(self)
52
+
53
+ self.rank = global_rank
54
+ self.local_rank = local_rank
55
+
56
+ self_device_group = None
57
+ self_cpu_group = None
58
+ self_tcp_store_group = None
59
+
60
+ from vllm.platforms import current_platform
61
+
62
+ backend = str(torch_distributed_backend)
63
+ self.backend = backend
64
+ assert group_ports is not None, "group_ports is not provided"
65
+ for idx, ranks in enumerate(group_ranks):
66
+ if self.rank in ranks:
67
+ self.ranks = ranks
68
+ self.world_size = len(ranks)
69
+ self.rank_in_group = ranks.index(self.rank)
70
+
71
+ ports = group_ports[idx]
72
+ device_port = ports[0]
73
+ cpu_port = ports[1]
74
+ tcp_store_port = ports[2]
75
+
76
+ device_group = stateless_init_torch_distributed_process_group(
77
+ host=host,
78
+ port=device_port,
79
+ rank=self.rank_in_group,
80
+ world_size=self.world_size,
81
+ backend=backend,
82
+ group_name=f"{self.unique_name}_device",
83
+ )
84
+ cpu_group = stateless_init_torch_distributed_process_group(
85
+ host=host,
86
+ port=cpu_port,
87
+ rank=self.rank_in_group,
88
+ world_size=self.world_size,
89
+ backend="gloo",
90
+ group_name=f"{self.unique_name}_cpu",
91
+ )
92
+ tcp_store_group = StatelessProcessGroup.create(
93
+ host=host,
94
+ port=tcp_store_port,
95
+ rank=self.rank_in_group,
96
+ world_size=self.world_size,
97
+ )
98
+
99
+ self_device_group = device_group
100
+ self_cpu_group = cpu_group
101
+ self_tcp_store_group = tcp_store_group
102
+
103
+ assert self_cpu_group is not None
104
+ assert self_device_group is not None
105
+ assert self_tcp_store_group is not None
106
+
107
+ self.cpu_group = self_cpu_group
108
+ self.device_group = self_device_group
109
+ self.tcp_store_group = self_tcp_store_group
110
+
111
+ if current_platform.is_cuda_alike():
112
+ self.device = torch.device(f"cuda:{local_rank}")
113
+ elif current_platform.is_xpu():
114
+ self.device = torch.device(f"xpu:{local_rank}")
115
+ elif current_platform.is_out_of_tree():
116
+ self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
117
+ else:
118
+ self.device = torch.device("cpu")
119
+
120
+ self.use_device_communicator = use_device_communicator
121
+ self.device_communicator = None
122
+ if use_device_communicator and self.world_size > 1:
123
+ device_comm_cls = resolve_obj_by_qualname(
124
+ current_platform.get_device_communicator_cls()
125
+ )
126
+ assert device_comm_cls == CudaCommunicator
127
+ self.device_communicator = CudaCommunicator(
128
+ cpu_group=self.cpu_group,
129
+ device=self.device,
130
+ device_group=self.device_group,
131
+ unique_name=self.unique_name,
132
+ global_ranks=self.ranks,
133
+ global_world_size=global_world_size,
134
+ tcp_store_group=self.tcp_store_group,
135
+ )
136
+
137
+ self.mq_broadcaster = None
138
+
139
+ self.use_custom_op_call = (
140
+ current_platform.is_cuda_alike() or current_platform.is_tpu()
141
+ )
142
+ self.use_cpu_custom_send_recv = False
143
+
144
+ def destroy(self):
145
+ if self.device_communicator:
146
+ self.device_communicator.destroy()
147
+ if self.device_group:
148
+ stateless_destroy_torch_distributed_process_group(self.device_group)
149
+ if self.cpu_group:
150
+ stateless_destroy_torch_distributed_process_group(self.cpu_group)
151
+
152
+ def size(self) -> int:
153
+ """Return the world size of this group."""
154
+ return self.world_size
155
+
156
+ def broadcast(self, input_: torch.Tensor, src: int = 0):
157
+ if self.world_size == 1:
158
+ return input_
159
+
160
+ if self.device_communicator and input_.is_cuda:
161
+ return self.device_communicator.broadcast(input_, src)
162
+ else:
163
+ return self.tcp_store_group.broadcast(input_, src)
164
+
165
+ def broadcast_object(self, obj=None, src: int = 0):
166
+ if self.world_size == 1:
167
+ return obj
168
+ return self.tcp_store_group.broadcast_obj(obj, src)
169
+
170
+ def broadcast_object_list(
171
+ self, obj_list: list[Any], src: int = 0, group: ProcessGroup | None = None
172
+ ):
173
+ assert src < self.world_size
174
+
175
+ if self.world_size == 1:
176
+ return obj_list
177
+
178
+ if self.rank_in_group == src:
179
+ for obj in obj_list:
180
+ self.tcp_store_group.broadcast_obj(obj, src)
181
+ else:
182
+ for i in range(len(obj_list)):
183
+ obj_list[i] = self.tcp_store_group.broadcast_obj(None, src)
184
+
185
+ return obj_list
186
+
187
+ def broadcast_tensor_dict(
188
+ self,
189
+ tensor_dict: dict[str, torch.Tensor | Any] | None = None,
190
+ src: int = 0,
191
+ group: ProcessGroup | None = None,
192
+ metadata_group: ProcessGroup | None = None,
193
+ ) -> dict[str, torch.Tensor | Any] | None:
194
+ if self.world_size == 1:
195
+ return tensor_dict
196
+
197
+ if self.rank_in_group == src:
198
+ assert isinstance(tensor_dict, dict), (
199
+ f"Expecting a dictionary, got {type(tensor_dict)}"
200
+ )
201
+ metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
202
+ else:
203
+ metadata_list = None
204
+ tensor_list = []
205
+
206
+ recv_metadata_list: list[tuple[str, Any]] = self.tcp_store_group.broadcast_obj(
207
+ metadata_list, src
208
+ )
209
+
210
+ if self.rank_in_group != src:
211
+ tensor_dict = {}
212
+ for key, value in recv_metadata_list:
213
+ if isinstance(value, TensorMetadata):
214
+ tensor = torch.empty(
215
+ value.size, dtype=value.dtype, device=value.device
216
+ )
217
+ tensor_list.append(tensor)
218
+ tensor_dict[key] = tensor
219
+ else:
220
+ tensor_dict[key] = value
221
+
222
+ for tensor in tensor_list:
223
+ if tensor.numel() == 0:
224
+ continue
225
+ if self.device_communicator and tensor.is_cuda:
226
+ tensor.copy_(self.device_communicator.broadcast(tensor, src))
227
+ else:
228
+ tensor.copy_(self.tcp_store_group.broadcast(tensor, src))
229
+
230
+ return tensor_dict
231
+
232
+ def send_object(self, obj, dst: int) -> None:
233
+ assert dst < self.world_size
234
+ assert dst != self.rank_in_group
235
+ self.tcp_store_group.send_obj(obj, dst)
236
+
237
+ def recv_object(self, src: int):
238
+ assert src < self.world_size
239
+ assert src != self.rank_in_group
240
+ return self.tcp_store_group.recv_obj(src)
241
+
242
+ def send_tensor_dict(
243
+ self,
244
+ tensor_dict: dict[str, torch.Tensor | Any],
245
+ dst: int | None = None,
246
+ all_gather_group: Optional["GroupCoordinator"] = None,
247
+ all_gather_tensors: dict[str, bool] | None = None,
248
+ ) -> dict[str, torch.Tensor | Any] | None:
249
+ if self.world_size == 1:
250
+ return tensor_dict
251
+
252
+ if dst is None:
253
+ dst = (self.rank_in_group + 1) % self.world_size
254
+ assert dst < self.world_size
255
+
256
+ metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
257
+ self.tcp_store_group.send_obj(metadata_list, dst)
258
+
259
+ for tensor in tensor_list:
260
+ if tensor.numel() == 0:
261
+ continue
262
+ if self.device_communicator and tensor.is_cuda:
263
+ self.device_communicator.send(tensor, dst)
264
+ else:
265
+ self.tcp_store_group.send(tensor, dst)
266
+
267
+ return None
268
+
269
+ def recv_tensor_dict(
270
+ self,
271
+ src: int | None = None,
272
+ all_gather_group: Optional["GroupCoordinator"] = None,
273
+ all_gather_tensors: dict[str, bool] | None = None,
274
+ ) -> dict[str, torch.Tensor | Any] | None:
275
+ if self.world_size == 1:
276
+ return None
277
+
278
+ if src is None:
279
+ src = (self.rank_in_group - 1) % self.world_size
280
+ assert src < self.world_size
281
+
282
+ recv_metadata_list = self.tcp_store_group.recv_obj(src)
283
+ tensor_dict = {}
284
+ for key, value in recv_metadata_list:
285
+ if isinstance(value, TensorMetadata):
286
+ tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)
287
+ if tensor.numel() > 0:
288
+ if self.device_communicator and tensor.is_cuda:
289
+ tensor = self.device_communicator.recv(
290
+ tensor.size(), tensor.dtype, src
291
+ )
292
+ else:
293
+ tensor = self.tcp_store_group.recv(tensor, src)
294
+ tensor_dict[key] = tensor
295
+ else:
296
+ tensor_dict[key] = value
297
+ return tensor_dict
298
+
299
+ def barrier(self):
300
+ self.tcp_store_group.barrier()
301
+
302
+ def gather(
303
+ self, input_: torch.Tensor, dst: int = 0, dim: int = -1
304
+ ) -> torch.Tensor | None:
305
+ if self.world_size == 1:
306
+ return input_
307
+
308
+ if self.device_communicator is None:
309
+ raise ValueError("No device communicator found")
310
+
311
+ if self.rank_in_group == dst:
312
+ gathered_list = [torch.empty_like(input_) for _ in range(self.world_size)]
313
+ gathered_list[self.rank_in_group] = input_
314
+ for src_rank in range(self.world_size):
315
+ if src_rank != self.rank_in_group:
316
+ gathered_list[src_rank] = self.device_communicator.recv(
317
+ input_.size(), input_.dtype, src_rank
318
+ )
319
+ return torch.cat(gathered_list, dim=dim)
320
+ else:
321
+ self.device_communicator.send(input_, dst)
322
+ return None
vllm/distributed/utils.py
+76
-11
@@ -18,7 +18,7 @@ from datetime import timedelta
18
from typing import Any
19
20
import torch
21
-from torch.distributed import ProcessGroup, TCPStore
21
+from torch.distributed import ProcessGroup, Store, TCPStore
22
from torch.distributed.distributed_c10d import (
23
Backend,
24
PrefixStore,
@@ -228,6 +228,55 @@ class StatelessProcessGroup:
228
gathered_objs.append(recv_obj)
229
return gathered_objs
230
231
+ def broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
232
+ """Broadcast a tensor from source rank to all other ranks."""
233
+ if self.rank == src:
234
+ tensor_bytes = pickle.dumps(tensor)
235
+ self.expire_data()
236
+ key = f"broadcast_tensor/{src}/{self.broadcast_send_counter}"
237
+ self.store.set(key, tensor_bytes)
238
+ self.broadcast_send_counter += 1
239
+ self.entries.append((key, time.time()))
240
+ return tensor
241
+ else:
242
+ key = f"broadcast_tensor/{src}/{self.broadcast_recv_src_counter[src]}"
243
+ tensor = pickle.loads(self.store.get(key))
244
+ self.broadcast_recv_src_counter[src] += 1
245
+ return tensor
246
+
247
+ def send(self, tensor: torch.Tensor, dst: int):
248
+ """Send a tensor to a destination rank."""
249
+ self.expire_data()
250
+ key = f"send_tensor/{dst}/{self.send_dst_counter[dst]}"
251
+ self.store.set(key, pickle.dumps(tensor))
252
+ self.send_dst_counter[dst] += 1
253
+ self.entries.append((key, time.time()))
254
+
255
+ def recv(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
256
+ """Receive a tensor from a source rank."""
257
+ key = f"send_tensor/{self.rank}/{self.recv_src_counter[src]}"
258
+ received = pickle.loads(self.store.get(key))
259
+ self.recv_src_counter[src] += 1
260
+ tensor.copy_(received)
261
+ return tensor
262
+
263
+ def all_reduce(
264
+ self, tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM
265
+ ) -> torch.Tensor:
266
+ """All-reduce a tensor across all ranks."""
267
+ tensors = self.all_gather_obj(tensor)
268
+ result = tensors[0].clone()
269
+ for t in tensors[1:]:
270
+ if op == torch.distributed.ReduceOp.SUM:
271
+ result.add_(t)
272
+ elif op == torch.distributed.ReduceOp.PRODUCT:
273
+ result.mul_(t)
274
+ elif op == torch.distributed.ReduceOp.MAX:
275
+ result = torch.maximum(result, t)
276
+ elif op == torch.distributed.ReduceOp.MIN:
277
+ result = torch.minimum(result, t)
278
+ return result
279
+
280
def barrier(self, timeout: float = 30.0):
281
"""A robust barrier to synchronize all ranks.
282
@@ -448,8 +497,14 @@ def init_gloo_process_group(
497
498
499
def stateless_init_torch_distributed_process_group(
451
- host: str, port: int, rank: int, world_size: int, backend: str
452
-) -> ProcessGroup:
500
+ host: str,
501
+ port: int,
502
+ rank: int,
503
+ world_size: int,
504
+ backend: str,
505
+ group_name: str | None = None,
506
+ return_store: bool = False,
507
+) -> ProcessGroup | tuple[ProcessGroup, Store]:
508
"""
509
A replacement for `torch.distributed.init_process_group` that does not
510
pollute the global state. The created ProcessGroup object can be used for
@@ -496,26 +551,36 @@ def stateless_init_torch_distributed_process_group(
551
# Use a PrefixStore to avoid accidental overrides of keys used by
552
# different systems (e.g. RPC) in case the store is multi-tenant.
553
prefix_store = PrefixStore(init_method, store)
499
- try:
500
- from vllm.platforms import current_platform
554
502
- return current_platform.stateless_init_device_torch_dist_pg(
503
- backend=backend,
555
+ if backend == "gloo":
556
+ pg = init_gloo_process_group(
557
prefix_store=prefix_store,
558
group_rank=group_rank,
559
group_size=group_size,
560
timeout=timeout,
561
)
509
- except NotImplementedError:
510
- # If platform doesn't implement stateless_init_device_torch_dist_pg, it
511
- # will raise a NotImplementedError. In this case, we fall back to gloo.
512
- return init_gloo_process_group(
562
+ else:
563
+ from vllm.platforms import current_platform
564
+
565
+ pg = current_platform.stateless_init_device_torch_dist_pg(
566
+ backend=backend,
567
prefix_store=prefix_store,
568
group_rank=group_rank,
569
group_size=group_size,
570
timeout=timeout,
571
)
572
573
+ if group_name is not None:
574
+ from torch._C._distributed_c10d import _register_process_group
575
+
576
+ pg._set_group_name(group_name)
577
+ _register_process_group(group_name, pg)
578
+
579
+ if return_store:
580
+ return pg, store
581
+ else:
582
+ return pg
583
+
584
585
def stateless_destroy_torch_distributed_process_group(pg: ProcessGroup) -> None:
586
"""
vllm/engine/arg_utils.py
+5
@@ -419,6 +419,7 @@ class EngineArgs:
419
enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
420
moe_backend: MoEBackend = KernelConfig.moe_backend
421
all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
422
+ enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
423
enable_dbo: bool = ParallelConfig.enable_dbo
424
ubatch_size: int = ParallelConfig.ubatch_size
425
dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold
@@ -896,6 +897,9 @@ class EngineArgs:
897
"--ubatch-size",
898
**parallel_kwargs["ubatch_size"],
899
)
900
+ parallel_group.add_argument(
901
+ "--enable-elastic-ep", **parallel_kwargs["enable_elastic_ep"]
902
+ )
903
parallel_group.add_argument(
904
"--dbo-decode-token-threshold",
905
**parallel_kwargs["dbo_decode_token_threshold"],
@@ -1698,6 +1702,7 @@ class EngineArgs:
1702
is_moe_model=model_config.is_moe,
1703
enable_expert_parallel=self.enable_expert_parallel,
1704
all2all_backend=self.all2all_backend,
1705
+ enable_elastic_ep=self.enable_elastic_ep,
1706
enable_dbo=self.enable_dbo,
1707
ubatch_size=self.ubatch_size,
1708
dbo_decode_token_threshold=self.dbo_decode_token_threshold,
vllm/entrypoints/cli/serve.py
+5
-1
@@ -246,8 +246,12 @@ def run_multi_api_server(args: argparse.Namespace):
246
247
api_server_manager: APIServerProcessManager | None = None
248
249
+ from vllm.v1.engine.utils import get_engine_zmq_addresses
250
+
251
+ addresses = get_engine_zmq_addresses(vllm_config, num_api_servers)
252
+
253
with launch_core_engines(
250
- vllm_config, executor_class, log_stats, num_api_servers
254
+ vllm_config, executor_class, log_stats, addresses, num_api_servers
255
) as (local_engine_manager, coordinator, addresses):
256
# Construct common args for the APIServerProcessManager up-front.
257
api_server_manager_kwargs = dict(
vllm/envs.py
+12
@@ -243,6 +243,8 @@ if TYPE_CHECKING:
243
VLLM_LORA_DISABLE_PDL: bool = False
244
VLLM_ENABLE_CUDA_COMPATIBILITY: bool = False
245
VLLM_CUDA_COMPATIBILITY_PATH: str | None = None
246
+ VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False
247
+ VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False
248
249
250
def get_default_cache_root():
@@ -1617,6 +1619,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
1619
"VLLM_CUDA_COMPATIBILITY_PATH": lambda: os.environ.get(
1620
"VLLM_CUDA_COMPATIBILITY_PATH", None
1621
),
1622
+ # Whether it is a scale up launch engine for elastic EP,
1623
+ # Should only be set by EngineCoreClient.
1624
+ "VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": lambda: bool(
1625
+ int(os.getenv("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH", "0"))
1626
+ ),
1627
+ # Whether to wait for all requests to drain before sending the
1628
+ # scaling command in elastic EP.
1629
+ "VLLM_ELASTIC_EP_DRAIN_REQUESTS": lambda: bool(
1630
+ int(os.getenv("VLLM_ELASTIC_EP_DRAIN_REQUESTS", "0"))
1631
+ ),
1632
}
1633
1634
vllm/model_executor/layers/fused_moe/layer.py
+3
-2
@@ -627,6 +627,7 @@ class FusedMoE(CustomOp):
627
moe_quant_params["intermediate_size_full"] = intermediate_size
628
629
self.quant_method.create_weights(layer=self, **moe_quant_params)
630
+ self.base_quant_method = self.quant_method
631
632
# Disable shared expert overlap if:
633
# - we are using eplb with non-default backend, because of correctness issues
@@ -683,7 +684,7 @@ class FusedMoE(CustomOp):
684
# routing_tables only needed for round-robin expert placement with
685
# DeepEP all2all backend.
686
routing_tables = self._maybe_init_expert_routing_tables()
686
- prepare_finalize = self.quant_method.maybe_make_prepare_finalize(
687
+ prepare_finalize = self.base_quant_method.maybe_make_prepare_finalize(
688
routing_tables=routing_tables
689
)
690
if prepare_finalize is not None:
@@ -693,7 +694,7 @@ class FusedMoE(CustomOp):
694
self._replace_quant_method(
695
FusedMoEModularMethod.make(
696
self,
696
- self.quant_method,
697
+ self.base_quant_method,
698
prepare_finalize,
699
self.shared_experts,
700
inplace=not self.moe_config.disable_inplace,
vllm/platforms/cuda.py
+34
@@ -6,10 +6,13 @@ pynvml. However, it should not initialize cuda context.
6
7
import os
8
from collections.abc import Callable
9
+from datetime import timedelta
10
from functools import cache, wraps
11
from typing import TYPE_CHECKING, TypeVar
12
13
import torch
14
+from torch.distributed import PrefixStore, ProcessGroup
15
+from torch.distributed.distributed_c10d import is_nccl_available
16
from typing_extensions import ParamSpec
17
18
# import custom ops, trigger op registration
@@ -482,6 +485,37 @@ class CudaPlatformBase(Platform):
485
def get_static_graph_wrapper_cls(cls) -> str:
486
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
487
488
+ @classmethod
489
+ def stateless_init_device_torch_dist_pg(
490
+ cls,
491
+ backend: str,
492
+ prefix_store: PrefixStore,
493
+ group_rank: int,
494
+ group_size: int,
495
+ timeout: timedelta,
496
+ ) -> ProcessGroup:
497
+ assert is_nccl_available()
498
+ pg: ProcessGroup = ProcessGroup(
499
+ prefix_store,
500
+ group_rank,
501
+ group_size,
502
+ )
503
+ from torch.distributed.distributed_c10d import ProcessGroupNCCL
504
+
505
+ backend_options = ProcessGroupNCCL.Options()
506
+ backend_options._timeout = timeout
507
+
508
+ backend_class = ProcessGroupNCCL(
509
+ prefix_store, group_rank, group_size, backend_options
510
+ )
511
+ backend_type = ProcessGroup.BackendType.NCCL
512
+ device = torch.device("cuda")
513
+ pg._set_default_backend(backend_type)
514
+ backend_class._set_sequence_number_for_group()
515
+
516
+ pg._register_backend(device, backend_type, backend_class)
517
+ return pg
518
+
519
@classmethod
520
def device_count(cls) -> int:
521
return cuda_device_count_stateless()
vllm/platforms/rocm.py
+34
@@ -2,10 +2,13 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import os
5
+from datetime import timedelta
6
from functools import cache, lru_cache, wraps
7
from typing import TYPE_CHECKING
8
9
import torch
10
+from torch.distributed import PrefixStore, ProcessGroup
11
+from torch.distributed.distributed_c10d import is_nccl_available
12
13
import vllm.envs as envs
14
from vllm.logger import init_logger
@@ -656,6 +659,37 @@ class RocmPlatform(Platform):
659
def get_static_graph_wrapper_cls(cls) -> str:
660
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
661
662
+ @classmethod
663
+ def stateless_init_device_torch_dist_pg(
664
+ cls,
665
+ backend: str,
666
+ prefix_store: PrefixStore,
667
+ group_rank: int,
668
+ group_size: int,
669
+ timeout: timedelta,
670
+ ) -> ProcessGroup:
671
+ assert is_nccl_available()
672
+ pg: ProcessGroup = ProcessGroup(
673
+ prefix_store,
674
+ group_rank,
675
+ group_size,
676
+ )
677
+ from torch.distributed.distributed_c10d import ProcessGroupNCCL
678
+
679
+ backend_options = ProcessGroupNCCL.Options()
680
+ backend_options._timeout = timeout
681
+
682
+ backend_class = ProcessGroupNCCL(
683
+ prefix_store, group_rank, group_size, backend_options
684
+ )
685
+ backend_type = ProcessGroup.BackendType.NCCL
686
+ device = torch.device("cuda")
687
+ pg._set_default_backend(backend_type)
688
+ backend_class._set_sequence_number_for_group()
689
+
690
+ pg._register_backend(device, backend_type, backend_class)
691
+ return pg
692
+
693
@classmethod
694
def device_count(cls) -> int:
695
return cuda_device_count_stateless()
vllm/v1/engine/__init__.py
+14
@@ -29,6 +29,15 @@ PauseMode = Literal["abort", "wait", "keep"]
29
# so form part of the external API.
30
FINISH_REASON_STRINGS = ("stop", "length", "abort", "error")
31
32
+EEP_NOTIFICATION_CALL_ID = -1
33
+
34
+
35
+class EEPNotificationType(enum.Enum):
36
+ NEW_CORE_ENGINES_INIT_READY = "NEW_CORE_ENGINES_INIT_READY"
37
+ NEW_CORE_ENGINES_WEIGHTS_INIT_READY = "NEW_CORE_ENGINES_WEIGHTS_INIT_READY"
38
+ RECONFIGURE_FINISHED = "RECONFIGURE_FINISHED"
39
+ SHUTDOWN_COMPLETE = "SHUTDOWN_COMPLETE"
40
+
41
42
class FinishReason(enum.IntEnum):
43
"""
@@ -235,6 +244,11 @@ class ReconfigureDistributedRequest(msgspec.Struct):
244
new_data_parallel_rank_local: int
245
new_data_parallel_master_ip: str
246
new_data_parallel_master_port: int
247
+ new_data_parallel_master_port_list: list[int]
248
+ new_stateless_world_group_port_list: list[list[int]]
249
+ new_stateless_dp_group_port_list: list[list[int]]
250
+ new_stateless_ep_group_port_list: list[list[int]]
251
+ new_stateless_eplb_group_port_list: list[list[int]]
252
253
254
class ReconfigureRankType(enum.IntEnum):
vllm/v1/engine/async_llm.py
+27
-14
@@ -20,6 +20,7 @@ from vllm.distributed.weight_transfer.base import (
20
)
21
from vllm.engine.arg_utils import AsyncEngineArgs
22
from vllm.engine.protocol import EngineClient, StreamingInput
23
+from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep
24
from vllm.inputs import ProcessorInputs, PromptType
25
from vllm.logger import init_logger
26
from vllm.lora.request import LoRARequest
@@ -647,7 +648,11 @@ class AsyncLLM(EngineClient):
648
engine_core = self.engine_core
649
output_processor = self.output_processor
650
log_stats = self.log_stats
650
- logger_manager = self.logger_manager
651
+ # We use a mutable list for logger_manager so that it can be updated
652
+ # during elastic EP scaling (see scale_elastic_ep) without creating
653
+ # a circular reference via self.
654
+ self._logger_ref = [self.logger_manager]
655
+ logger_ref = self._logger_ref
656
renderer = self.renderer
657
chunk_size = envs.VLLM_V1_OUTPUT_PROC_CHUNK_SIZE
658
@@ -691,8 +696,8 @@ class AsyncLLM(EngineClient):
696
# 4) Logging.
697
# TODO(rob): make into a coroutine and launch it in
698
# background thread once Prometheus overhead is non-trivial.
694
- if logger_manager:
695
- logger_manager.record(
699
+ if logger_ref[0]:
700
+ logger_ref[0].record(
701
engine_idx=outputs.engine_index,
702
scheduler_stats=outputs.scheduler_stats,
703
iteration_stats=iteration_stats,
@@ -976,17 +981,13 @@ class AsyncLLM(EngineClient):
981
new_data_parallel_size,
982
)
983
return
979
- logger.info(
980
- "Waiting for requests to drain before scaling up to %s engines...",
981
- new_data_parallel_size,
982
- )
983
- await self.wait_for_requests_to_drain(drain_timeout)
984
- logger.info(
985
- "Requests have been drained, proceeding with scale to %s engines",
986
- new_data_parallel_size,
987
- )
988
- await self.engine_core.scale_elastic_ep(new_data_parallel_size)
989
- self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
984
+
985
+ if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS:
986
+ logger.info(
987
+ "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, "
988
+ "waiting for requests to drain before scaling"
989
+ )
990
+ await self.wait_for_requests_to_drain(drain_timeout)
991
992
# recreate stat loggers
993
if new_data_parallel_size > old_data_parallel_size and self.log_stats:
@@ -999,6 +1000,18 @@ class AsyncLLM(EngineClient):
1000
engine_idxs=list(range(new_data_parallel_size)),
1001
custom_stat_loggers=None,
1002
)
1003
+ # Update the mutable ref so output_handler picks up the
1004
+ # new logger without creating a circular reference via self.
1005
+ if hasattr(self, "_logger_ref"):
1006
+ self._logger_ref[0] = self.logger_manager
1007
+ self.logger_manager.log_engine_initialized()
1008
+
1009
+ set_scaling_elastic_ep(True)
1010
+ try:
1011
+ await self.engine_core.scale_elastic_ep(new_data_parallel_size)
1012
+ self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1013
+ finally:
1014
+ set_scaling_elastic_ep(False)
1015
1016
@property
1017
def is_running(self) -> bool:
vllm/v1/engine/coordinator.py
+20
-1
@@ -71,6 +71,9 @@ class DPCoordinator:
71
)
72
73
local_only_eng = dp_size == parallel_config.data_parallel_size_local
74
+ # NOTE(yongji): handling scaling from intra-node to inter-node
75
+ if parallel_config.enable_elastic_ep:
76
+ local_only_eng = False
77
back_publish_address = get_engine_client_zmq_addr(local_only_eng, host)
78
back_output_address = get_engine_client_zmq_addr(local_only_eng, host)
79
@@ -201,6 +204,7 @@ class DPCoordinatorProc:
204
205
poller = zmq.Poller()
206
poller.register(publish_front, zmq.POLLIN)
207
+ poller.register(publish_back, zmq.POLLIN)
208
poller.register(output_back, zmq.POLLIN)
209
last_publish_time = 0
210
while True:
@@ -231,6 +235,22 @@ class DPCoordinatorProc:
235
events = dict(events)
236
wave_state_changed = False
237
238
+ if publish_back in events:
239
+ buffer = publish_back.recv()
240
+ if buffer == b"\x01":
241
+ # NOTE(yongji): newly started engine subscribed
242
+ # We need to send READY message here instead of receiving
243
+ # SCALE_ELASTIC_EP notification from engine core client
244
+ # as SCALE_ELASTIC_EP is only sent when
245
+ # new engines finished initialization.
246
+ # Subscription message, on the other hand, is sent
247
+ # by each engine during initialization
248
+ publish_back.send(b"READY")
249
+ else:
250
+ logger.error(
251
+ "DP Coordinator receives unexpected message from engines"
252
+ )
253
+
254
if publish_front in events:
255
buffer = publish_front.recv()
256
if buffer in (b"\x01", b"\x00"):
@@ -259,7 +279,6 @@ class DPCoordinatorProc:
279
# current_wave
280
# we note that 0 is the wave number for the new
281
# engine
262
- engines_running = False
282
logger.info(
283
"DPCoordinator scaled up from %s to %s engines",
284
current_count,
vllm/v1/engine/core.py
+167
-48
@@ -17,6 +17,7 @@ from typing import Any, TypeVar, cast
17
import msgspec
18
import zmq
19
20
+import vllm.envs as envs
21
from vllm.config import ParallelConfig, VllmConfig
22
from vllm.distributed import stateless_destroy_torch_distributed_process_group
23
from vllm.envs import enable_envs_cache
@@ -44,6 +45,8 @@ from vllm.v1.core.kv_cache_utils import (
45
from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
46
from vllm.v1.core.sched.output import SchedulerOutput
47
from vllm.v1.engine import (
48
+ EEP_NOTIFICATION_CALL_ID,
49
+ EEPNotificationType,
50
EngineCoreOutput,
51
EngineCoreOutputs,
52
EngineCoreRequest,
@@ -110,6 +113,9 @@ class EngineCore:
113
114
self.available_gpu_memory_for_kv_cache = -1
115
116
+ if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
117
+ self._eep_scale_up_before_kv_init()
118
+
119
# Setup KV Caches and update CacheConfig after profiling.
120
num_gpu_blocks, num_cpu_blocks, kv_cache_config = self._initialize_kv_caches(
121
vllm_config
@@ -233,12 +239,10 @@ class EngineCore:
239
240
has_kv_cache = any(kv_cache_spec for kv_cache_spec in kv_cache_specs)
241
if has_kv_cache:
236
- if os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1":
237
- dp_group = getattr(self, "dp_group", None)
238
- assert dp_group is not None
239
- self.available_gpu_memory_for_kv_cache = (
240
- ParallelConfig.sync_kv_cache_memory_size(dp_group, -1)
241
- )
242
+ if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
243
+ # NOTE(yongji): should already be set
244
+ # during _eep_scale_up_before_kv_init
245
+ assert self.available_gpu_memory_for_kv_cache > 0
246
available_gpu_memory = [self.available_gpu_memory_for_kv_cache] * len(
247
kv_cache_specs
248
)
@@ -752,11 +756,22 @@ class EngineCore:
756
self.structured_output_manager.grammar_init(req)
757
return req, request.current_wave
758
759
+ def _eep_scale_up_before_kv_init(self):
760
+ raise NotImplementedError
761
+
762
+ def _eep_send_engine_core_notification(
763
+ self,
764
+ notification_type: EEPNotificationType,
765
+ vllm_config: VllmConfig | None = None,
766
+ ):
767
+ raise NotImplementedError
768
+
769
770
class EngineCoreProc(EngineCore):
771
"""ZMQ-wrapper for running EngineCore in background process."""
772
773
ENGINE_CORE_DEAD = b"ENGINE_CORE_DEAD"
774
+ addresses: EngineZmqAddresses
775
776
@instrument(span_name="EngineCoreProc init")
777
def __init__(
@@ -807,6 +822,13 @@ class EngineCoreProc(EngineCore):
822
# and "hybrid" LB modes.
823
self.publish_dp_lb_stats = internal_dp_balancing
824
825
+ self.addresses = addresses
826
+ self.process_input_queue_block = True
827
+ if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
828
+ self._eep_send_engine_core_notification(
829
+ EEPNotificationType.NEW_CORE_ENGINES_INIT_READY,
830
+ vllm_config=vllm_config,
831
+ )
832
self._init_data_parallel(vllm_config)
833
834
super().__init__(
@@ -1119,8 +1141,14 @@ class EngineCoreProc(EngineCore):
1141
if logger.isEnabledFor(DEBUG):
1142
logger.debug("EngineCore waiting for work.")
1143
waited = True
1122
- req = self.input_queue.get()
1123
- self._handle_client_request(*req)
1144
+ block = self.process_input_queue_block
1145
+ try:
1146
+ req = self.input_queue.get(block=block)
1147
+ self._handle_client_request(*req)
1148
+ except queue.Empty:
1149
+ break
1150
+ if not block:
1151
+ break
1152
1153
if waited:
1154
logger.debug("EngineCore loop active.")
@@ -1290,6 +1318,11 @@ class EngineCoreProc(EngineCore):
1318
for input_socket, _ in poller.poll():
1319
# (RequestType, RequestData)
1320
type_frame, *data_frames = input_socket.recv_multipart(copy=False)
1321
+ # NOTE(yongji): ignore READY message sent by DP coordinator
1322
+ # that is used to notify newly started engines
1323
+ if type_frame.buffer == b"READY":
1324
+ assert input_socket == coord_socket
1325
+ continue
1326
request_type = EngineCoreRequestType(bytes(type_frame.buffer))
1327
1328
# Deserialize the request data.
@@ -1488,6 +1521,10 @@ class DPEngineCoreProc(EngineCoreProc):
1521
self.current_wave = 0
1522
self.last_counts = (0, 0)
1523
1524
+ from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
1525
+
1526
+ self.eep_scaling_state: ElasticEPScalingState | None = None
1527
+
1528
# Initialize the engine.
1529
dp_rank = vllm_config.parallel_config.data_parallel_rank
1530
super().__init__(
@@ -1511,7 +1548,9 @@ class DPEngineCoreProc(EngineCoreProc):
1548
assert 0 <= local_dp_rank <= dp_rank < dp_size
1549
1550
self.dp_rank = dp_rank
1514
- self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()
1551
+ self.dp_group, self.dp_store = (
1552
+ vllm_config.parallel_config.stateless_init_dp_group(return_store=True)
1553
+ )
1554
1555
def shutdown(self):
1556
super().shutdown()
@@ -1574,7 +1613,12 @@ class DPEngineCoreProc(EngineCoreProc):
1613
# 1) Poll the input queue until there is work to do.
1614
self._process_input_queue()
1615
1577
- # 2) Step the engine core.
1616
+ if self.eep_scaling_state is not None:
1617
+ _ = self.eep_scaling_state.progress()
1618
+ if self.eep_scaling_state.is_complete():
1619
+ self.process_input_queue_block = True
1620
+ self.eep_scaling_state = None
1621
+
1622
executed = self._process_engine_step()
1623
self._maybe_publish_request_counts()
1624
@@ -1624,54 +1668,129 @@ class DPEngineCoreProc(EngineCoreProc):
1668
def reinitialize_distributed(
1669
self, reconfig_request: ReconfigureDistributedRequest
1670
) -> None:
1627
- stateless_destroy_torch_distributed_process_group(self.dp_group)
1628
- self.shutdown()
1629
-
1630
- parallel_config = self.vllm_config.parallel_config
1631
- old_dp_size = parallel_config.data_parallel_size
1632
- parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size
1633
- if reconfig_request.new_data_parallel_rank != -1:
1634
- parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
1635
- # local rank specifies device visibility, it should not be changed
1636
- assert (
1637
- reconfig_request.new_data_parallel_rank_local
1638
- == ReconfigureRankType.KEEP_CURRENT_RANK
1639
- )
1640
- parallel_config.data_parallel_master_ip = (
1671
+ from copy import deepcopy
1672
+
1673
+ from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
1674
+
1675
+ new_parallel_config = deepcopy(self.vllm_config.parallel_config)
1676
+ old_dp_size = new_parallel_config.data_parallel_size
1677
+ new_parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size
1678
+ if (
1679
+ reconfig_request.new_data_parallel_rank
1680
+ != ReconfigureRankType.KEEP_CURRENT_RANK
1681
+ ):
1682
+ new_parallel_config.data_parallel_rank = (
1683
+ reconfig_request.new_data_parallel_rank
1684
+ )
1685
+ new_parallel_config.data_parallel_master_ip = (
1686
reconfig_request.new_data_parallel_master_ip
1687
)
1643
- parallel_config.data_parallel_master_port = (
1688
+ new_parallel_config.data_parallel_master_port = (
1689
reconfig_request.new_data_parallel_master_port
1690
)
1646
- if reconfig_request.new_data_parallel_rank != -2:
1647
- self.dp_rank = parallel_config.data_parallel_rank
1648
- self.dp_group = parallel_config.stateless_init_dp_group()
1649
- reconfig_request.new_data_parallel_master_port = (
1650
- parallel_config.data_parallel_master_port
1691
+ new_parallel_config._data_parallel_master_port_list = (
1692
+ reconfig_request.new_data_parallel_master_port_list
1693
)
1694
1653
- self.model_executor.reinitialize_distributed(reconfig_request)
1654
- if reconfig_request.new_data_parallel_size > old_dp_size:
1655
- assert self.available_gpu_memory_for_kv_cache > 0
1656
- # pass available_gpu_memory_for_kv_cache from existing
1657
- # engine-cores to new engine-cores so they can directly
1658
- # use it in _initialize_kv_caches() rather than profiling.
1659
- ParallelConfig.sync_kv_cache_memory_size(
1660
- self.dp_group, self.available_gpu_memory_for_kv_cache
1661
- )
1662
- # NOTE(yongji): newly joined workers require dummy_run even
1663
- # CUDA graph is not used
1664
- self.model_executor.collective_rpc("compile_or_warm_up_model")
1665
- if (
1695
+ is_scale_down = reconfig_request.new_data_parallel_size < old_dp_size
1696
+ is_shutdown = (
1697
reconfig_request.new_data_parallel_rank
1698
== ReconfigureRankType.SHUTDOWN_CURRENT_RANK
1668
- ):
1669
- self.shutdown()
1670
- logger.info("DPEngineCoreProc %s shutdown", self.dp_rank)
1699
+ )
1700
+
1701
+ self.eep_scaling_state = ElasticEPScalingState(
1702
+ model_executor=self.model_executor,
1703
+ engine_core=self,
1704
+ vllm_config=self.vllm_config,
1705
+ new_parallel_config=new_parallel_config,
1706
+ worker_type="removing" if is_shutdown else "existing",
1707
+ scale_type="scale_down" if is_scale_down else "scale_up",
1708
+ reconfig_request=reconfig_request,
1709
+ )
1710
+ self.process_input_queue_block = False
1711
+ logger.info(
1712
+ "[Elastic EP] Received reconfiguration request and starting scaling up/down"
1713
+ )
1714
+
1715
+ def _eep_send_engine_core_notification(
1716
+ self,
1717
+ notification_type: EEPNotificationType,
1718
+ vllm_config: VllmConfig | None = None,
1719
+ ):
1720
+ """
1721
+ Send notifications to EngineCoreClient, which can then forward
1722
+ the notifications to other engine core processes. It is used for:
1723
+ 1) In scale up: new core engines to notify exisiting core engines
1724
+ that they are ready;
1725
+ 2) In scale down: removing core engines to notify EngineCoreClient
1726
+ so EngineCoreClient can release their ray placement groups;
1727
+ 3) Both scale up/down: to notify EngineCoreClient that exisiting
1728
+ core engines have already switched to the new parallel setup.
1729
+ """
1730
+ if vllm_config is None:
1731
+ dp_rank = self.vllm_config.parallel_config.data_parallel_rank
1732
else:
1672
- logger.info(
1673
- "Distributed environment reinitialized for DP rank %s", self.dp_rank
1733
+ dp_rank = vllm_config.parallel_config.data_parallel_rank
1734
+ notification_data = (notification_type.value, dp_rank)
1735
+ outputs = EngineCoreOutputs(
1736
+ utility_output=UtilityOutput(
1737
+ call_id=EEP_NOTIFICATION_CALL_ID,
1738
+ result=UtilityResult(notification_data),
1739
)
1740
+ )
1741
+ outputs.engine_index = self.engine_index
1742
+
1743
+ if hasattr(self, "output_thread") and self.output_thread.is_alive():
1744
+ self.output_queue.put_nowait((0, outputs))
1745
+ else:
1746
+ encoder = MsgpackEncoder()
1747
+ with (
1748
+ zmq.Context() as ctx,
1749
+ make_zmq_socket(
1750
+ ctx, self.addresses.outputs[0], zmq.PUSH, linger=4000
1751
+ ) as socket,
1752
+ ):
1753
+ socket.send_multipart(encoder.encode(outputs))
1754
+
1755
+ def eep_handle_engine_core_notification(
1756
+ self, notification_type: str | EEPNotificationType
1757
+ ):
1758
+ """
1759
+ Handle notification received from EngineCoreClient
1760
+ (forwarded from new core engines).
1761
+ """
1762
+ assert self.eep_scaling_state is not None
1763
+ if isinstance(notification_type, str):
1764
+ notification_type = EEPNotificationType(notification_type)
1765
+ self.eep_scaling_state.handle_notification(notification_type)
1766
+
1767
+ def _eep_scale_up_before_kv_init(self):
1768
+ from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
1769
+
1770
+ self.eep_scaling_state = ElasticEPScalingState(
1771
+ model_executor=self.model_executor,
1772
+ engine_core=self,
1773
+ vllm_config=self.vllm_config,
1774
+ new_parallel_config=self.vllm_config.parallel_config,
1775
+ worker_type="new",
1776
+ scale_type="scale_up",
1777
+ reconfig_request=None,
1778
+ )
1779
+ self.model_executor.collective_rpc("init_device")
1780
+ self.model_executor.collective_rpc("load_model")
1781
+ self._eep_send_engine_core_notification(
1782
+ EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY
1783
+ )
1784
+ self.model_executor.collective_rpc(
1785
+ "elastic_ep_execute", args=("receive_weights",)
1786
+ )
1787
+ self.available_gpu_memory_for_kv_cache = (
1788
+ ParallelConfig.sync_kv_cache_memory_size(self.dp_group, -1)
1789
+ )
1790
+ self.model_executor.collective_rpc(
1791
+ "elastic_ep_execute", args=("prepare_new_worker",)
1792
+ )
1793
+ self.process_input_queue_block = False
1794
1795
1796
class EngineCoreActorMixin:
vllm/v1/engine/core_client.py
+268
-51
@@ -28,11 +28,12 @@ from vllm.tracing import instrument
28
from vllm.utils.async_utils import in_loop
29
from vllm.utils.network_utils import (
30
close_sockets,
31
- get_open_port,
31
get_open_zmq_inproc_path,
32
make_zmq_socket,
33
)
34
from vllm.v1.engine import (
35
+ EEP_NOTIFICATION_CALL_ID,
36
+ EEPNotificationType,
37
EngineCoreOutputs,
38
EngineCoreRequest,
39
EngineCoreRequestType,
@@ -47,6 +48,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
48
from vllm.v1.engine.utils import (
49
CoreEngineActorManager,
50
CoreEngineProcManager,
51
+ get_engine_zmq_addresses,
52
launch_core_engines,
53
)
54
from vllm.v1.executor import Executor
@@ -445,6 +447,63 @@ class BackgroundResources:
447
raise EngineDeadError()
448
449
450
+@dataclass
451
+class ElasticScalingCache:
452
+ existing_core_engines: list[EngineIdentity]
453
+ num_new_core_engines: int
454
+ pending_notifications: dict[EEPNotificationType, set[int]]
455
+
456
+
457
+def allocate_stateless_group_ports(parallel_config, new_data_parallel_size: int):
458
+ """
459
+ Allocate stateless group ports for elastic EP.
460
+ """
461
+ from vllm.utils.network_utils import get_open_ports_list
462
+
463
+ assert parallel_config.enable_elastic_ep, "Elastic EP must be enabled"
464
+ world_size = parallel_config.world_size
465
+ new_world_size_across_dp = world_size * new_data_parallel_size
466
+ num_world_groups = 1
467
+ num_dp_groups = max(1, new_world_size_across_dp // new_data_parallel_size)
468
+ num_ep_groups = max(
469
+ 1,
470
+ new_world_size_across_dp
471
+ // (new_data_parallel_size * parallel_config.tensor_parallel_size),
472
+ )
473
+ num_eplb_groups = num_ep_groups
474
+ total_ports_needed = (
475
+ num_world_groups + num_dp_groups + num_ep_groups + num_eplb_groups
476
+ ) * 3 + 5
477
+ all_ports = get_open_ports_list(total_ports_needed)
478
+ new_data_parallel_master_port_list = all_ports[-5:]
479
+ all_ports = all_ports[:-5]
480
+ new_stateless_world_group_port_list = [
481
+ all_ports[i : i + 3] for i in range(0, num_world_groups * 3, 3)
482
+ ]
483
+ start_idx = num_world_groups * 3
484
+ new_stateless_dp_group_port_list = [
485
+ all_ports[i : i + 3] for i in range(start_idx, start_idx + num_dp_groups * 3, 3)
486
+ ]
487
+ start_idx += num_dp_groups * 3
488
+ new_stateless_ep_group_port_list = [
489
+ all_ports[i : i + 3] for i in range(start_idx, start_idx + num_ep_groups * 3, 3)
490
+ ]
491
+ start_idx += num_ep_groups * 3
492
+ new_stateless_eplb_group_port_list = [
493
+ all_ports[i : i + 3]
494
+ for i in range(start_idx, start_idx + num_eplb_groups * 3, 3)
495
+ ]
496
+
497
+ parallel_config._stateless_world_group_port_list = (
498
+ new_stateless_world_group_port_list
499
+ )
500
+ parallel_config._stateless_dp_group_port_list = new_stateless_dp_group_port_list
501
+ parallel_config._stateless_ep_group_port_list = new_stateless_ep_group_port_list
502
+ parallel_config._stateless_eplb_group_port_list = new_stateless_eplb_group_port_list
503
+ parallel_config.data_parallel_master_port = new_data_parallel_master_port_list.pop()
504
+ parallel_config._data_parallel_master_port_list = new_data_parallel_master_port_list
505
+
506
+
507
class MPClient(EngineCoreClient):
508
"""
509
MPClient: base client for multi-proc EngineCore.
@@ -491,32 +550,37 @@ class MPClient(EngineCoreClient):
550
input_address = client_addresses["input_address"]
551
output_address = client_addresses["output_address"]
552
self.stats_update_address = client_addresses.get("stats_update_address")
553
+ self.input_socket = self.resources.input_socket = make_zmq_socket(
554
+ self.ctx, input_address, zmq.ROUTER, bind=True
555
+ )
556
+ self.resources.output_socket = make_zmq_socket(
557
+ self.ctx, output_address, zmq.PULL
558
+ )
559
else:
560
# Engines are managed by this client.
496
- with launch_core_engines(vllm_config, executor_class, log_stats) as (
497
- engine_manager,
498
- coordinator,
561
+ addresses = get_engine_zmq_addresses(vllm_config)
562
+ self.input_socket = self.resources.input_socket = make_zmq_socket(
563
+ self.ctx, addresses.inputs[0], zmq.ROUTER, bind=True
564
+ )
565
+ self.resources.output_socket = make_zmq_socket(
566
+ self.ctx, addresses.outputs[0], zmq.PULL
567
+ )
568
+
569
+ with launch_core_engines(
570
+ vllm_config,
571
+ executor_class,
572
+ log_stats,
573
addresses,
500
- ):
574
+ ) as (engine_manager, coordinator, addresses):
575
self.resources.coordinator = coordinator
576
self.resources.engine_manager = engine_manager
577
504
- (input_address,) = addresses.inputs
505
- (output_address,) = addresses.outputs
578
self.stats_update_address = addresses.frontend_stats_publish_address
579
if coordinator is not None:
580
assert self.stats_update_address == (
581
coordinator.get_stats_publish_address()
582
)
583
512
- # Create input and output sockets.
513
- self.input_socket = self.resources.input_socket = make_zmq_socket(
514
- self.ctx, input_address, zmq.ROUTER, bind=True
515
- )
516
- self.resources.output_socket = make_zmq_socket(
517
- self.ctx, output_address, zmq.PULL
518
- )
519
-
584
parallel_config = vllm_config.parallel_config
585
dp_size = parallel_config.data_parallel_size
586
dp_rank = parallel_config.data_parallel_index
@@ -877,6 +941,10 @@ class AsyncMPClient(MPClient):
941
output_socket = resources.output_socket
942
assert output_socket is not None
943
944
+ notification_callback_handler: (
945
+ Callable[[AsyncMPClient, Sequence[Any]], Any] | None
946
+ ) = getattr(self.__class__, "eep_process_engine_core_notification", None)
947
+
948
async def process_outputs_socket():
949
try:
950
while True:
@@ -884,7 +952,26 @@ class AsyncMPClient(MPClient):
952
resources.validate_alive(frames)
953
outputs: EngineCoreOutputs = decoder.decode(frames)
954
if outputs.utility_output:
887
- _process_utility_output(outputs.utility_output, utility_results)
955
+ if (
956
+ outputs.utility_output.call_id == EEP_NOTIFICATION_CALL_ID
957
+ and notification_callback_handler is not None
958
+ ):
959
+ assert _self_ref is not None
960
+ _self = _self_ref()
961
+ if not _self:
962
+ return
963
+ if outputs.utility_output.result is None:
964
+ continue
965
+ notification_data = outputs.utility_output.result.result
966
+ assert isinstance(notification_data, Sequence)
967
+ assert len(notification_data) == 2
968
+ asyncio.create_task(
969
+ notification_callback_handler(_self, notification_data)
970
+ )
971
+ else:
972
+ _process_utility_output(
973
+ outputs.utility_output, utility_results
974
+ )
975
continue
976
977
if output_handler is not None:
@@ -1081,6 +1168,8 @@ class DPAsyncMPClient(AsyncMPClient):
1168
# Used only by DPLBAsyncMPClient subclass.
1169
self.lb_engines: list[list[int]] = [[0, 0] for _ in self.core_engines]
1170
1171
+ self.eep_scaling_cache: ElasticScalingCache | None = None
1172
+
1173
self.first_req_sock_addr = get_open_zmq_inproc_path()
1174
self.first_req_send_socket = self.resources.first_req_send_socket = (
1175
make_zmq_socket(self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=True)
@@ -1101,12 +1190,6 @@ class DPAsyncMPClient(AsyncMPClient):
1190
assert self.stats_update_address is not None
1191
stats_addr: str = self.stats_update_address
1192
assert len(self.engine_ranks_managed) > 0
1104
- # NOTE: running and waiting counts are all global from
1105
- # the Coordinator include all global EngineCores. This
1106
- # slice includes just the cores managed by this client.
1107
- count_slice = slice(
1108
- self.engine_ranks_managed[0], self.engine_ranks_managed[-1] + 1
1109
- )
1193
1194
async def run_engine_stats_update_task():
1195
with (
@@ -1145,6 +1228,29 @@ class DPAsyncMPClient(AsyncMPClient):
1228
):
1229
# Extract new engine count from the decoded message
1230
new_engine_count = decoded[1]
1231
+ # Update engine_ranks_managed and count_slice
1232
+ parallel_config = self.vllm_config.parallel_config
1233
+ dp_size = parallel_config.data_parallel_size
1234
+ dp_rank = parallel_config.data_parallel_rank
1235
+ assert dp_rank == 0
1236
+ assert dp_size == new_engine_count
1237
+ assert not (
1238
+ parallel_config.data_parallel_hybrid_lb
1239
+ or parallel_config.data_parallel_external_lb
1240
+ )
1241
+ num_ranks = dp_size
1242
+ self.engine_ranks_managed = list(
1243
+ range(dp_rank, dp_rank + num_ranks)
1244
+ )
1245
+ if len(self.lb_engines) < new_engine_count:
1246
+ self.lb_engines = self.lb_engines + [
1247
+ [0, 0]
1248
+ for _ in range(
1249
+ new_engine_count - len(self.lb_engines)
1250
+ )
1251
+ ]
1252
+ else:
1253
+ self.lb_engines = self.lb_engines[:new_engine_count]
1254
# Send scale up notification to coordinator
1255
scale_msg = msgspec.msgpack.encode(
1256
("SCALE_ELASTIC_EP", new_engine_count)
@@ -1178,6 +1284,11 @@ class DPAsyncMPClient(AsyncMPClient):
1284
self.current_wave = wave
1285
self.engines_running = running
1286
if counts is not None:
1287
+ # Running and waiting counts are global from the
1288
+ # Coordinator including all EngineCores. Slice to get
1289
+ # just the cores managed by this client.
1290
+ ranks = self.engine_ranks_managed
1291
+ count_slice = slice(ranks[0], ranks[-1] + 1)
1292
sliced_counts = counts[count_slice]
1293
self.lb_engines = sliced_counts
1294
logger.debug(
@@ -1287,6 +1398,67 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1398
for req_id in outputs.finished_requests:
1399
self.reqs_in_flight.pop(req_id, None)
1400
1401
+ @staticmethod
1402
+ async def eep_process_engine_core_notification(
1403
+ self: "DPLBAsyncMPClient", notification_data: tuple[str, int]
1404
+ ):
1405
+ cache = self.eep_scaling_cache
1406
+ notification_type_str, dp_rank = notification_data
1407
+ try:
1408
+ notification_type = EEPNotificationType(notification_type_str)
1409
+ except ValueError as e:
1410
+ raise ValueError(
1411
+ f"Unknown EEP notification type: {notification_type_str}"
1412
+ ) from e
1413
+
1414
+ if notification_type == EEPNotificationType.RECONFIGURE_FINISHED:
1415
+ from vllm.v1.engine import UtilityResult
1416
+
1417
+ # NOTE(yongji): process a dummy UtilityOutput to resolve the future
1418
+ # awaited in _eep_wait_for_setup_switch_complete(), signaling that
1419
+ # all engine cores have completed reconfiguration.
1420
+ dummy_output = UtilityOutput(
1421
+ call_id=EEP_NOTIFICATION_CALL_ID, result=UtilityResult(None)
1422
+ )
1423
+ _process_utility_output(dummy_output, self.utility_results)
1424
+ return
1425
+ assert cache is not None
1426
+ if notification_type not in cache.pending_notifications:
1427
+ cache.pending_notifications[notification_type] = set()
1428
+ if dp_rank in cache.pending_notifications[notification_type]:
1429
+ raise ValueError(
1430
+ f"Duplicate notification {notification_type} from dp_rank {dp_rank}"
1431
+ )
1432
+ cache.pending_notifications[notification_type].add(dp_rank)
1433
+ if len(cache.pending_notifications[notification_type]) >= abs(
1434
+ cache.num_new_core_engines
1435
+ ):
1436
+ if notification_type == EEPNotificationType.SHUTDOWN_COMPLETE:
1437
+ assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
1438
+ assert cache.num_new_core_engines < 0
1439
+ old_dp_size = len(cache.existing_core_engines)
1440
+ new_dp_size = old_dp_size + cache.num_new_core_engines
1441
+ self.resources.engine_manager.scale_down_elastic_ep(
1442
+ old_dp_size, new_dp_size
1443
+ )
1444
+ else:
1445
+ await asyncio.gather(
1446
+ *[
1447
+ self._call_utility_async(
1448
+ "eep_handle_engine_core_notification",
1449
+ notification_type,
1450
+ engine=engine,
1451
+ )
1452
+ for engine in cache.existing_core_engines
1453
+ ]
1454
+ )
1455
+ cache.pending_notifications[notification_type] = set()
1456
+ if notification_type in [
1457
+ EEPNotificationType.SHUTDOWN_COMPLETE,
1458
+ EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY,
1459
+ ]:
1460
+ self.eep_scaling_cache = None
1461
+
1462
async def abort_requests_async(self, request_ids: list[str]) -> None:
1463
if not request_ids or self.resources.engine_dead:
1464
return
@@ -1333,6 +1505,20 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1505
cur_data_parallel_size, new_data_parallel_size
1506
)
1507
1508
+ async def _eep_wait_for_setup_switch_complete(self) -> None:
1509
+ """
1510
+ Wait for core engines to switch to the new setup.
1511
+
1512
+ In eep_process_engine_core_notification(), a dummy UtilityOutput with
1513
+ EEP_NOTIFICATION_CALL_ID will be set when RECONFIGURE_FINISHED
1514
+ notification is received from engine 0. We create a future with
1515
+ that call_id and wait for it to be resolved.
1516
+ """
1517
+ future = asyncio.get_running_loop().create_future()
1518
+ self.utility_results[EEP_NOTIFICATION_CALL_ID] = future
1519
+ self._ensure_output_queue_task()
1520
+ await future
1521
+
1522
async def _scale_up_elastic_ep(
1523
self, cur_data_parallel_size: int, new_data_parallel_size: int
1524
) -> None:
@@ -1340,38 +1526,57 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1526
and reconfiguring existing ones."""
1527
cur_data_parallel_size = len(self.core_engines)
1528
1343
- # Phase 1: Send reconfigure messages to all existing engines and wait
1344
- # for them to be sent
1529
+ self.eep_scaling_cache = ElasticScalingCache(
1530
+ existing_core_engines=self.core_engines.copy(),
1531
+ num_new_core_engines=new_data_parallel_size - cur_data_parallel_size,
1532
+ pending_notifications=dict(),
1533
+ )
1534
+
1535
+ parallel_config = self.vllm_config.parallel_config
1536
+ allocate_stateless_group_ports(parallel_config, new_data_parallel_size)
1537
+
1538
+ # Phase 1: Send reconfig messages to existing engines
1539
reconfig_futures = []
1346
- self.vllm_config.parallel_config.data_parallel_master_port = get_open_port()
1540
for engine in self.core_engines:
1541
reconfig_request = ReconfigureDistributedRequest(
1542
new_data_parallel_size=new_data_parallel_size,
1543
new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,
1544
new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,
1352
- new_data_parallel_master_ip=self.vllm_config.parallel_config.data_parallel_master_ip,
1353
- new_data_parallel_master_port=self.vllm_config.parallel_config.data_parallel_master_port,
1545
+ new_data_parallel_master_ip=parallel_config.data_parallel_master_ip,
1546
+ new_data_parallel_master_port=parallel_config.data_parallel_master_port,
1547
+ new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,
1548
+ new_stateless_world_group_port_list=parallel_config._stateless_world_group_port_list,
1549
+ new_stateless_dp_group_port_list=parallel_config._stateless_dp_group_port_list,
1550
+ new_stateless_ep_group_port_list=parallel_config._stateless_ep_group_port_list,
1551
+ new_stateless_eplb_group_port_list=parallel_config._stateless_eplb_group_port_list,
1552
)
1553
coro = self._call_utility_async(
1554
"reinitialize_distributed", reconfig_request, engine=engine
1555
)
1556
reconfig_futures.append(asyncio.create_task(coro))
1557
1360
- logger.info("All reconfigure messages sent, starting engine creation")
1361
-
1362
- # Phase 2: Create new engines now that reconfig messages have been sent
1363
- # self.resources.engine_manager is guaranteed to be
1364
- # CoreEngineActorManager for RayDPClient
1558
+ # Phase 2: Create new engines
1559
assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
1366
- self.resources.engine_manager.scale_up_elastic_ep(
1367
- self.vllm_config, new_data_parallel_size
1560
+ parallel_config.eplb_config.num_redundant_experts = 0
1561
+ start_new_worker_future = asyncio.to_thread(
1562
+ self.resources.engine_manager.scale_up_elastic_ep,
1563
+ self.vllm_config,
1564
+ new_data_parallel_size,
1565
)
1566
+ wait_future = self._eep_wait_for_setup_switch_complete()
1567
+
1568
+ # Phase 3: Wait for new engines to be created
1569
+ # and reconfig messages to be received
1570
+ await asyncio.gather(start_new_worker_future, *reconfig_futures)
1571
+ logger.info("[Elastic EP] Successfully started new engines")
1572
1573
# Create new CoreEngine objects for the new engines
1574
new_engine_identities = set()
1575
for i in range(cur_data_parallel_size, new_data_parallel_size):
1576
new_engine = i.to_bytes(2, "little")
1577
self.core_engines.append(new_engine)
1578
+ # NOTE(yongji): we don't update lb_engines here,
1579
+ # we let run_engine_stats_update_task to update it.
1580
new_engine_identities.add(new_engine)
1581
1582
# Wait for ready messages from new engines on the input socket
@@ -1387,10 +1592,11 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1592
identity, _ = sync_input_socket.recv_multipart()
1593
new_engine_identities.discard(identity)
1594
1390
- # Phase 3: Wait for all existing engines to complete reconfiguration
1391
- logger.info("Waiting for existing engines to complete reconfiguration")
1392
- await asyncio.gather(*reconfig_futures)
1393
-
1595
+ # NOTE(yongji): Before we schedule any requests on the new workers,
1596
+ # we should wait for them to switch to the new setup.
1597
+ await wait_future
1598
+ # Update the parallel config
1599
+ self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1600
# Notify coordinator about scale up through existing
1601
# stats_update_task connection
1602
self._ensure_stats_update_task()
@@ -1399,8 +1605,6 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1605
)
1606
await self.first_req_send_socket.send(scale_up_marker)
1607
1402
- # Update the parallel config
1403
- self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1608
logger.info(
1609
"[Elastic EP] Scale up completed, new data parallel size: %s",
1610
new_data_parallel_size,
@@ -1413,7 +1617,14 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1617
reconfiguring existing engine cores."""
1618
cur_data_parallel_size = len(self.core_engines)
1619
1416
- self.vllm_config.parallel_config.data_parallel_master_port = get_open_port()
1620
+ self.eep_scaling_cache = ElasticScalingCache(
1621
+ existing_core_engines=self.core_engines.copy(),
1622
+ num_new_core_engines=new_data_parallel_size - cur_data_parallel_size,
1623
+ pending_notifications=dict(),
1624
+ )
1625
+
1626
+ parallel_config = self.vllm_config.parallel_config
1627
+ allocate_stateless_group_ports(parallel_config, new_data_parallel_size)
1628
1629
reconfig_futures = []
1630
for cur_dp_rank, engine in enumerate(self.core_engines):
@@ -1421,8 +1632,13 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1632
new_data_parallel_size=new_data_parallel_size,
1633
new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,
1634
new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,
1424
- new_data_parallel_master_ip=self.vllm_config.parallel_config.data_parallel_master_ip,
1425
- new_data_parallel_master_port=self.vllm_config.parallel_config.data_parallel_master_port,
1635
+ new_data_parallel_master_ip=parallel_config.data_parallel_master_ip,
1636
+ new_data_parallel_master_port=parallel_config.data_parallel_master_port,
1637
+ new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,
1638
+ new_stateless_world_group_port_list=parallel_config._stateless_world_group_port_list,
1639
+ new_stateless_dp_group_port_list=parallel_config._stateless_dp_group_port_list,
1640
+ new_stateless_ep_group_port_list=parallel_config._stateless_ep_group_port_list,
1641
+ new_stateless_eplb_group_port_list=parallel_config._stateless_eplb_group_port_list,
1642
)
1643
if cur_dp_rank >= new_data_parallel_size:
1644
reconfig_request.new_data_parallel_rank = (
@@ -1433,23 +1649,24 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1649
)
1650
reconfig_futures.append(asyncio.create_task(coro))
1651
1436
- for _ in range(new_data_parallel_size, cur_data_parallel_size):
1437
- self.core_engines.pop()
1652
+ # NOTE(yongji): Immediately stop sending requests to the removing engines.
1653
+ self.core_engines = self.core_engines[:new_data_parallel_size]
1654
+ self.lb_engines = self.lb_engines[:new_data_parallel_size]
1655
+ wait_future = self._eep_wait_for_setup_switch_complete()
1656
1657
await asyncio.gather(*reconfig_futures)
1658
1441
- assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
1442
- self.resources.engine_manager.scale_down_elastic_ep(
1443
- cur_data_parallel_size, new_data_parallel_size
1444
- )
1445
-
1659
+ self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1660
self._ensure_stats_update_task()
1661
scale_down_marker = msgspec.msgpack.encode(
1662
("SCALE_ELASTIC_EP", new_data_parallel_size)
1663
)
1664
await self.first_req_send_socket.send(scale_down_marker)
1665
1452
- self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1666
+ # NOTE(yongji): Unlike scaling up,
1667
+ # here we don't actually need to wait for the setup switch to complete.
1668
+ # We may want to remove it in the future.
1669
+ await wait_future
1670
logger.info(
1671
"[Elastic EP] Scale down completed, new data parallel size: %s",
1672
new_data_parallel_size,
vllm/v1/engine/utils.py
+43
-17
@@ -277,6 +277,8 @@ class CoreEngineActorManager:
277
else:
278
ray.init()
279
280
+ vllm_config.parallel_config.allocate_elastic_ep_ports()
281
+
282
if placement_groups is not None:
283
assert local_dp_ranks is not None, (
284
"local_dp_ranks must be provided if placement_groups is provided"
@@ -584,6 +586,8 @@ class CoreEngineActorManager:
586
587
node_ip = node.node_ip
588
node_id = node.node_id
589
+ if device_str not in available_resources[node_id]:
590
+ continue
591
available_gpus = int(available_resources[node_id][device_str])
592
593
# Get total GPUs on this node from the node's resources
@@ -773,26 +777,15 @@ class CoreEngineActorManager:
777
ray.util.remove_placement_group(pg)
778
779
776
-@contextlib.contextmanager
777
-def launch_core_engines(
780
+def get_engine_zmq_addresses(
781
vllm_config: VllmConfig,
779
- executor_class: type[Executor],
780
- log_stats: bool,
782
num_api_servers: int = 1,
782
-) -> Iterator[
783
- tuple[
784
- CoreEngineProcManager | CoreEngineActorManager | None,
785
- DPCoordinator | None,
786
- EngineZmqAddresses,
787
- ]
788
-]:
789
- """Launch engine and DP coordinator processes as needed."""
790
-
783
+) -> EngineZmqAddresses:
784
+ """Allocate ZMQ addresses for engine-client communication."""
785
parallel_config = vllm_config.parallel_config
792
- dp_size = parallel_config.data_parallel_size
786
local_engine_count = parallel_config.data_parallel_size_local
787
local_start_index = parallel_config.data_parallel_rank_local
795
- dp_rank = parallel_config.data_parallel_rank
788
+ dp_size = parallel_config.data_parallel_size
789
host = parallel_config.data_parallel_master_ip
790
local_engines_only = parallel_config.local_engines_only
791
@@ -806,9 +799,11 @@ def launch_core_engines(
799
client_local_only = (
800
offline_mode or local_engines_only or (local_engine_count == dp_size)
801
)
802
+ # NOTE(yongji): handling scaling from intra-node to inter-node
803
+ if parallel_config.enable_elastic_ep:
804
+ client_local_only = False
805
810
- # Set up input and output addresses.
811
- addresses = EngineZmqAddresses(
806
+ return EngineZmqAddresses(
807
inputs=[
808
get_engine_client_zmq_addr(client_local_only, host)
809
for _ in range(num_api_servers)
@@ -819,6 +814,33 @@ def launch_core_engines(
814
],
815
)
816
817
+
818
+@contextlib.contextmanager
819
+def launch_core_engines(
820
+ vllm_config: VllmConfig,
821
+ executor_class: type[Executor],
822
+ log_stats: bool,
823
+ addresses: EngineZmqAddresses,
824
+ num_api_servers: int = 1,
825
+) -> Iterator[
826
+ tuple[
827
+ CoreEngineProcManager | CoreEngineActorManager | None,
828
+ DPCoordinator | None,
829
+ EngineZmqAddresses,
830
+ ]
831
+]:
832
+ """Launch engine and DP coordinator processes as needed."""
833
+
834
+ parallel_config = vllm_config.parallel_config
835
+ dp_size = parallel_config.data_parallel_size
836
+ local_engine_count = parallel_config.data_parallel_size_local
837
+ local_start_index = parallel_config.data_parallel_rank_local
838
+ dp_rank = parallel_config.data_parallel_rank
839
+ host = parallel_config.data_parallel_master_ip
840
+ local_engines_only = parallel_config.local_engines_only
841
+
842
+ offline_mode = local_start_index is not None
843
+
844
# Run the DP Coordinator process with rank 0 when in online DP mode.
845
# The coordinator is needed for:
846
# 1. Internal/hybrid LB: collecting and publishing queue stats for load balancing
@@ -885,6 +907,10 @@ def launch_core_engines(
907
# will be False.
908
handshake_local_only = offline_mode or local_engine_count == dp_size
909
910
+ # NOTE(yongji): handling scaling from intra-node to inter-node
911
+ if parallel_config.enable_elastic_ep:
912
+ handshake_local_only = False
913
+
914
handshake_address = get_engine_client_zmq_addr(
915
handshake_local_only, host, parallel_config.data_parallel_rpc_port
916
)
vllm/v1/executor/multiproc_executor.py
+16
-5
@@ -38,6 +38,7 @@ from vllm.distributed.parallel_state import (
38
get_pcp_group,
39
get_pp_group,
40
get_tp_group,
41
+ model_parallel_is_initialized,
42
)
43
from vllm.envs import enable_envs_cache
44
from vllm.logger import init_logger
@@ -580,17 +581,20 @@ class WorkerProc:
581
)
582
self.async_output_copy_thread.start()
583
583
- # Initialize device
584
- self.worker.init_device()
585
-
586
- # Set process title and log prefix
584
self.setup_proc_title_and_log_prefix(
585
enable_ep=vllm_config.parallel_config.enable_expert_parallel
586
)
587
588
# Load model
589
self._init_message_queues(input_shm_handle, vllm_config)
593
- self.worker.load_model()
590
+ is_eep_new_worker = envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH
591
+ if not is_eep_new_worker:
592
+ self.worker.init_device()
593
+ # Update process title now that parallel groups are initialized
594
+ self.setup_proc_title_and_log_prefix(
595
+ enable_ep=vllm_config.parallel_config.enable_expert_parallel
596
+ )
597
+ self.worker.load_model()
598
599
# Enable environment variable cache (e.g. assume no more
600
# environment variable overrides after this point)
@@ -885,6 +889,13 @@ class WorkerProc:
889
890
@staticmethod
891
def setup_proc_title_and_log_prefix(enable_ep: bool) -> None:
892
+ # Check if parallel groups are initialized first
893
+ if not model_parallel_is_initialized():
894
+ # Parallel groups not yet initialized, use default process name
895
+ set_process_title(name="Worker")
896
+ decorate_logs("Worker")
897
+ return
898
+
899
dp_size = get_dp_group().world_size
900
dp_rank = get_dp_group().rank_in_group
901
pp_size = get_pp_group().world_size
vllm/v1/executor/ray_executor.py
+4
-2
@@ -382,8 +382,10 @@ class RayDistributedExecutor(Executor):
382
all_kwargs.append(kwargs)
383
self.collective_rpc("init_worker", args=(all_kwargs,))
384
385
- self.collective_rpc("init_device")
386
- self.collective_rpc("load_model")
385
+ is_eep_new_worker = envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH
386
+ if not is_eep_new_worker:
387
+ self.collective_rpc("init_device")
388
+ self.collective_rpc("load_model")
389
390
for pp_rank in range(self.parallel_config.pipeline_parallel_size):
391
self.pp_tp_workers.append([])
vllm/v1/executor/uniproc_executor.py
+4
-13
@@ -14,7 +14,6 @@ import vllm.envs as envs
14
from vllm.logger import init_logger
15
from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
16
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
17
-from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
17
from vllm.v1.executor.abstract import Executor
18
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
19
from vllm.v1.serial_utils import run_method
@@ -43,9 +42,11 @@ class UniProcExecutor(Executor):
42
max_workers=1, thread_name_prefix="WorkerAsyncOutput"
43
)
44
45
+ is_eep_new_worker = envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH
46
self.driver_worker.init_worker(all_kwargs=[kwargs])
47
- self.driver_worker.init_device()
48
- self.driver_worker.load_model()
47
+ if not is_eep_new_worker:
48
+ self.driver_worker.init_device()
49
+ self.driver_worker.load_model()
50
51
def _distributed_args(self) -> tuple[str, int, int]:
52
"""Return (distributed_init_method, rank, local_rank)."""
@@ -122,16 +123,6 @@ class UniProcExecutor(Executor):
123
# it's running.
124
return
125
125
- def reinitialize_distributed(
126
- self, reconfig_request: ReconfigureDistributedRequest
127
- ) -> None:
128
- self.driver_worker.reinitialize_distributed(reconfig_request)
129
- if (
130
- reconfig_request.new_data_parallel_rank
131
- == ReconfigureRankType.SHUTDOWN_CURRENT_RANK
132
- ):
133
- self.shutdown()
134
-
126
def shutdown(self) -> None:
127
if worker := self.driver_worker:
128
worker.shutdown()
vllm/v1/worker/cpu_model_runner.py
+6
-1
@@ -53,7 +53,12 @@ class CPUModelRunner(GPUModelRunner):
53
v.gpu = v.cpu
54
55
@instrument(span_name="Loading (CPU)")
56
- def load_model(self, eep_scale_up: bool = False) -> None:
56
+ def load_model(self, load_dummy_weights: bool = False) -> None:
57
+ if load_dummy_weights:
58
+ raise ValueError(
59
+ "Loading dummy weights (needed for elastic EP scale-up) "
60
+ "Is not supported by the CPU Model Runner."
61
+ )
62
logger.info("Starting to load model %s...", self.model_config.model)
63
self.model = get_model(vllm_config=self.vllm_config)
64
vllm/v1/worker/gpu_model_runner.py
+39
-40
@@ -461,6 +461,8 @@ class GPUModelRunner(
461
self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
462
463
self.eplb_state: EplbState | None = None
464
+ # NOTE(yongji): flag to temporarily disable EPLB during scaling up/down
465
+ self.eep_eplb_suppressed = False
466
"""
467
State of the expert parallelism load balancer.
468
@@ -2702,7 +2704,7 @@ class GPUModelRunner(
2704
"""
2705
Step for the EPLB (Expert Parallelism Load Balancing) state.
2706
"""
2705
- if not self.parallel_config.enable_eplb:
2707
+ if not self.parallel_config.enable_eplb or self.eep_eplb_suppressed:
2708
return
2709
2710
assert self.eplb_state is not None
@@ -2714,6 +2716,23 @@ class GPUModelRunner(
2716
log_stats=self.parallel_config.eplb_config.log_balancedness,
2717
)
2718
2719
+ def setup_eplb_from_mapping(
2720
+ self,
2721
+ expanded_physical_to_logical: torch.Tensor,
2722
+ old_num_physical_experts: int,
2723
+ ) -> None:
2724
+ model = self.get_model()
2725
+ assert is_mixture_of_experts(model)
2726
+
2727
+ self.eplb_state = EplbState.from_mapping(
2728
+ model=model,
2729
+ model_config=self.model_config,
2730
+ device=self.device,
2731
+ parallel_config=self.parallel_config,
2732
+ expanded_physical_to_logical=expanded_physical_to_logical,
2733
+ num_valid_physical_experts=old_num_physical_experts,
2734
+ )
2735
+
2736
def _pool(
2737
self,
2738
hidden_states: torch.Tensor,
@@ -4175,21 +4194,16 @@ class GPUModelRunner(
4194
setattr(self, config_name, new_config)
4195
4196
@instrument(span_name="Loading (GPU)")
4178
- def load_model(self, eep_scale_up: bool = False) -> None:
4197
+ def load_model(self, load_dummy_weights: bool = False) -> None:
4198
"""
4199
Args:
4181
- eep_scale_up: the model loading is for elastic EP scale up.
4200
+ load_dummy_weights: load dummy weights instead of real weights.
4201
"""
4202
logger.info_once(
4203
"Starting to load model %s...",
4204
self.model_config.model,
4205
scope="global",
4206
)
4188
- global_expert_loads, old_global_expert_indices_per_model, rank_mapping = (
4189
- EplbState.get_eep_state(self.parallel_config)
4190
- if eep_scale_up
4191
- else (None, None, None)
4192
- )
4207
4208
if self.parallel_config.enable_eplb:
4209
self.eplb_state = EplbState(self.parallel_config, self.device)
@@ -4198,6 +4212,8 @@ class GPUModelRunner(
4212
try:
4213
with DeviceMemoryProfiler() as m:
4214
time_before_load = time.perf_counter()
4215
+ if load_dummy_weights:
4216
+ self.load_config.load_format = "dummy"
4217
model_loader = get_model_loader(self.load_config)
4218
self.model = model_loader.load_model(
4219
vllm_config=self.vllm_config, model_config=self.model_config
@@ -4214,6 +4230,9 @@ class GPUModelRunner(
4230
and is_mixture_of_experts(self.drafter.model)
4231
and self.parallel_config.enable_eplb
4232
):
4233
+ assert not self.parallel_config.enable_elastic_ep, (
4234
+ "Elastic EP is not supported with drafter model."
4235
+ )
4236
spec_config = self.vllm_config.speculative_config
4237
assert spec_config is not None
4238
assert spec_config.draft_model_config is not None
@@ -4221,17 +4240,6 @@ class GPUModelRunner(
4240
"EPLB is enabled for drafter model %s.",
4241
spec_config.draft_model_config.model,
4242
)
4224
-
4225
- global_expert_load = (
4226
- global_expert_loads[eplb_models]
4227
- if global_expert_loads
4228
- else None
4229
- )
4230
- old_global_expert_indices = (
4231
- old_global_expert_indices_per_model[eplb_models]
4232
- if old_global_expert_indices_per_model
4233
- else None
4234
- )
4243
if self.eplb_state is None:
4244
self.eplb_state = EplbState(
4245
self.parallel_config, self.device
@@ -4239,9 +4247,6 @@ class GPUModelRunner(
4247
self.eplb_state.add_model(
4248
self.drafter.model,
4249
spec_config.draft_model_config,
4242
- global_expert_load,
4243
- old_global_expert_indices,
4244
- rank_mapping,
4250
)
4251
eplb_models += 1
4252
@@ -4283,11 +4288,12 @@ class GPUModelRunner(
4288
time_after_load - time_before_load,
4289
scope="local",
4290
)
4286
- prepare_communication_buffer_for_model(self.model)
4287
- if (drafter := getattr(self, "drafter", None)) and (
4288
- drafter_model := getattr(drafter, "model", None)
4289
- ):
4290
- prepare_communication_buffer_for_model(drafter_model)
4291
+ if not load_dummy_weights:
4292
+ prepare_communication_buffer_for_model(self.model)
4293
+ if (drafter := getattr(self, "drafter", None)) and (
4294
+ drafter_model := getattr(drafter, "model", None)
4295
+ ):
4296
+ prepare_communication_buffer_for_model(drafter_model)
4297
mm_config = self.model_config.multimodal_config
4298
self.is_multimodal_pruning_enabled = (
4299
supports_multimodal_pruning(self.get_model())
@@ -4295,26 +4301,19 @@ class GPUModelRunner(
4301
and mm_config.is_multimodal_pruning_enabled()
4302
)
4303
4298
- if is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb:
4304
+ if (
4305
+ is_mixture_of_experts(self.model)
4306
+ and self.parallel_config.enable_eplb
4307
+ and not load_dummy_weights
4308
+ ):
4309
logger.info_once("EPLB is enabled for model %s.", self.model_config.model)
4300
- global_expert_load = (
4301
- global_expert_loads[eplb_models] if global_expert_loads else None
4302
- )
4303
- old_global_expert_indices = (
4304
- old_global_expert_indices_per_model[eplb_models]
4305
- if old_global_expert_indices_per_model
4306
- else None
4307
- )
4310
assert self.eplb_state is not None
4311
self.eplb_state.add_model(
4312
self.model,
4313
self.model_config,
4312
- global_expert_load,
4313
- old_global_expert_indices,
4314
- rank_mapping,
4314
)
4315
if self.eplb_state.is_async:
4317
- self.eplb_state.start_async_loop(rank_mapping=rank_mapping)
4316
+ self.eplb_state.start_async_loop()
4317
4318
if (
4319
self.vllm_config.compilation_config.mode
vllm/v1/worker/gpu_worker.py
+27
-228
@@ -7,11 +7,10 @@ import os
7
from collections.abc import Callable
8
from contextlib import AbstractContextManager, nullcontext
9
from types import NoneType
10
-from typing import TYPE_CHECKING, Any, cast
10
+from typing import TYPE_CHECKING, Any
11
12
import numpy as np
13
import torch
14
-import torch.distributed
14
import torch.nn as nn
15
16
import vllm.envs as envs
@@ -32,14 +31,12 @@ from vllm.distributed.kv_transfer import (
31
)
32
from vllm.distributed.parallel_state import (
33
Handle,
35
- get_pcp_group,
34
get_pp_group,
35
get_tp_group,
36
)
37
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
38
from vllm.logger import init_logger
39
from vllm.lora.request import LoRARequest
42
-from vllm.model_executor.models.interfaces import is_mixture_of_experts
40
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
41
from vllm.platforms import current_platform
42
from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper
@@ -49,7 +46,6 @@ from vllm.tracing import instrument
46
from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
47
from vllm.utils.torch_utils import set_random_seed
48
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
52
-from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
49
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
50
from vllm.v1.outputs import (
51
AsyncModelRunnerOutput,
@@ -124,6 +120,10 @@ class Worker(WorkerBase):
120
precision = envs.VLLM_FLOAT32_MATMUL_PRECISION
121
torch.set_float32_matmul_precision(precision)
122
123
+ from vllm.distributed.elastic_ep.elastic_execute import ElasticEPScalingExecutor
124
+
125
+ self.elastic_ep_executor = ElasticEPScalingExecutor(self)
126
+
127
# Buffers saved before sleep
128
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
129
@@ -317,12 +317,29 @@ class Worker(WorkerBase):
317
# FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool
318
# to hijack tensor allocation.
319
def load_model(self) -> None:
320
- eep_scale_up = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"
320
+ dummy_weights = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"
321
+ if dummy_weights:
322
+ (
323
+ expanded_physical_to_logical,
324
+ num_logical_experts,
325
+ old_num_physical_experts,
326
+ ) = self.elastic_ep_executor.receive_expert_mapping()
327
+ num_physical_experts = expanded_physical_to_logical.shape[1]
328
+ self.parallel_config.eplb_config.num_redundant_experts = (
329
+ num_physical_experts - num_logical_experts
330
+ )
331
+
332
with (
333
self._maybe_get_memory_pool_context(tag="weights"),
334
set_current_vllm_config(self.vllm_config),
335
):
325
- self.model_runner.load_model(eep_scale_up=eep_scale_up)
336
+ self.model_runner.load_model(load_dummy_weights=dummy_weights)
337
+
338
+ if dummy_weights:
339
+ self.model_runner.setup_eplb_from_mapping(
340
+ expanded_physical_to_logical, old_num_physical_experts
341
+ )
342
+ self.model_runner.eep_eplb_suppressed = True
343
344
def update_config(self, overrides: dict[str, Any]) -> None:
345
self.model_runner.update_config(overrides)
@@ -801,227 +818,6 @@ class Worker(WorkerBase):
818
# worker will always be healthy as long as it's running.
819
return
820
804
- def _eplb_before_scale_down(self, old_ep_size: int, new_ep_size: int) -> None:
805
- from vllm.distributed.parallel_state import get_ep_group
806
-
807
- if get_ep_group().rank == 0:
808
- logger.info(
809
- "[Elastic EP] Starting expert resharding before scaling down..."
810
- )
811
- rank_mapping = {
812
- old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
813
- for old_ep_rank in range(old_ep_size)
814
- }
815
- assert self.model_runner.eplb_state is not None
816
- self.model_runner.eplb_state.rearrange(
817
- execute_shuffle=True,
818
- global_expert_loads=None,
819
- rank_mapping=rank_mapping,
820
- )
821
- torch.cuda.synchronize()
822
- if get_ep_group().rank == 0:
823
- logger.info("[Elastic EP] Expert resharding completed!")
824
-
825
- def _eplb_after_scale_up(
826
- self,
827
- old_ep_size: int,
828
- new_ep_size: int,
829
- global_expert_loads: list[torch.Tensor] | None,
830
- ) -> None:
831
- from vllm.distributed.parallel_state import get_ep_group
832
-
833
- if get_ep_group().rank == 0:
834
- logger.info("[Elastic EP] Starting expert resharding after scaling up...")
835
- rank_mapping = {old_ep_rank: old_ep_rank for old_ep_rank in range(old_ep_size)}
836
- assert self.model_runner.eplb_state is not None
837
- self.model_runner.eplb_state.rearrange(
838
- execute_shuffle=True,
839
- global_expert_loads=global_expert_loads,
840
- rank_mapping=rank_mapping,
841
- )
842
- if get_ep_group().rank == 0:
843
- logger.info("[Elastic EP] Expert resharding completed!")
844
-
845
- def _reconfigure_parallel_config(
846
- self, reconfig_request: ReconfigureDistributedRequest
847
- ) -> None:
848
- """
849
- Update parallel config with provided reconfig_request
850
- """
851
- parallel_config = self.vllm_config.parallel_config
852
- parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size
853
- if (
854
- reconfig_request.new_data_parallel_rank
855
- != ReconfigureRankType.KEEP_CURRENT_RANK
856
- ):
857
- parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
858
- if (
859
- reconfig_request.new_data_parallel_rank_local
860
- != ReconfigureRankType.KEEP_CURRENT_RANK
861
- ):
862
- parallel_config.data_parallel_rank_local = (
863
- reconfig_request.new_data_parallel_rank_local
864
- )
865
- parallel_config.data_parallel_master_ip = (
866
- reconfig_request.new_data_parallel_master_ip
867
- )
868
- parallel_config.data_parallel_master_port = (
869
- reconfig_request.new_data_parallel_master_port
870
- )
871
-
872
- def _reconfigure_moe(
873
- self, old_ep_size: int, new_ep_size: int
874
- ) -> list[torch.Tensor] | None:
875
- """
876
- Reconfigure MoE modules with provided reconfig_request
877
-
878
- Return the global expert load if new_ep_size > old_ep_size,
879
- otherwise None
880
- """
881
- from vllm.distributed.parallel_state import (
882
- get_dp_group,
883
- get_ep_group,
884
- prepare_communication_buffer_for_model,
885
- )
886
- from vllm.model_executor.layers.fused_moe.layer import (
887
- FusedMoE,
888
- FusedMoEParallelConfig,
889
- )
890
-
891
- parallel_config = self.vllm_config.parallel_config
892
-
893
- def get_moe_modules(model: torch.nn.Module) -> list[FusedMoE]:
894
- return [
895
- module
896
- for module in model.modules()
897
- if (
898
- module.__class__.__name__ == "FusedMoE"
899
- or module.__class__.__name__ == "SharedFusedMoE"
900
- )
901
- ]
902
-
903
- def update_moe_modules(moe_modules: list[FusedMoE], num_local_experts: int):
904
- assert all(
905
- module.moe_config.num_local_experts == num_local_experts
906
- for module in moe_modules
907
- ), "All MoE modules must have the same number of experts"
908
- for module in moe_modules:
909
- module.moe_config.num_experts = num_local_experts * new_ep_size
910
- module.global_num_experts = module.moe_config.num_experts
911
- tp_size = get_tp_group().world_size
912
- is_sequence_parallel = parallel_config.use_sequence_parallel_moe
913
- sp_size = tp_size if is_sequence_parallel else 1
914
- module.moe_parallel_config = FusedMoEParallelConfig.make(
915
- tp_size_=tp_size,
916
- pcp_size_=get_pcp_group().world_size,
917
- dp_size_=get_dp_group().world_size,
918
- sp_size_=sp_size,
919
- vllm_parallel_config=parallel_config,
920
- )
921
- module.moe_config.moe_parallel_config = module.moe_parallel_config
922
- return moe_modules
923
-
924
- model_moe_modules = get_moe_modules(self.model_runner.model)
925
- num_local_experts = model_moe_modules[0].moe_config.num_local_experts
926
-
927
- update_moe_modules(model_moe_modules, num_local_experts)
928
- drafter_model = None
929
- if hasattr(self.model_runner, "drafter") and hasattr(
930
- self.model_runner.drafter, "model"
931
- ):
932
- drafter_model = self.model_runner.drafter.model
933
- if drafter_model is not None and is_mixture_of_experts(drafter_model):
934
- drafter_moe_modules = get_moe_modules(drafter_model)
935
- # Check if drafter and model have matching configs
936
- assert (
937
- drafter_moe_modules[0].moe_config.num_local_experts == num_local_experts
938
- ), "Drafter and model configs should be the same"
939
- update_moe_modules(drafter_moe_modules, num_local_experts)
940
-
941
- if new_ep_size < old_ep_size:
942
- num_local_physical_experts = num_local_experts
943
- assert self.model_runner.eplb_state is not None
944
- new_physical_experts = (
945
- self.model_runner.eplb_state.physical_to_logical_map.shape[1] # type: ignore[attr-defined]
946
- )
947
- parallel_config.eplb_config.num_redundant_experts = (
948
- new_physical_experts
949
- - self.model_runner.eplb_state.logical_replica_count.shape[1] # type: ignore[attr-defined]
950
- )
951
- global_expert_loads = None
952
- else:
953
- num_local_physical_experts_tensor = torch.tensor(
954
- [num_local_experts], dtype=torch.int32, device="cpu"
955
- )
956
- torch.distributed.broadcast(
957
- num_local_physical_experts_tensor,
958
- group=get_ep_group().cpu_group,
959
- group_src=0,
960
- )
961
- num_local_physical_experts = int(num_local_physical_experts_tensor.item())
962
- new_physical_experts = num_local_physical_experts * new_ep_size
963
- assert self.model_runner.eplb_state is not None
964
- global_expert_loads_any = self.model_runner.eplb_state.rearrange(
965
- execute_shuffle=False
966
- )
967
- global_expert_loads = cast(list[torch.Tensor], global_expert_loads_any)
968
- parallel_config.eplb_config.num_redundant_experts = (
969
- new_physical_experts - global_expert_loads[0].shape[1]
970
- )
971
- prepare_communication_buffer_for_model(self.model_runner.model)
972
- if drafter_model is not None:
973
- prepare_communication_buffer_for_model(drafter_model)
974
- self.model_runner.model.update_physical_experts_metadata(
975
- num_physical_experts=new_physical_experts,
976
- num_local_physical_experts=num_local_physical_experts,
977
- )
978
- return global_expert_loads
979
-
980
- def reinitialize_distributed(
981
- self, reconfig_request: ReconfigureDistributedRequest
982
- ) -> None:
983
- from vllm.config import set_current_vllm_config
984
- from vllm.distributed.parallel_state import (
985
- cleanup_dist_env_and_memory,
986
- get_ep_group,
987
- )
988
-
989
- old_ep_size = get_ep_group().world_size
990
- old_ep_rank = get_ep_group().rank
991
- new_ep_size = (
992
- reconfig_request.new_data_parallel_size
993
- * get_tp_group().world_size
994
- * get_pp_group().world_size
995
- )
996
- if new_ep_size < old_ep_size:
997
- self._eplb_before_scale_down(old_ep_size, new_ep_size)
998
-
999
- cleanup_dist_env_and_memory()
1000
-
1001
- if (
1002
- reconfig_request.new_data_parallel_rank
1003
- == ReconfigureRankType.SHUTDOWN_CURRENT_RANK
1004
- ):
1005
- assert old_ep_rank >= new_ep_size
1006
- # shutdown
1007
- return
1008
-
1009
- self._reconfigure_parallel_config(reconfig_request)
1010
-
1011
- with set_current_vllm_config(self.vllm_config):
1012
- init_worker_distributed_environment(
1013
- self.vllm_config,
1014
- self.rank,
1015
- self.distributed_init_method,
1016
- self.local_rank,
1017
- )
1018
-
1019
- global_expert_loads = self._reconfigure_moe(old_ep_size, new_ep_size)
1020
-
1021
- if new_ep_size > old_ep_size:
1022
- assert global_expert_loads is not None
1023
- self._eplb_after_scale_up(old_ep_size, new_ep_size, global_expert_loads)
1024
-
821
def save_sharded_state(
822
self,
823
path: str,
@@ -1118,6 +914,9 @@ class Worker(WorkerBase):
914
if weight_transfer_engine := getattr(self, "weight_transfer_engine", None):
915
weight_transfer_engine.shutdown()
916
917
+ def elastic_ep_execute(self, execute_method: str, *args, **kwargs):
918
+ return self.elastic_ep_executor.execute(execute_method, *args, **kwargs)
919
+
920
921
def init_worker_distributed_environment(
922
vllm_config: VllmConfig,
vllm/v1/worker/workspace.py
+28
@@ -66,6 +66,23 @@ class WorkspaceManager:
66
],
67
)
68
69
+ def unlock(self) -> None:
70
+ """Unlock the workspace to allow growth.
71
+
72
+ This is used during elastic EP scaling when the workspace size
73
+ needs to grow due to changes in the number of experts.
74
+ """
75
+ self._locked = False
76
+ if envs.VLLM_DEBUG_WORKSPACE:
77
+ logger.info(
78
+ "[WORKSPACE DEBUG] Workspace unlocked. Current sizes: %s",
79
+ [
80
+ self._workspace_size_bytes(ws) / _MB
81
+ for ws in self._current_workspaces
82
+ if ws is not None
83
+ ],
84
+ )
85
+
86
def is_locked(self) -> bool:
87
"""Check if workspace is locked."""
88
return self._locked
@@ -242,6 +259,17 @@ def lock_workspace() -> None:
259
current_workspace_manager().lock()
260
261
262
+def unlock_workspace() -> None:
263
+ """Unlock the workspace to allow growth.
264
+
265
+ This is used during elastic EP scaling when the workspace size
266
+ needs to grow due to changes in the number of experts.
267
+ After scaling operations complete, lock_workspace() should be
268
+ called again to prevent unexpected allocations.
269
+ """
270
+ current_workspace_manager().unlock()
271
+
272
+
273
def reset_workspace_manager() -> None:
274
"""Reset the workspace manager to uninitialized state.
275