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 15, 04:45 UTC
[Feature] Expert Parallelism Load Balancer (EPLB) (#18343)
Signed-off-by: Bowen Wang <abmfy@icloud.com>
Bowen Wang committed
Jun 26, 2025 at 15:30 UTC
e9fd658a736a4d30f7a367c317506c87ad7f5359
24 files changed
+2446
-54
.buildkite/test-pipeline.yaml
+17
@@ -168,6 +168,23 @@ steps:
168
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
169
- popd
170
171
+- label: EPLB Algorithm Test
172
+ working_dir: "/vllm-workspace/tests"
173
+ source_file_dependencies:
174
+ - vllm/distributed/eplb
175
+ - tests/distributed/test_eplb_algo.py
176
+ commands:
177
+ - pytest -v -s distributed/test_eplb_algo.py
178
+
179
+- label: EPLB Execution Test # 5min
180
+ working_dir: "/vllm-workspace/tests"
181
+ num_gpus: 4
182
+ source_file_dependencies:
183
+ - vllm/distributed/eplb
184
+ - tests/distributed/test_eplb_execute.py
185
+ commands:
186
+ - pytest -v -s distributed/test_eplb_execute.py
187
+
188
- label: Metrics, Tracing Test # 10min
189
mirror_hardwares: [amdexperimental, amdproduction]
190
num_gpus: 2
tests/distributed/test_eplb_algo.py
new
+292
@@ -0,0 +1,292 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.distributed.eplb.rebalance_algo import rebalance_experts
8
+
9
+
10
+def test_basic_rebalance():
11
+ """Test basic rebalancing functionality"""
12
+ # Example from https://github.com/deepseek-ai/eplb
13
+ weight = torch.tensor([
14
+ [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
15
+ [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
16
+ ])
17
+
18
+ num_layers = weight.shape[0]
19
+ num_replicas = 16
20
+ num_groups = 4
21
+ num_nodes = 2
22
+ num_gpus = 8
23
+
24
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
25
+ num_groups, num_nodes,
26
+ num_gpus)
27
+
28
+ # Verify output shapes
29
+ assert phy2log.shape == (
30
+ 2,
31
+ 16,
32
+ ), f"Expected `phy2log` shape (2, 16), got {phy2log.shape}"
33
+ assert (log2phy.shape[0] == 2
34
+ ), f"Expected `log2phy` first dimension 2, got {log2phy.shape[0]}"
35
+ assert (
36
+ log2phy.shape[1] == 12
37
+ ), f"Expected `log2phy` second dimension 12, got {log2phy.shape[1]}"
38
+ assert logcnt.shape == (
39
+ 2,
40
+ 12,
41
+ ), f"Expected `logcnt` shape (2, 12), got {logcnt.shape}"
42
+
43
+ # Verify physical to logical expert mapping range is correct
44
+ assert torch.all(phy2log >= 0) and torch.all(
45
+ phy2log < 12), "Physical to logical mapping should be in range [0, 12)"
46
+
47
+ # Verify expert count reasonableness
48
+ assert torch.all(
49
+ logcnt >= 1), "Each logical expert should have at least 1 replica"
50
+ assert (
51
+ torch.sum(logcnt, dim=1).sum() == num_replicas *
52
+ num_layers), f"Total replicas should be {num_replicas * num_layers}"
53
+
54
+ # Verify expected output
55
+ expected_phy2log = torch.tensor([
56
+ [5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1],
57
+ [7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1],
58
+ ])
59
+ assert torch.all(phy2log == expected_phy2log)
60
+
61
+ expected_logcnt = torch.tensor([[1, 2, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1],
62
+ [1, 2, 1, 1, 1, 2, 2, 1, 2, 1, 1, 1]])
63
+ assert torch.all(logcnt == expected_logcnt)
64
+
65
+
66
+def test_single_gpu_case():
67
+ """Test single GPU case"""
68
+ weight = torch.tensor([[10, 20, 30, 40]])
69
+ num_replicas = 4
70
+ num_groups = 1
71
+ num_nodes = 1
72
+ num_gpus = 1
73
+
74
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
75
+ num_groups, num_nodes,
76
+ num_gpus)
77
+
78
+ # Verify shapes
79
+ assert phy2log.shape == (1, 4)
80
+ assert log2phy.shape[0] == 1
81
+ assert log2phy.shape[1] == 4
82
+ assert logcnt.shape == (1, 4)
83
+
84
+ # Verify all logical experts are mapped
85
+ assert set(phy2log[0].tolist()) == {0, 1, 2, 3}
86
+
87
+
88
+def test_equal_weights():
89
+ """Test case with equal weights"""
90
+ weight = torch.tensor([[50, 50, 50, 50, 50, 50, 50, 50]])
91
+ num_replicas = 8
92
+ num_groups = 2
93
+ num_nodes = 2
94
+ num_gpus = 4
95
+
96
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
97
+ num_groups, num_nodes,
98
+ num_gpus)
99
+
100
+ # Verify shapes
101
+ assert phy2log.shape == (1, 8)
102
+ assert logcnt.shape == (1, 8)
103
+
104
+ # With equal weights, each expert should have exactly one replica
105
+ assert torch.all(
106
+ logcnt == 1
107
+ ), "With equal weights and no replication, " \
108
+ "each expert should have exactly 1 replica"
109
+
110
+
111
+def test_extreme_weight_imbalance():
112
+ """Test extreme weight imbalance case"""
113
+ weight = torch.tensor([[1000, 1, 1, 1, 1, 1, 1, 1]])
114
+ num_replicas = 12
115
+ num_groups = 2
116
+ num_nodes = 2
117
+ num_gpus = 4
118
+
119
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
120
+ num_groups, num_nodes,
121
+ num_gpus)
122
+
123
+ # Verify shapes
124
+ assert phy2log.shape == (1, 12)
125
+ assert logcnt.shape == (1, 8)
126
+
127
+ # Expert with highest weight (index 0) should have more replicas
128
+ assert (
129
+ logcnt[0, 0]
130
+ > logcnt[0, 1]), "Expert with highest weight should have more replicas"
131
+
132
+
133
+def test_multiple_layers():
134
+ """Test multiple layers case"""
135
+ weight = torch.tensor([
136
+ [10, 20, 30, 40, 50, 60], # First layer
137
+ [60, 50, 40, 30, 20, 10], # Second layer (opposite weight pattern)
138
+ [25, 25, 25, 25, 25, 25], # Third layer (equal weights)
139
+ ])
140
+ num_replicas = 8
141
+ num_groups = 2
142
+ num_nodes = 2
143
+ num_gpus = 4
144
+
145
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
146
+ num_groups, num_nodes,
147
+ num_gpus)
148
+
149
+ # Verify shapes
150
+ assert phy2log.shape == (3, 8)
151
+ assert logcnt.shape == (3, 6)
152
+
153
+ # Verify expert allocation is reasonable for each layer
154
+ for layer in range(3):
155
+ assert torch.all(phy2log[layer] >= 0) and torch.all(
156
+ phy2log[layer] < 6
157
+ ), f"Layer {layer} physical to logical mapping" \
158
+ "should be in range [0, 6)"
159
+ assert (torch.sum(logcnt[layer]) == num_replicas
160
+ ), f"Layer {layer} total replicas should be {num_replicas}"
161
+
162
+
163
+def test_parameter_validation():
164
+ """Test parameter validation"""
165
+ weight = torch.tensor([[10, 20, 30, 40]])
166
+
167
+ # Test non-divisible case - this should handle normally without throwing
168
+ # errors because the function will fall back to global load balancing
169
+ # strategy
170
+ phy2log, log2phy, logcnt = rebalance_experts(weight, 8, 3, 2, 4)
171
+ assert phy2log.shape == (1, 8)
172
+ assert logcnt.shape == (1, 4)
173
+
174
+ # Test cases that will actually cause errors:
175
+ # num_physical_experts not divisible by num_gpus
176
+ with pytest.raises(AssertionError):
177
+ rebalance_experts(weight, 7, 2, 2, 4) # 7 not divisible by 4
178
+
179
+
180
+def test_small_scale_hierarchical():
181
+ """Test small-scale hierarchical load balancing"""
182
+ weight = torch.tensor([
183
+ [100, 50, 200, 75, 150, 25, 300, 80], # 8 experts
184
+ ])
185
+ num_replicas = 12
186
+ num_groups = 4 # 4 groups, 2 experts each
187
+ num_nodes = 2 # 2 nodes
188
+ num_gpus = 4 # 4 GPUs
189
+
190
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
191
+ num_groups, num_nodes,
192
+ num_gpus)
193
+
194
+ # Verify basic constraints
195
+ assert phy2log.shape == (1, 12)
196
+ assert logcnt.shape == (1, 8)
197
+ assert torch.sum(logcnt) == num_replicas
198
+ assert torch.all(logcnt >= 1)
199
+
200
+ # Expert with highest weight should have more replicas
201
+ max_weight_expert = torch.argmax(weight[0])
202
+ assert (logcnt[0, max_weight_expert]
203
+ >= 2), "Highest weight expert should have multiple replicas"
204
+
205
+
206
+def test_global_load_balance_fallback():
207
+ """Test global load balancing fallback case"""
208
+ # When num_groups % num_nodes != 0, should fall back to global load
209
+ # balancing
210
+ weight = torch.tensor([[10, 20, 30, 40, 50, 60]])
211
+ num_replicas = 8
212
+ num_groups = 3 # Cannot be divided evenly by num_nodes=2
213
+ num_nodes = 2
214
+ num_gpus = 4
215
+
216
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
217
+ num_groups, num_nodes,
218
+ num_gpus)
219
+
220
+ # Should work normally, just using global load balancing strategy
221
+ assert phy2log.shape == (1, 8)
222
+ assert logcnt.shape == (1, 6)
223
+ assert torch.sum(logcnt) == num_replicas
224
+
225
+
226
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
227
+def test_device_compatibility(device):
228
+ """Test device compatibility"""
229
+ if device == "cuda" and not torch.cuda.is_available():
230
+ pytest.skip("CUDA not available")
231
+
232
+ weight = torch.tensor([[10, 20, 30, 40]], device=device)
233
+ num_replicas = 6
234
+ num_groups = 2
235
+ num_nodes = 1
236
+ num_gpus = 2
237
+
238
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
239
+ num_groups, num_nodes,
240
+ num_gpus)
241
+
242
+ # Function will convert to CPU internally, but should handle different
243
+ # device inputs normally
244
+ assert phy2log.shape == (1, 6)
245
+ assert logcnt.shape == (1, 4)
246
+
247
+
248
+def test_additional_cases():
249
+ """Test more edge cases and different parameter combinations"""
250
+
251
+ # Test case 1: Large-scale distributed setup
252
+ weight1 = torch.tensor(
253
+ [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]])
254
+ phy2log1, log2phy1, logcnt1 = rebalance_experts(weight1, 24, 8, 4, 8)
255
+
256
+ assert phy2log1.shape == (1, 24)
257
+ assert logcnt1.shape == (1, 16)
258
+ assert torch.sum(logcnt1) == 24
259
+
260
+ # Test case 2: Different weight distributions
261
+ weight2 = torch.tensor([
262
+ [200, 150, 100, 50, 25, 12], # Decreasing weights
263
+ [12, 25, 50, 100, 150, 200], # Increasing weights
264
+ ])
265
+ phy2log2, log2phy2, logcnt2 = rebalance_experts(weight2, 10, 3, 1, 2)
266
+
267
+ assert phy2log2.shape == (2, 10)
268
+ assert logcnt2.shape == (2, 6)
269
+
270
+ # Verify high-weight experts have more replicas
271
+ for layer in range(2):
272
+ max_weight_idx = torch.argmax(weight2[layer])
273
+ assert logcnt2[layer, max_weight_idx] >= 2
274
+
275
+
276
+if __name__ == "__main__":
277
+ weight = torch.tensor([
278
+ [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
279
+ [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
280
+ ])
281
+
282
+ num_replicas = 16
283
+ num_groups = 4
284
+ num_nodes = 2
285
+ num_gpus = 8
286
+
287
+ phy2log, log2phy, logcnt = rebalance_experts(weight, num_replicas,
288
+ num_groups, num_nodes,
289
+ num_gpus)
290
+ print(phy2log)
291
+
292
+ test_basic_rebalance()
tests/distributed/test_eplb_execute.py
new
+504
@@ -0,0 +1,504 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import multiprocessing
5
+import os
6
+import random
7
+
8
+import pytest
9
+import torch
10
+import torch.distributed
11
+
12
+from vllm.distributed.eplb.rebalance_execute import (
13
+ rearrange_expert_weights_inplace)
14
+from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
15
+ get_tp_group,
16
+ init_distributed_environment)
17
+from vllm.utils import update_environment_variables
18
+
19
+
20
+def distributed_run(fn, world_size):
21
+ number_of_processes = world_size
22
+ processes: list[multiprocessing.Process] = []
23
+ for i in range(number_of_processes):
24
+ env: dict[str, str] = {}
25
+ env['RANK'] = str(i)
26
+ env['LOCAL_RANK'] = str(i)
27
+ env['WORLD_SIZE'] = str(number_of_processes)
28
+ env['LOCAL_WORLD_SIZE'] = str(number_of_processes)
29
+ env['MASTER_ADDR'] = 'localhost'
30
+ env['MASTER_PORT'] = '12345'
31
+ p = multiprocessing.Process(target=fn, args=(env, ))
32
+ processes.append(p)
33
+ p.start()
34
+
35
+ for p in processes:
36
+ p.join()
37
+
38
+ for p in processes:
39
+ assert p.exitcode == 0
40
+
41
+
42
+def worker_fn_wrapper(fn):
43
+ # `multiprocessing.Process` cannot accept environment variables directly
44
+ # so we need to pass the environment variables as arguments
45
+ # and update the environment variables in the function
46
+ def wrapped_fn(env):
47
+ update_environment_variables(env)
48
+ local_rank = os.environ['LOCAL_RANK']
49
+ device = torch.device(f"cuda:{local_rank}")
50
+ torch.cuda.set_device(device)
51
+ init_distributed_environment()
52
+
53
+ # Ensure each worker process has the same random seed
54
+ random.seed(42)
55
+ torch.manual_seed(42)
56
+
57
+ fn()
58
+
59
+ return wrapped_fn
60
+
61
+
62
+def create_expert_indices_with_redundancy(
63
+ num_layers: int,
64
+ num_logical_experts: int,
65
+ total_physical_experts: int,
66
+ redundancy_config: list[int], # redundancy for each logical expert
67
+) -> torch.Tensor:
68
+ """
69
+ Create expert indices with redundancy.
70
+
71
+ Args:
72
+ num_layers: number of layers
73
+ num_logical_experts: number of logical experts
74
+ total_physical_experts: total number of physical experts
75
+ redundancy_config: redundancy for each logical expert
76
+
77
+ Returns:
78
+ indices: Shape (num_layers, total_physical_experts)
79
+ """
80
+ assert sum(redundancy_config) == total_physical_experts
81
+ assert len(redundancy_config) == num_logical_experts
82
+
83
+ indices = torch.zeros(num_layers, total_physical_experts, dtype=torch.long)
84
+
85
+ for layer in range(num_layers):
86
+ physical_pos = 0
87
+ for logical_expert_id, redundancy in enumerate(redundancy_config):
88
+ for _ in range(redundancy):
89
+ indices[layer, physical_pos] = logical_expert_id
90
+ physical_pos += 1
91
+
92
+ # Shuffle the indices at dim 1
93
+ for layer in range(num_layers):
94
+ indices[layer] = indices[layer][torch.randperm(indices.shape[1])]
95
+
96
+ return indices
97
+
98
+
99
+def create_expert_weights(
100
+ num_layers: int,
101
+ num_local_experts: int,
102
+ hidden_sizes: list[int],
103
+ rank: int,
104
+ device: torch.device,
105
+ physical_to_logical_mapping: torch.Tensor,
106
+) -> list[list[torch.Tensor]]:
107
+ """
108
+ Create fake expert weights tensor for testing.
109
+
110
+ Use `arange` to generate predictable weights values, based on logical
111
+ expert ID.
112
+ All replicas of the same logical expert should have the same weights.
113
+
114
+ Args:
115
+ physical_to_logical_mapping: Shape (num_layers, num_local_experts)
116
+ mapping[layer, physical_pos] = logical_expert_id
117
+ """
118
+ expert_weights = []
119
+
120
+ for layer in range(num_layers):
121
+ layer_weights = []
122
+ for weight_idx, hidden_size in enumerate(hidden_sizes):
123
+ weight_tensor = torch.zeros(num_local_experts,
124
+ hidden_size,
125
+ device=device,
126
+ dtype=torch.float32)
127
+
128
+ for local_expert in range(num_local_experts):
129
+ # Get the logical expert ID for this physical expert
130
+ global_pos = rank * num_local_experts + local_expert
131
+ logical_expert_id = physical_to_logical_mapping[
132
+ layer, global_pos].item()
133
+
134
+ # Generate weights based on logical expert ID
135
+ # (so that all replicas of the same logical expert have the
136
+ # same weights)
137
+ base_value = (logical_expert_id * 1000 + layer * 100 +
138
+ weight_idx * 10)
139
+ weight_tensor[local_expert] = torch.arange(base_value,
140
+ base_value +
141
+ hidden_size,
142
+ device=device,
143
+ dtype=torch.float32)
144
+
145
+ layer_weights.append(weight_tensor)
146
+ expert_weights.append(layer_weights)
147
+
148
+ return expert_weights
149
+
150
+
151
+def create_redundancy_config(
152
+ num_logical_experts: int,
153
+ num_physical_experts: int,
154
+) -> list[int]:
155
+ """Create a redundancy configuration."""
156
+ redundancy_config = [1] * num_logical_experts
157
+ remaining = num_physical_experts - num_logical_experts
158
+ # Randomly assign the remaining physical experts to the logical experts
159
+ for _ in range(remaining):
160
+ redundancy_config[random.choice(range(num_logical_experts))] += 1
161
+ return redundancy_config
162
+
163
+
164
+def verify_expert_weights_after_shuffle(
165
+ expert_weights: list[list[torch.Tensor]],
166
+ new_indices: torch.Tensor,
167
+ hidden_sizes: list[int],
168
+ ep_rank: int,
169
+ num_local_experts: int,
170
+):
171
+ """Verify the weights after shuffling are correct."""
172
+ num_layers = len(expert_weights)
173
+
174
+ for layer in range(num_layers):
175
+ for weight_idx, hidden_size in enumerate(hidden_sizes):
176
+ weight_tensor = expert_weights[layer][weight_idx]
177
+
178
+ for local_expert in range(num_local_experts):
179
+ # Calculate the global expert ID for this local expert
180
+ global_pos = ep_rank * num_local_experts + local_expert
181
+ expected_logical_expert = new_indices[layer, global_pos].item()
182
+
183
+ # Check if the weights are correct
184
+ actual_weights = weight_tensor[local_expert]
185
+ expected_base = (expected_logical_expert * 1000 + layer * 100 +
186
+ weight_idx * 10)
187
+ expected_weights = torch.arange(expected_base,
188
+ expected_base + hidden_size,
189
+ device=actual_weights.device,
190
+ dtype=actual_weights.dtype)
191
+
192
+ torch.testing.assert_close(
193
+ actual_weights,
194
+ expected_weights,
195
+ msg=f"Layer {layer}, weight {weight_idx},"
196
+ f"local expert {local_expert}: "
197
+ f"weights do not match. "
198
+ f"Expected logical expert {expected_logical_expert}")
199
+
200
+
201
+def verify_redundant_experts_have_same_weights(
202
+ expert_weights: list[list[torch.Tensor]],
203
+ indices: torch.Tensor,
204
+ hidden_sizes: list[int],
205
+ world_size: int,
206
+ num_local_experts: int,
207
+):
208
+ """
209
+ Verify that all replicas of the same logical expert have the same weights.
210
+ """
211
+ num_layers = len(expert_weights)
212
+ total_physical_experts = world_size * num_local_experts
213
+
214
+ for layer in range(num_layers):
215
+ # Collect weights for all physical experts for each weight matrix
216
+ all_weights: list[torch.Tensor] = []
217
+
218
+ for weight_idx, hidden_size in enumerate(hidden_sizes):
219
+ # Create tensor to store all expert weights
220
+ # Shape: [total_physical_experts, hidden_size]
221
+ gathered_weights = torch.zeros(
222
+ total_physical_experts,
223
+ hidden_size,
224
+ device=expert_weights[layer][weight_idx].device,
225
+ dtype=expert_weights[layer][weight_idx].dtype)
226
+
227
+ # Use all_gather to collect expert weights from current node
228
+ # expert_weights[layer][weight_idx] shape:
229
+ # [num_local_experts, hidden_size]
230
+ local_weights = expert_weights[layer][
231
+ weight_idx] # [num_local_experts, hidden_size]
232
+
233
+ # Split tensor along dim 0 into a list for all_gather
234
+ gathered_weights_list = torch.chunk(gathered_weights,
235
+ world_size,
236
+ dim=0)
237
+
238
+ torch.distributed.all_gather(
239
+ # Output list: each element corresponds to one rank's weights
240
+ list(gathered_weights_list),
241
+ local_weights # Input: current rank's local weights
242
+ )
243
+
244
+ all_weights.append(gathered_weights)
245
+
246
+ # Verify that all replicas of the same logical expert have the same
247
+ # weights
248
+ logical_expert_weights: dict[int, dict[int, torch.Tensor]] = {}
249
+
250
+ for physical_pos in range(total_physical_experts):
251
+ logical_expert_id = int(indices[layer, physical_pos].item())
252
+
253
+ if logical_expert_id not in logical_expert_weights:
254
+ # First time encountering this logical expert, save its weights
255
+ logical_expert_weights[logical_expert_id] = {
256
+ weight_idx: all_weights[weight_idx][physical_pos]
257
+ for weight_idx in range(len(hidden_sizes))
258
+ }
259
+ else:
260
+ # Verify that current physical expert's weights match the
261
+ # previously saved logical expert weights
262
+ for weight_idx in range(len(hidden_sizes)):
263
+ torch.testing.assert_close(
264
+ all_weights[weight_idx][physical_pos],
265
+ logical_expert_weights[logical_expert_id][weight_idx],
266
+ msg=f"Layer {layer}, weight {weight_idx},"
267
+ f"logical expert {logical_expert_id}: "
268
+ f"Physical expert {physical_pos} has different weights"
269
+ f"than expected")
270
+
271
+
272
+@pytest.mark.parametrize(
273
+ "world_size,num_layers,num_local_experts,num_logical_experts",
274
+ [
275
+ # 2 GPU, 2 experts per GPU
276
+ # 3 logical experts, 4 physical experts, 1 redundant experts
277
+ (2, 1, 2, 3),
278
+ # 2 GPU, 3 experts per GPU
279
+ # 4 logical experts, 6 physical experts, 2 redundant experts
280
+ (2, 2, 3, 4),
281
+ # 2 GPU, 8 experts per GPU
282
+ # 16 logical experts, 16 physical experts, 0 redundant experts
283
+ (2, 4, 8, 16),
284
+ # 4 GPU, 2 experts per GPU
285
+ # 6 logical experts, 8 physical experts, 2 redundant experts
286
+ (4, 1, 2, 6),
287
+ # 4 GPU, 2 experts per GPU
288
+ # 5 logical experts, 8 physical experts, 3 redundant experts
289
+ (4, 2, 2, 5),
290
+ # 4 GPU, 8 experts per GPU
291
+ # 16 logical experts, 32 physical experts, 16 redundant experts
292
+ (4, 8, 8, 16),
293
+ ])
294
+def test_rearrange_expert_weights_with_redundancy(world_size, num_layers,
295
+ num_local_experts,
296
+ num_logical_experts):
297
+ """Test the functionality of rearranging expert weights with redundancy."""
298
+
299
+ if torch.cuda.device_count() < world_size:
300
+ pytest.skip(f"Need at least {world_size} GPUs to run the test")
301
+
302
+ @worker_fn_wrapper
303
+ def worker_fn():
304
+ # Initialize model parallel (using tensor parallel as an entrypoint
305
+ # to expert parallel)
306
+ ensure_model_parallel_initialized(
307
+ tensor_model_parallel_size=world_size,
308
+ pipeline_model_parallel_size=1)
309
+
310
+ ep_group = get_tp_group().cpu_group
311
+ ep_rank = torch.distributed.get_rank()
312
+ device = torch.device(f"cuda:{ep_rank}")
313
+
314
+ # Test parameters
315
+ total_physical_experts = world_size * num_local_experts
316
+ hidden_sizes = [32, 64] # Two different weight matrices
317
+
318
+ # Create old expert indices (with redundancy)
319
+ redundancy_config = create_redundancy_config(num_logical_experts,
320
+ total_physical_experts)
321
+
322
+ old_indices = create_expert_indices_with_redundancy(
323
+ num_layers,
324
+ num_logical_experts,
325
+ total_physical_experts,
326
+ redundancy_config,
327
+ )
328
+
329
+ # Create new expert indices (with redundancy)
330
+ new_redundancy_config = create_redundancy_config(
331
+ num_logical_experts, total_physical_experts)
332
+ new_indices = create_expert_indices_with_redundancy(
333
+ num_layers,
334
+ num_logical_experts,
335
+ total_physical_experts,
336
+ new_redundancy_config,
337
+ )
338
+
339
+ # Create expert weights
340
+ expert_weights = create_expert_weights(num_layers, num_local_experts,
341
+ hidden_sizes, ep_rank, device,
342
+ old_indices)
343
+
344
+ # Execute weight rearrangement
345
+ rearrange_expert_weights_inplace(
346
+ old_indices,
347
+ new_indices,
348
+ expert_weights,
349
+ ep_group,
350
+ is_profile=False,
351
+ )
352
+
353
+ # Verify the rearrangement result
354
+ verify_expert_weights_after_shuffle(
355
+ expert_weights,
356
+ new_indices,
357
+ hidden_sizes,
358
+ ep_rank,
359
+ num_local_experts,
360
+ )
361
+
362
+ verify_redundant_experts_have_same_weights(
363
+ expert_weights,
364
+ new_indices,
365
+ hidden_sizes,
366
+ world_size,
367
+ num_local_experts,
368
+ )
369
+
370
+ distributed_run(worker_fn, world_size)
371
+
372
+
373
+@pytest.mark.parametrize("world_size", [2, 4])
374
+def test_rearrange_expert_weights_no_change(world_size):
375
+ """
376
+ Test that when the indices do not change, the weights should remain
377
+ unchanged.
378
+ """
379
+
380
+ if torch.cuda.device_count() < world_size:
381
+ pytest.skip(f"Need at least {world_size} GPUs to run the test")
382
+
383
+ @worker_fn_wrapper
384
+ def worker_fn():
385
+ ensure_model_parallel_initialized(
386
+ tensor_model_parallel_size=world_size,
387
+ pipeline_model_parallel_size=1)
388
+
389
+ ep_group = get_tp_group().cpu_group
390
+ ep_rank = torch.distributed.get_rank()
391
+ device = torch.device(f"cuda:{ep_rank}")
392
+
393
+ num_layers = 2
394
+ num_local_experts = 2
395
+ total_physical_experts = world_size * num_local_experts
396
+ num_logical_experts = total_physical_experts // 2 # Some redundancy
397
+ hidden_sizes = [32, 64]
398
+
399
+ # Create redundancy configuration
400
+ redundancy_config = [2] * num_logical_experts
401
+
402
+ # Same indices - no change
403
+ indices = create_expert_indices_with_redundancy(
404
+ num_layers, num_logical_experts, total_physical_experts,
405
+ redundancy_config)
406
+
407
+ expert_weights = create_expert_weights(num_layers, num_local_experts,
408
+ hidden_sizes, ep_rank, device,
409
+ indices)
410
+
411
+ # Save original weights
412
+ original_weights = []
413
+ for layer_weights in expert_weights:
414
+ layer_copy = []
415
+ for weight in layer_weights:
416
+ layer_copy.append(weight.clone())
417
+ original_weights.append(layer_copy)
418
+
419
+ # Execute rearrangement (should be no change)
420
+ rearrange_expert_weights_inplace(
421
+ indices,
422
+ indices, # Same indices
423
+ expert_weights,
424
+ ep_group,
425
+ is_profile=False)
426
+
427
+ # Verify that the weights have not changed
428
+ for layer in range(num_layers):
429
+ for weight_idx in range(len(hidden_sizes)):
430
+ torch.testing.assert_close(
431
+ expert_weights[layer][weight_idx],
432
+ original_weights[layer][weight_idx],
433
+ msg=f"Layer {layer}, weight {weight_idx} should remain "
434
+ f"unchanged")
435
+
436
+ distributed_run(worker_fn, world_size)
437
+
438
+
439
+@pytest.mark.parametrize("world_size", [2, 4])
440
+def test_rearrange_expert_weights_profile_mode(world_size):
441
+ """Test profile mode (should not copy actual weights)"""
442
+
443
+ if torch.cuda.device_count() < world_size:
444
+ pytest.skip(f"Need at least {world_size} GPUs to run the test")
445
+
446
+ @worker_fn_wrapper
447
+ def worker_fn():
448
+ ensure_model_parallel_initialized(
449
+ tensor_model_parallel_size=world_size,
450
+ pipeline_model_parallel_size=1)
451
+
452
+ ep_group = get_tp_group().cpu_group
453
+ ep_rank = torch.distributed.get_rank()
454
+ device = torch.device(f"cuda:{ep_rank}")
455
+
456
+ num_layers = 1
457
+ num_local_experts = 2
458
+ total_physical_experts = world_size * num_local_experts
459
+ num_logical_experts = total_physical_experts // 2
460
+ hidden_sizes = [32]
461
+
462
+ # Create different index distributions
463
+ old_redundancy = create_redundancy_config(num_logical_experts,
464
+ total_physical_experts)
465
+ new_redundancy = create_redundancy_config(num_logical_experts,
466
+ total_physical_experts)
467
+
468
+ old_indices = create_expert_indices_with_redundancy(
469
+ num_layers, num_logical_experts, total_physical_experts,
470
+ old_redundancy)
471
+ new_indices = create_expert_indices_with_redundancy(
472
+ num_layers, num_logical_experts, total_physical_experts,
473
+ new_redundancy)
474
+
475
+ expert_weights = create_expert_weights(num_layers, num_local_experts,
476
+ hidden_sizes, ep_rank, device,
477
+ old_indices)
478
+
479
+ # Save original weights
480
+ original_weights = []
481
+ for layer_weights in expert_weights:
482
+ layer_copy = []
483
+ for weight in layer_weights:
484
+ layer_copy.append(weight.clone())
485
+ original_weights.append(layer_copy)
486
+
487
+ # Execute profile mode rearrangement
488
+ rearrange_expert_weights_inplace(
489
+ old_indices,
490
+ new_indices,
491
+ expert_weights,
492
+ ep_group,
493
+ is_profile=True # Profile mode
494
+ )
495
+
496
+ # In profile mode, the weights should remain unchanged
497
+ for layer in range(num_layers):
498
+ for weight_idx in range(len(hidden_sizes)):
499
+ torch.testing.assert_close(
500
+ expert_weights[layer][weight_idx],
501
+ original_weights[layer][weight_idx],
502
+ msg="In profile mode, the weights should remain unchanged")
503
+
504
+ distributed_run(worker_fn, world_size)
tests/models/test_initialization.py
+10
-2
@@ -31,12 +31,20 @@ def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):
31
32
text_config = hf_config.get_text_config()
33
34
+ # Ensure at least 2 expert per group
35
+ # Since `grouped_topk` assums top-2
36
+ num_experts = getattr(text_config, 'n_group', 1) * 2
37
+
38
text_config.update({
39
"num_layers": 1,
40
"num_hidden_layers": 1,
37
- "num_experts": 2,
41
+ "num_experts": num_experts,
42
"num_experts_per_tok": 2,
39
- "num_local_experts": 2,
43
+ "num_local_experts": num_experts,
44
+ # Otherwise there will not be any expert layers
45
+ "first_k_dense_replace": 0,
46
+ # To avoid OOM on DeepSeek-V3
47
+ "n_routed_experts": num_experts,
48
})
49
50
if hasattr(hf_config, "vision_config"):
vllm/config.py
+33
@@ -1775,6 +1775,25 @@ class ParallelConfig:
1775
"""Backend to use for data parallel, either "mp" or "ray"."""
1776
enable_expert_parallel: bool = False
1777
"""Use expert parallelism instead of tensor parallelism for MoE layers."""
1778
+ enable_eplb: bool = False
1779
+ """Enable expert parallelism load balancing for MoE layers."""
1780
+ num_redundant_experts: int = 0
1781
+ """Number of redundant experts to use for expert parallelism."""
1782
+ eplb_window_size: int = 1000
1783
+ """Window size for expert load recording."""
1784
+ eplb_step_interval: int = 3000
1785
+ """
1786
+ Interval for rearranging experts in expert parallelism.
1787
+
1788
+ Note that if this is greater than the EPLB window size, only the metrics
1789
+ of the last `eplb_window_size` steps will be used for rearranging experts.
1790
+ """
1791
+ eplb_log_balancedness: bool = False
1792
+ """
1793
+ Log the balancedness each step of expert parallelism.
1794
+ This is turned off by default since it will cause communication overhead.
1795
+ """
1796
+
1797
max_parallel_loading_workers: Optional[int] = None
1798
"""Maximum number of parallel loading workers when loading model
1799
sequentially in multiple batches. To avoid RAM OOM when using tensor
@@ -1913,6 +1932,20 @@ class ParallelConfig:
1932
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
1933
logger.info("Disabling V1 multiprocessing for external launcher.")
1934
1935
+ if self.enable_eplb:
1936
+ if not current_platform.is_cuda():
1937
+ raise ValueError(
1938
+ "Expert parallelism load balancing is only supported on "
1939
+ "CUDA devices now.")
1940
+ if self.num_redundant_experts < 0:
1941
+ raise ValueError(
1942
+ "num_redundant_experts must be non-negative, but got "
1943
+ f"{self.num_redundant_experts}.")
1944
+ else:
1945
+ if self.num_redundant_experts != 0:
1946
+ raise ValueError(
1947
+ "num_redundant_experts should be used with EPLB."
1948
+ f"{self.num_redundant_experts}.")
1949
if self.distributed_executor_backend is None and self.world_size > 1:
1950
# We use multiprocessing by default if world_size fits on the
1951
# current node and we aren't in a ray placement group.
vllm/distributed/eplb/__init__.py
new
+7
@@ -0,0 +1,7 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+'''
3
+Expert parallelism load balancer (EPLB).
4
+'''
5
+
6
+from .eplb_state import *
7
+from .rebalance_algo import *
vllm/distributed/eplb/eplb_state.py
new
+431
@@ -0,0 +1,431 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Expert parallelism load balancer (EPLB) metrics and states.
4
+
5
+# Glossary
6
+
7
+- **Logical Expert**: An expert that is part of the model's logical structure.
8
+ It holds a set of weights and is replicated across multiple physical
9
+ experts.
10
+- **Redundant Expert**: To achieve load balancing, for some popular logical
11
+ experts, we create additional copies of the expert weights. During inference,
12
+ each of these copies can be routed to by the same set of tokens.
13
+- **Physical Expert**: An expert that is instantiated on a specific device.
14
+ It is a replica of a logical expert and can be rearranged across devices.
15
+ I.e., one logical expert may have multiple sets of weights initialized on
16
+ different devices, and each of these sets is a physical expert.
17
+- **Local Physical Expert**: A physical expert that is instantiated on the
18
+ current device.
19
+
20
+For example: DeepSeek-R1 has 256 logical experts, so each MoE layer
21
+has 256 sets of linear layer weights in the model parameters. If we add 32
22
+redundant experts, DeepSeek-R1 will have 256 + 32 = 288 physical experts in
23
+total. And when deploying, we'll have 288 sets of linear layer weights for each
24
+MoE layer. If we have 32 EP ranks, then each GPU will hold 288 / 32 = 9 local
25
+physical experts.
26
+"""
27
+
28
+import time
29
+from collections.abc import Sequence
30
+from dataclasses import dataclass
31
+
32
+import torch
33
+from torch.distributed import all_gather, all_reduce
34
+
35
+from vllm.config import ParallelConfig
36
+from vllm.distributed.parallel_state import get_ep_group, get_node_count
37
+from vllm.logger import init_logger
38
+from vllm.model_executor.models.interfaces import MixtureOfExperts
39
+
40
+from .rebalance_algo import rebalance_experts
41
+from .rebalance_execute import rearrange_expert_weights_inplace
42
+
43
+logger = init_logger(__name__)
44
+
45
+
46
+@dataclass
47
+class EplbState:
48
+ """EPLB metrics."""
49
+
50
+ physical_to_logical_map: torch.Tensor
51
+ """
52
+ Mapping from physical experts to logical experts.
53
+
54
+ Shape: (num_moe_layers, num_physical_experts)
55
+
56
+ # Example
57
+
58
+ For a 2-layer MoE model with 6 physical experts and 4 logical experts on 3
59
+ EP ranks, the mapping could look like this:
60
+
61
+ ```
62
+ [[0, 1, 2, 3, 0, 1],
63
+ [0, 2, 0, 1, 0, 3]]
64
+ ```
65
+ """
66
+ logical_to_physical_map: torch.Tensor
67
+ """
68
+ Mapping from logical experts to physical experts.
69
+
70
+ This is a sparse matrix, where -1 indicates no mapping.
71
+
72
+ Shape: (num_moe_layers, num_logical_experts, num_redundant_experts + 1)
73
+
74
+ # Example
75
+
76
+ For a 2-layer MoE model with 6 physical experts and 4 logical experts on 3
77
+ EP ranks, the mapping could look like this:
78
+
79
+ ```
80
+ [[[0, 4, -1],
81
+ [1, 5, -1],
82
+ [2, -1, -1],
83
+ [3, -1, -1]],
84
+ [[0, 2, 4],
85
+ [3, -1, -1],
86
+ [1, -1, -1],
87
+ [5, -1, -1]]]
88
+ ```
89
+ """
90
+ logical_replica_count: torch.Tensor
91
+ """
92
+ Number of replicas for each logical expert.
93
+ This is exactly the non-`-1` count in the `logical_to_physical_map`.
94
+
95
+ Shape: (num_moe_layers, num_logical_experts)
96
+
97
+ # Example
98
+ For a 2-layer MoE model with 6 physical experts and 4 logical experts on 3
99
+ EP ranks, the count could look like this:
100
+
101
+ ```
102
+ [[2, 2, 1, 1],
103
+ [3, 1, 1, 1]]
104
+ """
105
+
106
+ expert_load_pass: torch.Tensor
107
+ """
108
+ Expert load during this forward pass.
109
+ We use the token count each expert processes as the load.
110
+
111
+ Shape: (num_moe_layers, num_local_physical_experts)
112
+ """
113
+ expert_load_window: torch.Tensor
114
+ """
115
+ A sliding window of expert load.
116
+
117
+ Shape: (window_size, num_moe_layers, num_local_physical_experts)
118
+ """
119
+ expert_load_window_step: int = 0
120
+ """
121
+ Current step in the sliding window.
122
+
123
+ Different from `expert_rearrangement_step`, each EP rank may have its own
124
+ `expert_load_window_step`.
125
+ """
126
+ expert_load_window_size: int = 0
127
+ """
128
+ Size of the expert load sliding window.
129
+ This is a constant and is taken from the config.
130
+ """
131
+
132
+ expert_rearrangement_step: int = 0
133
+ """
134
+ Steps after last rearrangement.
135
+ Will trigger a rearrangement if it exceeds the threshold.
136
+
137
+ NOTE: Keep in mind that all EP ranks need to have the same
138
+ `expert_rearrangement_step` value to ensure synchronization.
139
+ Otherwise, the rearrangement will hang at collective
140
+ communication calls.
141
+ """
142
+ expert_rearrangement_step_interval: int = 0
143
+ """
144
+ Interval for expert rearrangement steps.
145
+ This is a constant and is taken from the config.
146
+ """
147
+
148
+ @staticmethod
149
+ def build_initial_global_physical_to_logical_map(
150
+ num_routed_experts: int,
151
+ num_redundant_experts: int,
152
+ ) -> Sequence[int]:
153
+ """
154
+ Build an initial expert arrangement using the following structure:
155
+ [original routed experts, redundant experts]
156
+
157
+ Returns:
158
+ physical_to_logical_map (Sequence[int]): A list of integers,
159
+ where each integer is the index of the logical expert
160
+ that the corresponding physical expert maps to.
161
+ """
162
+ global_physical_to_logical_map = list(range(num_routed_experts))
163
+ global_physical_to_logical_map += [
164
+ i % num_routed_experts for i in range(num_redundant_experts)
165
+ ]
166
+ return global_physical_to_logical_map
167
+
168
+ @classmethod
169
+ def build(
170
+ cls,
171
+ model: MixtureOfExperts,
172
+ device: torch.device,
173
+ parallel_config: ParallelConfig,
174
+ ) -> "EplbState":
175
+ """
176
+ Build the initial EPLB state.
177
+ """
178
+ physical_to_logical_map_list = (
179
+ cls.build_initial_global_physical_to_logical_map(
180
+ model.num_routed_experts,
181
+ model.num_redundant_experts,
182
+ ))
183
+ physical_to_logical_map = torch.tensor(
184
+ physical_to_logical_map_list,
185
+ device=device,
186
+ )
187
+ logical_to_physical_map = torch.full(
188
+ (model.num_logical_experts, model.num_redundant_experts + 1),
189
+ -1,
190
+ device=device,
191
+ )
192
+ logical_replica_count = torch.zeros(
193
+ (model.num_logical_experts, ),
194
+ device=device,
195
+ dtype=torch.long,
196
+ )
197
+
198
+ for i in range(model.num_physical_experts):
199
+ logical_idx = physical_to_logical_map[i]
200
+ logical_to_physical_map[logical_idx,
201
+ logical_replica_count[logical_idx]] = i
202
+ logical_replica_count[logical_idx] += 1
203
+
204
+ # Duplicate initial mapping for all layers
205
+ physical_to_logical_map = physical_to_logical_map.unsqueeze(0).expand(
206
+ model.num_moe_layers,
207
+ -1,
208
+ ).contiguous()
209
+ logical_to_physical_map = logical_to_physical_map.unsqueeze(0).expand(
210
+ model.num_moe_layers,
211
+ -1,
212
+ -1,
213
+ ).contiguous()
214
+ logical_replica_count = logical_replica_count.unsqueeze(0).expand(
215
+ model.num_moe_layers,
216
+ -1,
217
+ ).contiguous()
218
+
219
+ expert_load_pass = torch.zeros(
220
+ (model.num_moe_layers, model.num_local_physical_experts),
221
+ dtype=torch.int32,
222
+ device=device,
223
+ )
224
+ expert_load_window_size = parallel_config.eplb_window_size
225
+ expert_load_window = torch.zeros(
226
+ (expert_load_window_size, model.num_moe_layers,
227
+ model.num_local_physical_experts),
228
+ dtype=torch.int32,
229
+ device=device,
230
+ )
231
+
232
+ # Set the initial progress of rearrangement to 3/4
233
+ eplb_step_interval = parallel_config.eplb_step_interval
234
+ expert_rearrangement_step = max(
235
+ 0, eplb_step_interval - eplb_step_interval // 4)
236
+
237
+ model.set_eplb_state(
238
+ expert_load_pass,
239
+ logical_to_physical_map,
240
+ logical_replica_count,
241
+ )
242
+
243
+ return cls(
244
+ physical_to_logical_map,
245
+ logical_to_physical_map,
246
+ logical_replica_count,
247
+ expert_load_pass,
248
+ expert_load_window,
249
+ expert_load_window_size=expert_load_window_size,
250
+ expert_rearrangement_step=expert_rearrangement_step,
251
+ expert_rearrangement_step_interval=eplb_step_interval,
252
+ )
253
+
254
+ def step(self,
255
+ model: MixtureOfExperts,
256
+ is_dummy: bool = False,
257
+ is_profile: bool = False,
258
+ log_stats: bool = False) -> None:
259
+ """
260
+ Step the EPLB state.
261
+
262
+ Args:
263
+ model (MixtureOfExperts): The MoE model.
264
+ is_dummy (bool): If `True`, this is a dummy step and the load
265
+ metrics recorded in this forward pass will not count. Defaults
266
+ to `False`.
267
+ is_profile (bool): If `True`, perform a dummy rearrangement
268
+ with maximum communication cost. This is used in `profile_run`
269
+ to reserve enough memory for the communication buffer.
270
+ log_stats (bool): If `True`, log the expert load metrics.
271
+
272
+ # Stats
273
+ The metrics are all summed up across layers.
274
+ - `avg_tokens`: The average load across ranks.
275
+ - `max_tokens`: The maximum load across ranks.
276
+ - `balancedness`: The ratio of average load to maximum load.
277
+ """
278
+
279
+ if is_profile:
280
+ self.rearrange(model, is_profile=True)
281
+ return
282
+
283
+ if is_dummy:
284
+ # Do not record load metrics for dummy steps
285
+ self.expert_load_pass.zero_()
286
+
287
+ if log_stats:
288
+ # `num_tokens`: (num_moe_layers,)
289
+ num_tokens = self.expert_load_pass.sum(dim=-1)
290
+
291
+ # Collect load metrics from all ranks
292
+ ep_group = get_ep_group().device_group
293
+ num_tokens_list = [
294
+ torch.empty_like(num_tokens) for _ in range(ep_group.size())
295
+ ]
296
+ all_gather(num_tokens_list, num_tokens, group=ep_group)
297
+ # Stack to get (num_ranks, num_moe_layers)
298
+ num_tokens_per_rank = torch.stack(num_tokens_list).float()
299
+
300
+ # Compute balancedness ratio:
301
+ # for each layer:
302
+ # (mean load across ranks) / (max load across ranks)
303
+ avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
304
+ max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(
305
+ dim=0)
306
+
307
+ # Just to make type checker happy
308
+ tokens_tensors: list[float] = torch.stack(
309
+ [avg_tokens_tensor, max_tokens_tensor]).tolist()
310
+ avg_tokens, max_tokens = tokens_tensors
311
+ balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
312
+
313
+ if ep_group.rank() == 0:
314
+ logger.info(
315
+ "EPLB step: avg_tokens=%.2f, max_tokens=%d, "
316
+ "balancedness=%.4f", avg_tokens, max_tokens, balancedness)
317
+
318
+ # Update the expert load sliding window
319
+ if not is_dummy:
320
+ self.expert_load_window[self.expert_load_window_step] = (
321
+ self.expert_load_pass.clone())
322
+ self.expert_load_window_step += 1
323
+ if self.expert_load_window_step >= self.expert_load_window_size:
324
+ self.expert_load_window_step = 0
325
+ self.expert_load_pass.zero_()
326
+
327
+ # Step the expert rearrangement step
328
+ # Note that even if this is a dummy step, we still increment the
329
+ # rearrangement step and perform rearrangement to ensure all ranks are
330
+ # performing collective communication.
331
+ self.expert_rearrangement_step += 1
332
+ if (self.expert_rearrangement_step
333
+ >= self.expert_rearrangement_step_interval):
334
+ self.expert_rearrangement_step = 0
335
+ self.rearrange(model)
336
+
337
+ def rearrange(self,
338
+ model: MixtureOfExperts,
339
+ is_profile: bool = False) -> None:
340
+ """
341
+ Rearrange the experts according to the current load.
342
+ """
343
+
344
+ ep_group = get_ep_group().device_group
345
+ ep_rank = ep_group.rank()
346
+
347
+ time_start = None
348
+ is_main_rank = ep_rank == 0
349
+ if is_main_rank:
350
+ torch.cuda.synchronize()
351
+ time_start = time.perf_counter()
352
+ logger.info("Rearranging experts %s...",
353
+ "(profile)" if is_profile else "")
354
+
355
+ # This mapping is only used here, so we do not store it in the state
356
+ physical_expert_start = ep_rank * model.num_local_physical_experts
357
+ physical_expert_end = (physical_expert_start +
358
+ model.num_local_physical_experts)
359
+ # (num_moe_layers, num_local_physical_experts)
360
+ local_physical_to_logical_map = self.physical_to_logical_map[
361
+ :,
362
+ physical_expert_start:physical_expert_end,
363
+ ]
364
+
365
+ # Map the local physical expert load to global logical experts
366
+ logical_expert_load_window = torch.zeros(
367
+ self.expert_load_window_size,
368
+ model.num_moe_layers,
369
+ model.num_logical_experts,
370
+ dtype=self.expert_load_window.dtype,
371
+ device=self.expert_load_window.device,
372
+ )
373
+ logical_expert_load_window.scatter_add_(
374
+ dim=-1,
375
+ index=local_physical_to_logical_map.unsqueeze(0).expand_as(
376
+ self.expert_load_window).long(),
377
+ src=self.expert_load_window,
378
+ )
379
+
380
+ # Perform all-reduce to get the expert load across all ranks
381
+ global_expert_load_window = logical_expert_load_window.sum(dim=0)
382
+ all_reduce(global_expert_load_window, group=ep_group)
383
+
384
+ # TODO(bowen): Treat differently for prefill and decode nodes
385
+ num_replicas = model.num_physical_experts
386
+ num_groups = model.num_expert_groups
387
+ num_nodes = get_node_count()
388
+ num_gpus = ep_group.size()
389
+
390
+ if num_gpus % num_nodes != 0:
391
+ logger.warning_once(
392
+ f"num_gpus % num_nodes != 0, "
393
+ "not using hierarchical rearrangement algorithm.\n"
394
+ f"{num_gpus=}, {num_nodes=}")
395
+
396
+ # Get new expert mappings
397
+ (
398
+ new_physical_to_logical_map,
399
+ new_logical_to_physical_map,
400
+ new_logical_replica_count,
401
+ ) = (rebalance_experts(
402
+ global_expert_load_window,
403
+ num_replicas,
404
+ num_groups,
405
+ num_nodes,
406
+ num_gpus,
407
+ ))
408
+
409
+ # Update expert weights
410
+ rearrange_expert_weights_inplace(
411
+ self.physical_to_logical_map,
412
+ new_physical_to_logical_map,
413
+ model.expert_weights,
414
+ ep_group,
415
+ is_profile,
416
+ )
417
+
418
+ if not is_profile:
419
+ self.physical_to_logical_map.copy_(new_physical_to_logical_map)
420
+ self.logical_to_physical_map.copy_(new_logical_to_physical_map)
421
+ self.logical_replica_count.copy_(new_logical_replica_count)
422
+
423
+ if is_main_rank:
424
+ assert time_start is not None
425
+ torch.cuda.synchronize()
426
+ time_end = time.perf_counter()
427
+ logger.info(
428
+ "Rearranged experts%sin %.2f seconds.",
429
+ " (profile) " if is_profile else " ",
430
+ time_end - time_start,
431
+ )
vllm/distributed/eplb/rebalance_algo.py
new
+233
@@ -0,0 +1,233 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Expert parallelism load balancer (EPLB) for vLLM.
4
+
5
+This module implements the core rearrangement algorithm.
6
+
7
+The rearrangement algorithm is adapted from
8
+[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).
9
+
10
+Please find at [#12](https://github.com/deepseek-ai/EPLB/issues/12) an example
11
+on how the EPLB algorithm works.
12
+"""
13
+
14
+import torch
15
+
16
+
17
+def balanced_packing(weight: torch.Tensor,
18
+ num_packs: int) -> tuple[torch.Tensor, torch.Tensor]:
19
+ """
20
+ Pack n weighted objects to m packs, such that each bin contains exactly
21
+ n/m objects and the weights of all packs are as balanced as possible.
22
+
23
+ Parameters:
24
+ weight: [X, n], the weight of each item
25
+ num_packs: number of packs
26
+
27
+ Returns:
28
+ pack_index: [X, n], the pack index of each item
29
+ rank_in_pack: [X, n], the rank of the item in the pack
30
+ """
31
+ num_layers, num_groups = weight.shape
32
+ assert num_groups % num_packs == 0
33
+ groups_per_pack = num_groups // num_packs
34
+
35
+ if groups_per_pack == 1:
36
+ pack_index = torch.arange(weight.size(-1),
37
+ dtype=torch.int64,
38
+ device=weight.device).expand(weight.shape)
39
+ rank_in_pack = torch.zeros_like(weight, dtype=torch.int64)
40
+ return pack_index, rank_in_pack
41
+
42
+ indices = weight.float().sort(-1, descending=True).indices.cpu()
43
+ pack_index = torch.full_like(weight,
44
+ fill_value=-1,
45
+ dtype=torch.int64,
46
+ device="cpu")
47
+ rank_in_pack = torch.full_like(pack_index, fill_value=-1)
48
+ for i in range(num_layers):
49
+ pack_weights = [0] * num_packs
50
+ pack_items = [0] * num_packs
51
+ for group in indices[i]:
52
+ pack = min(
53
+ (i
54
+ for i in range(num_packs) if pack_items[i] < groups_per_pack),
55
+ key=pack_weights.__getitem__,
56
+ )
57
+ assert pack_items[pack] < groups_per_pack
58
+ pack_index[i, group] = pack
59
+ rank_in_pack[i, group] = pack_items[pack]
60
+ pack_weights[pack] += weight[i, group]
61
+ pack_items[pack] += 1
62
+ return pack_index, rank_in_pack
63
+
64
+
65
+def replicate_experts(
66
+ weight: torch.Tensor,
67
+ num_phy: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
68
+ """
69
+ Replicate `num_log` experts to `num_phy` replicas, such that the maximum
70
+ load of all replicas is minimized.
71
+
72
+ Parameters:
73
+ weight: [X, num_log]
74
+ num_phy: total number of experts after replication
75
+
76
+ Returns:
77
+ phy2log: [X, num_phy], logical expert id of each physical expert
78
+ rank: [X, num_phy], the replica rank
79
+ logcnt: [X, num_log], number of replicas for each logical expert
80
+ """
81
+ n, num_log = weight.shape
82
+ num_redundant = num_phy - num_log
83
+ assert num_redundant >= 0
84
+ device = weight.device
85
+ phy2log = torch.arange(num_phy, dtype=torch.int64,
86
+ device=device).repeat(n, 1)
87
+ rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
88
+ logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
89
+ arangen = torch.arange(n, dtype=torch.int64, device=device)
90
+ for i in range(num_log, num_phy):
91
+ redundant_indices = (weight / logcnt).max(dim=-1).indices
92
+ phy2log[:, i] = redundant_indices
93
+ rank[:, i] = logcnt[arangen, redundant_indices]
94
+ logcnt[arangen, redundant_indices] += 1
95
+ return phy2log, rank, logcnt
96
+
97
+
98
+def rebalance_experts_hierarchical(
99
+ weight: torch.Tensor,
100
+ num_physical_experts: int,
101
+ num_groups: int,
102
+ num_nodes: int,
103
+ num_gpus: int,
104
+):
105
+ """
106
+ Parameters:
107
+ weight: [num_moe_layers, num_logical_experts]
108
+ num_physical_experts: number of physical experts after replication
109
+ num_groups: number of expert groups
110
+ num_nodes: number of server nodes, where the intra-node network
111
+ (e.g, NVLink) is faster
112
+ num_gpus: number of GPUs, must be a multiple of `num_nodes`
113
+
114
+ Returns:
115
+ physical_to_logical_map: [num_moe_layers, num_physical_experts]
116
+ logical_to_physical_map: [num_moe_layers, num_logical_experts, X]
117
+ logical_count: [num_moe_layers, num_logical_experts]
118
+ """
119
+ num_layers, num_logical_experts = weight.shape
120
+ assert num_logical_experts % num_groups == 0
121
+ group_size = num_logical_experts // num_groups
122
+ assert num_groups % num_nodes == 0
123
+ groups_per_node = num_groups // num_nodes
124
+ assert num_gpus % num_nodes == 0
125
+ assert num_physical_experts % num_gpus == 0
126
+ phy_experts_per_gpu = num_physical_experts // num_gpus
127
+
128
+ def inverse(perm: torch.Tensor) -> torch.Tensor:
129
+ inv = torch.empty_like(perm)
130
+ inv.scatter_(
131
+ 1,
132
+ perm,
133
+ torch.arange(perm.size(1), dtype=torch.int64,
134
+ device=perm.device).expand(perm.shape),
135
+ )
136
+ return inv
137
+
138
+ # Step 1: pack groups to nodes
139
+ tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
140
+ group_pack_index, group_rank_in_pack = balanced_packing(
141
+ tokens_per_group, num_nodes)
142
+ log2mlog = (((group_pack_index * groups_per_node + group_rank_in_pack) *
143
+ group_size).unsqueeze(-1) +
144
+ torch.arange(group_size,
145
+ dtype=torch.int64,
146
+ device=group_pack_index.device)).flatten(-2)
147
+ mlog2log = inverse(log2mlog)
148
+
149
+ # Step 2: construct redundant experts within nodes
150
+ # [num_layers * num_nodes, num_logical_experts // num_nodes]
151
+ tokens_per_mlog = weight.gather(-1, mlog2log).view(
152
+ -1, num_logical_experts // num_nodes)
153
+ phy2mlog, phyrank, mlogcnt = replicate_experts(
154
+ tokens_per_mlog, num_physical_experts // num_nodes)
155
+
156
+ # Step 3: pack physical_experts to GPUs
157
+ # [num_layers * num_nodes, num_physical_experts // num_nodes]
158
+ tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
159
+ pack_index, rank_in_pack = balanced_packing(tokens_per_phy,
160
+ num_gpus // num_nodes)
161
+ phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
162
+ pphy2phy = inverse(phy2pphy)
163
+
164
+ pphy2mlog = phy2mlog.gather(
165
+ -1, pphy2phy) # [num_layers * num_nodes, num_log_per_nodes]
166
+ pphy2mlog = (pphy2mlog.view(num_layers, num_nodes, -1) + torch.arange(
167
+ 0,
168
+ num_logical_experts,
169
+ num_logical_experts // num_nodes,
170
+ device=group_pack_index.device,
171
+ ).view(1, -1, 1)).flatten(-2)
172
+ pphy2log = mlog2log.gather(-1, pphy2mlog)
173
+ pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
174
+ logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
175
+ return pphy2log, pphyrank, logcnt
176
+
177
+
178
+def rebalance_experts(
179
+ weight: torch.Tensor,
180
+ num_replicas: int,
181
+ num_groups: int,
182
+ num_nodes: int,
183
+ num_gpus: int,
184
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
185
+ """
186
+ Entry point for expert-parallelism load balancer.
187
+
188
+ Parameters:
189
+ weight: [layers, num_logical_experts], the load statistics for all
190
+ logical experts
191
+ num_replicas: number of physical experts, must be a multiple of
192
+ `num_gpus`
193
+ num_groups: number of expert groups
194
+ num_nodes: number of server nodes, where the intra-node network
195
+ (e.g, NVLink) is faster
196
+ num_gpus: number of GPUs, must be a multiple of `num_nodes`
197
+
198
+ Returns:
199
+ physical_to_logical_map: [layers, num_replicas], the expert index of
200
+ each replica
201
+ logical_to_physical_map: [layers, num_logical_experts, X], the replica
202
+ indices for each expert
203
+ expert_count: [layers, num_logical_experts], number of physical
204
+ replicas for each logical expert
205
+ """
206
+ num_layers, num_logical_experts = weight.shape
207
+ weight = weight.float().cpu()
208
+ if num_groups % num_nodes == 0:
209
+ # use hierarchical load-balance policy
210
+ phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
211
+ weight, num_replicas, num_groups, num_nodes, num_gpus)
212
+ else:
213
+ # use global load-balance policy
214
+ phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
215
+ weight, num_replicas, 1, 1, num_gpus)
216
+ num_redundant_experts = num_replicas - num_logical_experts
217
+ maxlogcnt = num_redundant_experts + 1
218
+ log2phy: torch.Tensor = torch.full(
219
+ (num_layers, num_logical_experts, maxlogcnt),
220
+ -1,
221
+ dtype=torch.int64,
222
+ device=logcnt.device,
223
+ )
224
+ log2phy.view(num_layers, -1).scatter_(
225
+ -1,
226
+ phy2log * maxlogcnt + phyrank,
227
+ torch.arange(num_replicas, dtype=torch.int64,
228
+ device=log2phy.device).expand(num_layers, -1),
229
+ )
230
+ return phy2log, log2phy, logcnt
231
+
232
+
233
+__all__ = ["rebalance_experts"]
vllm/distributed/eplb/rebalance_execute.py
new
+306
@@ -0,0 +1,306 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+The actual execution of the rearrangement.
4
+
5
+This involves the exchange of expert weights between GPUs.
6
+"""
7
+
8
+from collections.abc import Iterable, MutableSequence, Sequence
9
+from functools import partial
10
+
11
+import torch
12
+from torch.distributed import (P2POp, ProcessGroup, all_gather,
13
+ batch_isend_irecv, get_global_rank)
14
+
15
+
16
+def idx_local_to_global(
17
+ local_idx: int,
18
+ local_cnt: int,
19
+ ep_rank: int,
20
+) -> int:
21
+ """
22
+ Convert a local expert index to a global expert index.
23
+ """
24
+ return ep_rank * local_cnt + local_idx
25
+
26
+
27
+def idx_global_to_local(
28
+ global_idx: int,
29
+ local_cnt: int,
30
+ ep_rank: int,
31
+) -> int:
32
+ """
33
+ Convert a global expert index to a local expert index.
34
+ """
35
+ return global_idx - ep_rank * local_cnt
36
+
37
+
38
+def global_idx_to_rank(
39
+ global_idx: int,
40
+ local_cnt: int,
41
+) -> int:
42
+ """
43
+ Convert a global expert index to a rank index.
44
+ """
45
+ return global_idx // local_cnt
46
+
47
+
48
+def get_ep_ranks_with_expert(
49
+ idx: int,
50
+ num_local_experts: int,
51
+ old_indices: Sequence[int],
52
+ new_indices: Sequence[int],
53
+) -> tuple[MutableSequence[int], MutableSequence[int]]:
54
+ """
55
+ Get the ranks of the experts that need to be exchanged.
56
+
57
+ Args:
58
+ idx: The index of the expert.
59
+ num_local_experts: The number of local experts.
60
+ old_indices: The old indices of the experts.
61
+ new_indices: The new indices of the experts.
62
+
63
+ Returns:
64
+ A tuple of two lists:
65
+ - The ranks of the experts that need to be sent.
66
+ - The ranks of the experts that need to be received.
67
+ """
68
+ global2rank = partial(
69
+ global_idx_to_rank,
70
+ local_cnt=num_local_experts,
71
+ )
72
+
73
+ ranks_to_send: list[int] = []
74
+ ranks_to_recv: list[int] = []
75
+
76
+ for i, e in enumerate(old_indices):
77
+ if e == idx:
78
+ rank = global2rank(i)
79
+ if not ranks_to_send or ranks_to_send[-1] != rank:
80
+ ranks_to_send.append(rank)
81
+
82
+ for i, e in enumerate(new_indices):
83
+ if e == idx:
84
+ rank = global2rank(i)
85
+ if not ranks_to_recv or ranks_to_recv[-1] != rank:
86
+ ranks_to_recv.append(rank)
87
+
88
+ # Remove those ranks that can get this expert locally.
89
+ ranks_to_send_set = set(ranks_to_send)
90
+ ranks_to_recv_actual = [
91
+ rank for rank in ranks_to_recv if rank not in ranks_to_send_set
92
+ ]
93
+
94
+ return ranks_to_send, ranks_to_recv_actual
95
+
96
+
97
+def shuffle_layer(
98
+ num_local_experts: int,
99
+ ep_rank: int,
100
+ old_indices: Sequence[int],
101
+ new_indices: Sequence[int],
102
+ expert_weights: Iterable[torch.Tensor],
103
+ expert_weights_buffer: Sequence[torch.Tensor],
104
+ ep_group: ProcessGroup,
105
+) -> None:
106
+ """
107
+ Perform expert weights rearrangement of one layer.
108
+ """
109
+ local2global = partial(
110
+ idx_local_to_global,
111
+ local_cnt=num_local_experts,
112
+ ep_rank=ep_rank,
113
+ )
114
+
115
+ # 0. Do nothing for experts that did not change.
116
+ is_unchanged = [
117
+ old_indices[local2global(i)] == new_indices[local2global(i)]
118
+ for i in range(num_local_experts)
119
+ ]
120
+
121
+ # 1. Perform weight copy inside the local rank.
122
+ is_received_locally = is_unchanged[:]
123
+ for src in range(num_local_experts):
124
+ src_global = local2global(src)
125
+ for dst in range(num_local_experts):
126
+ dst_global = local2global(dst)
127
+ if is_received_locally[dst]:
128
+ continue
129
+ if old_indices[src_global] == new_indices[dst_global]:
130
+ is_received_locally[dst] = True
131
+ for weight, buffer in zip(expert_weights,
132
+ expert_weights_buffer):
133
+ buffer[dst].copy_(weight[src])
134
+
135
+ p2p_ops: list[P2POp] = []
136
+
137
+ # 2. Initiate sending of weights.
138
+ experts_send_loc: dict[int, int] = {}
139
+ for src in range(num_local_experts):
140
+ expert = old_indices[local2global(src)]
141
+ if expert in experts_send_loc:
142
+ continue
143
+ experts_send_loc[expert] = src
144
+
145
+ # We need to sort here to match send/recv
146
+ for expert, src in sorted(experts_send_loc.items()):
147
+ ranks_to_send, ranks_to_recv = get_ep_ranks_with_expert(
148
+ expert,
149
+ num_local_experts,
150
+ old_indices,
151
+ new_indices,
152
+ )
153
+
154
+ # Calculate the ranks to send by this rank
155
+ num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)
156
+ sender_pos = ranks_to_send.index(ep_rank)
157
+ recv_begin = sender_pos * num_dst_per_sender
158
+ recv_end = recv_begin + num_dst_per_sender
159
+ recv_ranks = ranks_to_recv[recv_begin:recv_end]
160
+
161
+ # Tackle remainders
162
+ remainder_start = len(ranks_to_send) * num_dst_per_sender
163
+ recver_pos = remainder_start + sender_pos
164
+ if recver_pos < len(ranks_to_recv):
165
+ recv_ranks.append(ranks_to_recv[recver_pos])
166
+
167
+ for dst in recv_ranks:
168
+ dst_global = get_global_rank(ep_group, dst)
169
+ p2p_ops += [
170
+ P2POp(
171
+ torch.distributed.isend,
172
+ weight[src],
173
+ dst_global,
174
+ ) for weight in expert_weights
175
+ ]
176
+
177
+ # 3. Initiate receiving of weights.
178
+ experts_recv_loc: dict[int, int] = {}
179
+ for dst in range(num_local_experts):
180
+ if is_received_locally[dst]:
181
+ continue
182
+ expert = new_indices[local2global(dst)]
183
+ if expert in experts_recv_loc:
184
+ continue
185
+ experts_recv_loc[expert] = dst
186
+
187
+ # We need to sort here to match send/recv
188
+ for expert, dst in sorted(experts_recv_loc.items()):
189
+ ranks_to_send, ranks_to_recv = get_ep_ranks_with_expert(
190
+ expert,
191
+ num_local_experts,
192
+ old_indices,
193
+ new_indices,
194
+ )
195
+
196
+ # Calculate the rank to recv by this rank
197
+ num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)
198
+ recver_pos = ranks_to_recv.index(ep_rank)
199
+ remainder_start = len(ranks_to_send) * num_dst_per_sender
200
+ if recver_pos < remainder_start:
201
+ src = ranks_to_send[recver_pos // num_dst_per_sender]
202
+ else:
203
+ src = ranks_to_send[recver_pos - remainder_start]
204
+
205
+ src_global = get_global_rank(ep_group, src)
206
+ p2p_ops += [
207
+ P2POp(
208
+ torch.distributed.irecv,
209
+ weight[dst],
210
+ src_global,
211
+ ) for weight in expert_weights_buffer
212
+ ]
213
+
214
+ # 4. Execute the P2P operations. The real communication happens here.
215
+ if p2p_ops:
216
+ reqs = batch_isend_irecv(p2p_ops)
217
+ for req in reqs:
218
+ req.wait()
219
+
220
+ # 5. Copy the weights from the buffer back to the original weights.
221
+ for dst in range(num_local_experts):
222
+ if is_unchanged[dst]:
223
+ continue
224
+ if is_received_locally[dst]:
225
+ for weight, buffer in zip(expert_weights, expert_weights_buffer):
226
+ weight[dst].copy_(buffer[dst])
227
+ else:
228
+ expert = new_indices[local2global(dst)]
229
+ src = experts_recv_loc[expert]
230
+ for weight, buffer in zip(expert_weights, expert_weights_buffer):
231
+ weight[dst].copy_(buffer[src])
232
+
233
+
234
+def rearrange_expert_weights_inplace(
235
+ old_global_expert_indices: torch.Tensor,
236
+ new_global_expert_indices: torch.Tensor,
237
+ expert_weights: Sequence[Iterable[torch.Tensor]],
238
+ ep_group: ProcessGroup,
239
+ is_profile: bool = False,
240
+) -> None:
241
+ """
242
+ Rearranges the expert weights in place according to the new expert indices.
243
+
244
+ The value of the indices arguments are logical indices of the experts,
245
+ while keys are physical.
246
+
247
+ Args:
248
+ old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
249
+ new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
250
+ expert_weights: A sequence of shape (num_moe_layers)(weight_count)
251
+ of tensors of shape (num_local_physical_experts, hidden_size_i).
252
+ For example, a linear layer may have up and down projection,
253
+ so weight_count = 2. Each weight's hidden size can be different.
254
+ ep_group: The device process group for expert parallelism.
255
+ is_profile (bool): If `True`, do not perform any actual weight copy.
256
+ This is used during profile run, where we only perform dummy
257
+ communications to reserve enough memory for the buffers.
258
+ """
259
+ num_moe_layers, num_physical_experts = old_global_expert_indices.shape
260
+ assert len(expert_weights) == num_moe_layers
261
+
262
+ num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
263
+ assert new_global_expert_indices.shape == (num_moe_layers,
264
+ num_physical_experts)
265
+
266
+ ep_rank = ep_group.rank()
267
+ ep_size = ep_group.size()
268
+ assert num_physical_experts == ep_size * num_local_physical_experts
269
+
270
+ # A buffer to hold the expert weights in one layer during the exchange.
271
+ # NOTE: Currently we assume the same weights across different layers
272
+ # have the same shape.
273
+ expert_weights_buffer = [torch.empty_like(w) for w in expert_weights[0]]
274
+
275
+ if is_profile:
276
+ # Maximum send size is to send all local experts to all ranks,
277
+ # So we use a dummy `all_gather` to reserve enough communication buffer
278
+ for weight, buffer in zip(expert_weights[0], expert_weights_buffer):
279
+ # A `/dev/null`-like buffer to avoid real memory allocation
280
+ dummy_recv_buffer = [buffer for _ in range(ep_size)]
281
+ # NOTE(bowen): Needed this barrier to avoid OOM during actual
282
+ # execution. I'm not very sure why this is needed
283
+ torch.distributed.barrier()
284
+ all_gather(
285
+ dummy_recv_buffer,
286
+ weight,
287
+ group=ep_group,
288
+ )
289
+ return
290
+
291
+ for layer in range(num_moe_layers):
292
+ # NOTE(bowen): We need this synchronize to run, but I don't know why.
293
+ # If you figure out the reason, please let me know -- thank you!
294
+ torch.cuda.synchronize()
295
+ shuffle_layer(
296
+ num_local_physical_experts,
297
+ ep_rank,
298
+ old_global_expert_indices[layer].tolist(),
299
+ new_global_expert_indices[layer].tolist(),
300
+ expert_weights[layer],
301
+ expert_weights_buffer,
302
+ ep_group,
303
+ )
304
+
305
+
306
+__all__ = ["rearrange_expert_weights_inplace"]
vllm/engine/arg_utils.py
+20
@@ -320,6 +320,11 @@ class EngineArgs:
320
data_parallel_rpc_port: Optional[int] = None
321
data_parallel_backend: str = ParallelConfig.data_parallel_backend
322
enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
323
+ enable_eplb: bool = ParallelConfig.enable_eplb
324
+ num_redundant_experts: int = ParallelConfig.num_redundant_experts
325
+ eplb_window_size: int = ParallelConfig.eplb_window_size
326
+ eplb_step_interval: int = ParallelConfig.eplb_step_interval
327
+ eplb_log_balancedness: bool = ParallelConfig.eplb_log_balancedness
328
max_parallel_loading_workers: Optional[
329
int] = ParallelConfig.max_parallel_loading_workers
330
block_size: Optional[BlockSize] = CacheConfig.block_size
@@ -666,6 +671,16 @@ class EngineArgs:
671
parallel_group.add_argument(
672
"--enable-expert-parallel",
673
**parallel_kwargs["enable_expert_parallel"])
674
+ parallel_group.add_argument("--enable-eplb",
675
+ **parallel_kwargs["enable_eplb"])
676
+ parallel_group.add_argument("--num-redundant-experts",
677
+ **parallel_kwargs["num_redundant_experts"])
678
+ parallel_group.add_argument("--eplb-window-size",
679
+ **parallel_kwargs["eplb_window_size"])
680
+ parallel_group.add_argument("--eplb-step-interval",
681
+ **parallel_kwargs["eplb_step_interval"])
682
+ parallel_group.add_argument("--eplb-log-balancedness",
683
+ **parallel_kwargs["eplb_log_balancedness"])
684
parallel_group.add_argument(
685
"--max-parallel-loading-workers",
686
**parallel_kwargs["max_parallel_loading_workers"])
@@ -1135,6 +1150,11 @@ class EngineArgs:
1150
data_parallel_rpc_port=data_parallel_rpc_port,
1151
data_parallel_backend=data_parallel_backend,
1152
enable_expert_parallel=self.enable_expert_parallel,
1153
+ enable_eplb=self.enable_eplb,
1154
+ num_redundant_experts=self.num_redundant_experts,
1155
+ eplb_window_size=self.eplb_window_size,
1156
+ eplb_step_interval=self.eplb_step_interval,
1157
+ eplb_log_balancedness=self.eplb_log_balancedness,
1158
max_parallel_loading_workers=self.max_parallel_loading_workers,
1159
disable_custom_all_reduce=self.disable_custom_all_reduce,
1160
ray_workers_use_nsight=self.ray_workers_use_nsight,
vllm/model_executor/layers/fused_moe/layer.py
+236
-28
@@ -3,9 +3,10 @@
3
4
import importlib
5
from abc import abstractmethod
6
+from collections.abc import Iterable
7
from dataclasses import dataclass
8
from enum import Enum
8
-from typing import Callable, Optional, Union
9
+from typing import Callable, Literal, Optional, Union, overload
10
11
import torch
12
import torch.nn.functional as F
@@ -20,6 +21,7 @@ from vllm.distributed import (get_dp_group, get_ep_group,
21
get_tensor_model_parallel_rank,
22
get_tensor_model_parallel_world_size,
23
tensor_model_parallel_all_reduce)
24
+from vllm.distributed.eplb.eplb_state import EplbState
25
from vllm.forward_context import ForwardContext, get_forward_context
26
from vllm.logger import init_logger
27
from vllm.model_executor.custom_op import CustomOp
@@ -435,6 +437,10 @@ class FusedMoEMethodBase(QuantizeMethodBase):
437
e_score_correction_bias: Optional[torch.Tensor] = None,
438
apply_router_weight_on_input: bool = False,
439
activation: str = "silu",
440
+ enable_eplb: bool = False,
441
+ expert_load_view: Optional[torch.Tensor] = None,
442
+ logical_to_physical_map: Optional[torch.Tensor] = None,
443
+ logical_replica_count: Optional[torch.Tensor] = None,
444
) -> torch.Tensor:
445
raise NotImplementedError
446
@@ -574,7 +580,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
580
e_score_correction_bias: Optional[torch.Tensor] = None,
581
apply_router_weight_on_input: bool = False,
582
activation: str = "silu",
583
+ enable_eplb: bool = False,
584
+ expert_load_view: Optional[torch.Tensor] = None,
585
+ logical_to_physical_map: Optional[torch.Tensor] = None,
586
+ logical_replica_count: Optional[torch.Tensor] = None,
587
) -> torch.Tensor:
588
+ if enable_eplb:
589
+ raise NotImplementedError(
590
+ "EPLB not supported for `UnquantizedFusedMoEMethod` yet.")
591
+
592
return self.forward(
593
x=x,
594
layer=layer,
@@ -821,6 +835,7 @@ class FusedMoE(torch.nn.Module):
835
reduce_results: Whether to all all_reduce on the output of the layer
836
renomalize: Whether to renormalize the logits in the fused_moe kernel
837
quant_config: Quantization configure.
838
+ enable_eplb: Whether to enable expert parallelism load balancer.
839
"""
840
841
def __init__(
@@ -845,6 +860,8 @@ class FusedMoE(torch.nn.Module):
860
e_score_correction_bias: Optional[torch.Tensor] = None,
861
apply_router_weight_on_input: bool = False,
862
activation: str = "silu",
863
+ enable_eplb: bool = False,
864
+ num_redundant_experts: int = 0,
865
):
866
super().__init__()
867
if params_dtype is None:
@@ -860,7 +877,7 @@ class FusedMoE(torch.nn.Module):
877
get_dp_group().world_size),
878
vllm_parallel_config=vllm_config.parallel_config))
879
863
- self.global_num_experts = num_experts
880
+ self.global_num_experts = num_experts + num_redundant_experts
881
882
# For smuggling this layer into the fused moe custom op
883
compilation_config = vllm_config.compilation_config
@@ -869,8 +886,20 @@ class FusedMoE(torch.nn.Module):
886
compilation_config.static_forward_context[prefix] = self
887
self.layer_name = prefix
888
889
+ self.enable_eplb = enable_eplb
890
+ self.expert_load_view: Optional[torch.Tensor] = None
891
+ self.logical_to_physical_map: Optional[torch.Tensor] = None
892
+ self.logical_replica_count: Optional[torch.Tensor] = None
893
+
894
# Determine expert maps
895
if self.use_ep:
896
+ if self.enable_eplb:
897
+ assert self.global_num_experts % self.ep_size == 0, \
898
+ "EPLB currently only supports even distribution of " \
899
+ "experts across ranks."
900
+ else:
901
+ assert num_redundant_experts == 0, \
902
+ "Redundant experts are only supported with EPLB."
903
self.local_num_experts, self.expert_map = determine_expert_map(
904
ep_size=self.ep_size,
905
ep_rank=self.ep_rank,
@@ -937,6 +966,20 @@ class FusedMoE(torch.nn.Module):
966
assert isinstance(quant_method, FusedMoEMethodBase)
967
self.quant_method = quant_method
968
969
+ if self.enable_eplb:
970
+ from vllm.model_executor.layers.quantization.fp8 import (
971
+ Fp8MoEMethod)
972
+ if not isinstance(quant_method, Fp8MoEMethod):
973
+ # TODO: Add support for additional quantization methods.
974
+ # The implementation for other quantization methods does not
975
+ # contain essential differences, but the current quant API
976
+ # design causes duplicated work when extending to new
977
+ # quantization methods, so I'm leaving it for now.
978
+ # If you plan to add support for more quantization methods,
979
+ # please refer to the implementation in `Fp8MoEMethod`.
980
+ raise NotImplementedError("EPLB is only supported for FP8 "
981
+ "quantization for now.")
982
+
983
moe_quant_params = {
984
"num_experts": self.local_num_experts,
985
"hidden_size": hidden_size,
@@ -965,8 +1008,9 @@ class FusedMoE(torch.nn.Module):
1008
dtype=act_dtype,
1009
device=torch.cuda.current_device())
1010
1011
+ # Note here we use `num_experts` which is logical expert count
1012
self.batched_router_logits = torch.zeros(
969
- (envs.VLLM_MOE_DP_CHUNK_SIZE, self.global_num_experts),
1013
+ (envs.VLLM_MOE_DP_CHUNK_SIZE, num_experts),
1014
dtype=act_dtype,
1015
device=torch.cuda.current_device())
1016
@@ -1130,13 +1174,33 @@ class FusedMoE(torch.nn.Module):
1174
return expert_id
1175
return self.expert_map[expert_id].item()
1176
1177
+ @overload
1178
def weight_loader(self, param: torch.nn.Parameter,
1179
loaded_weight: torch.Tensor, weight_name: str,
1135
- shard_id: str, expert_id: int) -> None:
1180
+ shard_id: str, expert_id: int,
1181
+ return_success: Literal[False]) -> None:
1182
+ ...
1183
1184
+ @overload
1185
+ def weight_loader(self, param: torch.nn.Parameter,
1186
+ loaded_weight: torch.Tensor, weight_name: str,
1187
+ shard_id: str, expert_id: int,
1188
+ return_success: Literal[True]) -> bool:
1189
+ ...
1190
+
1191
+ def weight_loader(self,
1192
+ param: torch.nn.Parameter,
1193
+ loaded_weight: torch.Tensor,
1194
+ weight_name: str,
1195
+ shard_id: str,
1196
+ expert_id: int,
1197
+ return_success: bool = False) -> Optional[bool]:
1198
expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
1199
if expert_id == -1:
1139
- return
1200
+ # Failed to load this param since it's not local to this rank
1201
+ return False if return_success else None
1202
+ # Hereafter, `expert_id` is local physical id
1203
+
1204
quant_method_name = self.quant_method.__class__.__name__
1205
# compressed-tensors checkpoints with packed weights are stored flipped
1206
# TODO (mgoin): check self.quant_method.quant_config.quant_format
@@ -1163,7 +1227,7 @@ class FusedMoE(torch.nn.Module):
1227
if is_gguf_weight_type:
1228
param.weight_type = loaded_weight.item()
1229
param.data.copy_(loaded_weight)
1166
- return
1230
+ return True if return_success else None
1231
1232
# is_transposed: if the dim to shard the weight
1233
# should be flipped. Required by GPTQ, compressed-tensors
@@ -1202,7 +1266,7 @@ class FusedMoE(torch.nn.Module):
1266
self._load_single_value(param=param,
1267
loaded_weight=loaded_weight,
1268
expert_id=expert_id)
1205
- return
1269
+ return True if return_success else None
1270
1271
# Case g_idx
1272
if "g_idx" in weight_name:
@@ -1211,7 +1275,7 @@ class FusedMoE(torch.nn.Module):
1275
loaded_weight=loaded_weight,
1276
expert_data=expert_data,
1277
tp_rank=self.tp_rank)
1214
- return
1278
+ return True if return_success else None
1279
1280
if "ModelOpt" in quant_method_name:
1281
if ('weight_scale_2' in weight_name
@@ -1227,7 +1291,7 @@ class FusedMoE(torch.nn.Module):
1291
loaded_weight=loaded_weight,
1292
expert_data=expert_data,
1293
tp_rank=self.tp_rank)
1230
- return
1294
+ return True if return_success else None
1295
1296
# Case weight scales, zero_points and offset
1297
if ("scale" in weight_name or "zero" in weight_name
@@ -1264,7 +1328,7 @@ class FusedMoE(torch.nn.Module):
1328
else:
1329
raise ValueError(
1330
f"quant method must be one of {WEIGHT_SCALE_SUPPORTED}")
1267
- return
1331
+ return True if return_success else None
1332
1333
# Case weight_shape
1334
if "weight_shape" in weight_name:
@@ -1272,7 +1336,7 @@ class FusedMoE(torch.nn.Module):
1336
self._load_single_value(param=param,
1337
loaded_weight=loaded_weight,
1338
expert_id=expert_id)
1275
- return
1339
+ return True if return_success else None
1340
1341
# Case model weights
1342
if "weight" in weight_name:
@@ -1282,23 +1346,77 @@ class FusedMoE(torch.nn.Module):
1346
loaded_weight=loaded_weight,
1347
expert_data=expert_data,
1348
tp_rank=self.tp_rank)
1285
- return
1349
+ return True if return_success else None
1350
+
1351
+ return False if return_success else None
1352
+
1353
+ def get_expert_weights(self) -> Iterable[torch.Tensor]:
1354
+ weights = list(self.named_parameters())
1355
+ assert all(weight.is_contiguous() for _, weight in weights)
1356
+
1357
+ # Filter out the non-expert weights.
1358
+ # `e_score_correction_bias` is a bias for each logical expert,
1359
+ # with shape (num_logical_experts,), not an expert weight.
1360
+ NON_EXPERT_WEIGHTS = {
1361
+ "e_score_correction_bias",
1362
+ }
1363
+
1364
+ return [
1365
+ weight.view(self.local_num_experts, -1) for name, weight in weights
1366
+ if name not in NON_EXPERT_WEIGHTS
1367
+ ]
1368
+
1369
+ def set_eplb_state(
1370
+ self,
1371
+ moe_layer_idx: int,
1372
+ expert_load_view: torch.Tensor,
1373
+ logical_to_physical_map: torch.Tensor,
1374
+ logical_replica_count: torch.Tensor,
1375
+ ) -> None:
1376
+ """
1377
+ Register the EPLB state in this layer.
1378
+
1379
+ This is used later in forward pass, where we get the expert mapping
1380
+ and record the load metrics in `expert_load_view`.
1381
+ """
1382
+ self.expert_load_view = expert_load_view[moe_layer_idx]
1383
+ self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
1384
+ self.logical_replica_count = logical_replica_count[moe_layer_idx]
1385
1386
@staticmethod
1288
- def select_experts(hidden_states: torch.Tensor,
1289
- router_logits: torch.Tensor,
1290
- top_k: int,
1291
- use_grouped_topk: bool,
1292
- renormalize: bool,
1293
- topk_group: Optional[int] = None,
1294
- num_expert_group: Optional[int] = None,
1295
- custom_routing_function: Optional[Callable] = None,
1296
- scoring_func: str = "softmax",
1297
- e_score_correction_bias: Optional[torch.Tensor] = None,
1298
- indices_type: Optional[torch.dtype] = None):
1387
+ def select_experts(
1388
+ hidden_states: torch.Tensor,
1389
+ router_logits: torch.Tensor,
1390
+ top_k: int,
1391
+ use_grouped_topk: bool,
1392
+ renormalize: bool,
1393
+ topk_group: Optional[int] = None,
1394
+ num_expert_group: Optional[int] = None,
1395
+ custom_routing_function: Optional[Callable] = None,
1396
+ scoring_func: str = "softmax",
1397
+ e_score_correction_bias: Optional[torch.Tensor] = None,
1398
+ indices_type: Optional[torch.dtype] = None,
1399
+ enable_eplb: bool = False,
1400
+ expert_map: Optional[torch.Tensor] = None,
1401
+ expert_load_view: Optional[torch.Tensor] = None,
1402
+ logical_to_physical_map: Optional[torch.Tensor] = None,
1403
+ logical_replica_count: Optional[torch.Tensor] = None,
1404
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1405
+ """
1406
+ Route the input hidden states to the top-k experts based on the
1407
+ router logits.
1408
+
1409
+ Returns:
1410
+ (topk_weights, topk_ids) (tuple[torch.Tensor, torch.Tensor]):
1411
+ The weights and *global physical* expert ids of the top-k experts.
1412
+
1413
+ **Compatibility**: When EPLB is not enabled, the returned ids are
1414
+ equivalent to global logical ids, so should be compatible with
1415
+ plain MoE implementations without redundant experts.
1416
+ """
1417
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
1418
1301
- # DeekSeekv2 uses grouped_top_k
1419
+ # DeepSeekv2 uses grouped_top_k
1420
if use_grouped_topk:
1421
assert topk_group is not None
1422
assert num_expert_group is not None
@@ -1330,6 +1448,74 @@ class FusedMoE(torch.nn.Module):
1448
if indices_type is not None:
1449
topk_ids = topk_ids.to(dtype=indices_type)
1450
1451
+ if enable_eplb:
1452
+ assert expert_load_view is not None
1453
+ assert logical_to_physical_map is not None
1454
+ assert logical_replica_count is not None
1455
+
1456
+ # 1. Convert the logical expert ids to physical expert ids
1457
+ # Directly select a random replica for each logical expert
1458
+
1459
+ # TODO: maybe optimize this by using specified kernels,
1460
+ # or compute pseudo-random indices by modulo
1461
+
1462
+ # In case `indices_type` is not `torch.long` or `torch.int`,
1463
+ # e.g. `torch.uint32` as required by dispatch/combine kernels
1464
+ topk_ids_long = topk_ids.long()
1465
+ replica_indices = (
1466
+ torch.rand_like(topk_ids, dtype=torch.float) *
1467
+ logical_replica_count[topk_ids_long]).long().unsqueeze(-1)
1468
+ physical_ids = logical_to_physical_map[topk_ids_long].gather(
1469
+ -1, replica_indices).squeeze(-1)
1470
+
1471
+ topk_ids = physical_ids
1472
+
1473
+ # 2. Record expert load metrics.
1474
+
1475
+ # TODO(bowen): When using `FusedMoEModularKernel`, this
1476
+ # can be done in a more unified way, since
1477
+ # `FusedMoEPrepareAndFinalize` will return the expert
1478
+ # token count, in some cases directly from the kernel.
1479
+ # However, now there are many code paths not using
1480
+ # the modular kernel, e.g. calling `fused_experts`,
1481
+ # so we decide to keep the logic here.
1482
+ #
1483
+ # If later refactor moved all the MoE kernel calls
1484
+ # to the modular kernel, we can move this logic there
1485
+ # to achieve better efficiency.
1486
+
1487
+ # `expert_load_view`: (num_logical_experts,)
1488
+
1489
+ # Mask out non-local experts
1490
+ if expert_map is not None:
1491
+ topk_ids_local = expert_map[topk_ids]
1492
+ topk_ids_flatten = topk_ids_local.flatten()
1493
+ else:
1494
+ topk_ids_flatten = topk_ids.flatten()
1495
+
1496
+ # Should be equivalent to:
1497
+ # ```
1498
+ # topk_ids_masked = topk_ids_local[topk_ids_local >= 0]
1499
+ # expert_load_view += topk_ids_masked.bincount(
1500
+ # minlength=expert_load_view.shape[0])
1501
+ # ```
1502
+ # We use `scatter_add_` since `bincount` cannot be compiled
1503
+
1504
+ # Performance optimization:
1505
+ # `masked_fill` is significantly faster than `masked_select`
1506
+ invalid_mask = topk_ids_flatten < 0
1507
+ # Replace invalid expert ids with 0 (just a dummy position)
1508
+ # to avoid out-of-bounds errors in scatter_add_
1509
+ index = topk_ids_flatten.masked_fill_(invalid_mask, 0)
1510
+ # `src` is the valid mask, which is 1 for valid and 0 for invalid
1511
+ src = ~invalid_mask
1512
+
1513
+ expert_load_view.scatter_add_(dim=0,
1514
+ index=index.long(),
1515
+ src=src.to(expert_load_view))
1516
+
1517
+ topk_ids = topk_ids.to(dtype=indices_type)
1518
+
1519
return topk_weights, topk_ids
1520
1521
def must_reduce_shared_expert_outputs(self) -> bool:
@@ -1410,6 +1596,10 @@ class FusedMoE(torch.nn.Module):
1596
scoring_func=self.scoring_func,
1597
e_score_correction_bias=self.e_score_correction_bias,
1598
activation=self.activation,
1599
+ enable_eplb=self.enable_eplb,
1600
+ expert_load_view=self.expert_load_view,
1601
+ logical_to_physical_map=self.logical_to_physical_map,
1602
+ logical_replica_count=self.logical_replica_count,
1603
)
1604
1605
if not skip_result_store:
@@ -1467,6 +1657,10 @@ class FusedMoE(torch.nn.Module):
1657
e_score_correction_bias=self.e_score_correction_bias,
1658
activation=self.activation,
1659
apply_router_weight_on_input=self.apply_router_weight_on_input,
1660
+ enable_eplb=self.enable_eplb,
1661
+ expert_load_view=self.expert_load_view,
1662
+ logical_to_physical_map=self.logical_to_physical_map,
1663
+ logical_replica_count=self.logical_replica_count,
1664
)
1665
1666
if do_naive_dispatch_combine:
@@ -1481,16 +1675,30 @@ class FusedMoE(torch.nn.Module):
1675
1676
@classmethod
1677
def make_expert_params_mapping(
1484
- cls, ckpt_gate_proj_name: str, ckpt_down_proj_name: str,
1678
+ cls,
1679
+ ckpt_gate_proj_name: str,
1680
+ ckpt_down_proj_name: str,
1681
ckpt_up_proj_name: str,
1486
- num_experts: int) -> list[tuple[str, str, int, str]]:
1682
+ num_experts: int,
1683
+ num_redundant_experts: int = 0) -> list[tuple[str, str, int, str]]:
1684
+
1685
+ num_physical_experts = num_experts + num_redundant_experts
1686
+
1687
+ # In the returned mapping:
1688
+ # - `expert_id` is the physical expert id
1689
+ # - `weight_name` contains the weight name of the logical expert
1690
+ # So that we should map the expert id to logical in `weight_name`
1691
+ physical_to_logical_map = \
1692
+ EplbState.build_initial_global_physical_to_logical_map(
1693
+ num_experts, num_redundant_experts)
1694
1695
return [
1696
# (param_name, weight_name, expert_id, shard_id)
1697
("experts.w13_" if weight_name
1698
in [ckpt_gate_proj_name, ckpt_up_proj_name] else "experts.w2_",
1492
- f"experts.{expert_id}.{weight_name}.", expert_id, shard_id)
1493
- for expert_id in range(num_experts) for shard_id, weight_name in [
1699
+ f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.",
1700
+ expert_id, shard_id) for expert_id in range(num_physical_experts)
1701
+ for shard_id, weight_name in [
1702
("w1", ckpt_gate_proj_name),
1703
("w2", ckpt_down_proj_name),
1704
("w3", ckpt_up_proj_name),
vllm/model_executor/layers/quantization/awq_marlin.py
+8
@@ -482,7 +482,15 @@ class AWQMoEMethod(FusedMoEMethodBase):
482
e_score_correction_bias: Optional[torch.Tensor] = None,
483
apply_router_weight_on_input: bool = False,
484
activation: str = "silu",
485
+ enable_eplb: bool = False,
486
+ expert_load_view: Optional[torch.Tensor] = None,
487
+ logical_to_physical_map: Optional[torch.Tensor] = None,
488
+ logical_replica_count: Optional[torch.Tensor] = None,
489
) -> torch.Tensor:
490
+ if enable_eplb:
491
+ raise NotImplementedError(
492
+ "EPLB not supported for `AWQMoEMethod` yet.")
493
+
494
assert activation == "silu", "Only SiLU activation is supported."
495
496
if apply_router_weight_on_input:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+42
@@ -331,7 +331,15 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
331
e_score_correction_bias: Optional[torch.Tensor] = None,
332
apply_router_weight_on_input: bool = False,
333
activation: str = "silu",
334
+ enable_eplb: bool = False,
335
+ expert_load_view: Optional[torch.Tensor] = None,
336
+ logical_to_physical_map: Optional[torch.Tensor] = None,
337
+ logical_replica_count: Optional[torch.Tensor] = None,
338
) -> torch.Tensor:
339
+ if enable_eplb:
340
+ raise NotImplementedError(
341
+ "EPLB not supported for "
342
+ "`CompressedTensorsW8A8Fp8MoEMethod` yet.")
343
344
topk_weights, topk_ids = FusedMoE.select_experts(
345
hidden_states=x,
@@ -593,7 +601,15 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
601
e_score_correction_bias: Optional[torch.Tensor] = None,
602
apply_router_weight_on_input: bool = False,
603
activation: str = "silu",
604
+ enable_eplb: bool = False,
605
+ expert_load_view: Optional[torch.Tensor] = None,
606
+ logical_to_physical_map: Optional[torch.Tensor] = None,
607
+ logical_replica_count: Optional[torch.Tensor] = None,
608
) -> torch.Tensor:
609
+ if enable_eplb:
610
+ raise NotImplementedError(
611
+ "EPLB not supported for "
612
+ "`CompressedTensorsW8A8Fp8MoECutlassMethod` yet.")
613
614
topk_weights, topk_ids = FusedMoE.select_experts(
615
hidden_states=x,
@@ -722,7 +738,16 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
738
e_score_correction_bias: Optional[torch.Tensor] = None,
739
apply_router_weight_on_input: bool = False,
740
activation: str = "silu",
741
+ enable_eplb: bool = False,
742
+ expert_load_view: Optional[torch.Tensor] = None,
743
+ logical_to_physical_map: Optional[torch.Tensor] = None,
744
+ logical_replica_count: Optional[torch.Tensor] = None,
745
) -> torch.Tensor:
746
+ if enable_eplb:
747
+ raise NotImplementedError(
748
+ "EPLB not supported for "
749
+ "`CompressedTensorsW8A8Int8MoEMethod` yet.")
750
+
751
from vllm.model_executor.layers.fused_moe import fused_experts
752
753
topk_weights, topk_ids = FusedMoE.select_experts(
@@ -1012,7 +1037,16 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1037
e_score_correction_bias: Optional[torch.Tensor] = None,
1038
apply_router_weight_on_input: bool = False,
1039
activation: str = "silu",
1040
+ enable_eplb: bool = False,
1041
+ expert_load_view: Optional[torch.Tensor] = None,
1042
+ logical_to_physical_map: Optional[torch.Tensor] = None,
1043
+ logical_replica_count: Optional[torch.Tensor] = None,
1044
) -> torch.Tensor:
1045
+ if enable_eplb:
1046
+ raise NotImplementedError(
1047
+ "EPLB not supported for "
1048
+ "`CompressedTensorsWNA16MarlinMoEMethod` yet.")
1049
+
1050
assert activation == "silu", (
1051
f"{activation} not supported for Marlin MoE.")
1052
assert not apply_router_weight_on_input, (
@@ -1228,7 +1262,15 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1262
e_score_correction_bias: Optional[torch.Tensor] = None,
1263
apply_router_weight_on_input: bool = False,
1264
activation: str = "silu",
1265
+ enable_eplb: bool = False,
1266
+ expert_load_view: Optional[torch.Tensor] = None,
1267
+ logical_to_physical_map: Optional[torch.Tensor] = None,
1268
+ logical_replica_count: Optional[torch.Tensor] = None,
1269
) -> torch.Tensor:
1270
+ if enable_eplb:
1271
+ raise NotImplementedError("EPLB not supported for "
1272
+ "`CompressedTensorsWNA16MoEMethod` yet.")
1273
+
1274
from vllm.model_executor.layers.fused_moe import fused_experts
1275
1276
topk_weights, topk_ids = FusedMoE.select_experts(
vllm/model_executor/layers/quantization/experts_int8.py
+8
@@ -117,7 +117,15 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
117
e_score_correction_bias: Optional[torch.Tensor] = None,
118
apply_router_weight_on_input: bool = False,
119
activation: str = "silu",
120
+ enable_eplb: bool = False,
121
+ expert_load_view: Optional[torch.Tensor] = None,
122
+ logical_to_physical_map: Optional[torch.Tensor] = None,
123
+ logical_replica_count: Optional[torch.Tensor] = None,
124
) -> torch.Tensor:
125
+ if enable_eplb:
126
+ raise NotImplementedError(
127
+ "EPLB not supported for `ExpertsInt8MoEMethod` yet.")
128
+
129
from vllm.model_executor.layers.fused_moe import fused_experts
130
131
topk_weights, topk_ids = FusedMoE.select_experts(
vllm/model_executor/layers/quantization/fp8.py
+14
@@ -825,7 +825,16 @@ class Fp8MoEMethod(FusedMoEMethodBase):
825
e_score_correction_bias: Optional[torch.Tensor] = None,
826
apply_router_weight_on_input: bool = False,
827
activation: str = "silu",
828
+ enable_eplb: bool = False,
829
+ expert_load_view: Optional[torch.Tensor] = None,
830
+ logical_to_physical_map: Optional[torch.Tensor] = None,
831
+ logical_replica_count: Optional[torch.Tensor] = None,
832
) -> torch.Tensor:
833
+ if enable_eplb:
834
+ assert expert_load_view is not None
835
+ assert logical_to_physical_map is not None
836
+ assert logical_replica_count is not None
837
+ assert isinstance(layer, FusedMoE)
838
839
topk_weights, topk_ids = FusedMoE.select_experts(
840
hidden_states=x,
@@ -839,6 +848,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
848
scoring_func=scoring_func,
849
e_score_correction_bias=e_score_correction_bias,
850
indices_type=self.topk_indices_dtype,
851
+ enable_eplb=enable_eplb,
852
+ expert_map=expert_map,
853
+ expert_load_view=expert_load_view,
854
+ logical_to_physical_map=logical_to_physical_map,
855
+ logical_replica_count=logical_replica_count,
856
)
857
858
if self.rocm_aiter_moe_enabled:
vllm/model_executor/layers/quantization/gguf.py
+8
@@ -520,7 +520,15 @@ class GGUFMoEMethod(FusedMoEMethodBase):
520
e_score_correction_bias: Optional[torch.Tensor] = None,
521
apply_router_weight_on_input: bool = False,
522
activation: str = "silu",
523
+ enable_eplb: bool = False,
524
+ expert_load_view: Optional[torch.Tensor] = None,
525
+ logical_to_physical_map: Optional[torch.Tensor] = None,
526
+ logical_replica_count: Optional[torch.Tensor] = None,
527
):
528
+ if enable_eplb:
529
+ raise NotImplementedError(
530
+ "EPLB not supported for `GGUFMoEMethod` yet.")
531
+
532
assert activation == "silu", "Only SiLU activation is supported."
533
if apply_router_weight_on_input:
534
raise NotImplementedError(
vllm/model_executor/layers/quantization/gptq_marlin.py
+8
@@ -635,7 +635,15 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
635
e_score_correction_bias: Optional[torch.Tensor] = None,
636
apply_router_weight_on_input: bool = False,
637
activation: str = "silu",
638
+ enable_eplb: bool = False,
639
+ expert_load_view: Optional[torch.Tensor] = None,
640
+ logical_to_physical_map: Optional[torch.Tensor] = None,
641
+ logical_replica_count: Optional[torch.Tensor] = None,
642
) -> torch.Tensor:
643
+ if enable_eplb:
644
+ raise NotImplementedError(
645
+ "EPLB not supported for `GPTQMarlinMoEMethod` yet.")
646
+
647
assert activation == "silu", "Only SiLU activation is supported."
648
if apply_router_weight_on_input:
649
raise NotImplementedError(
vllm/model_executor/layers/quantization/modelopt.py
+8
@@ -664,7 +664,15 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
664
e_score_correction_bias: Optional[torch.Tensor] = None,
665
apply_router_weight_on_input: bool = False,
666
activation: str = "silu",
667
+ enable_eplb: bool = False,
668
+ expert_load_view: Optional[torch.Tensor] = None,
669
+ logical_to_physical_map: Optional[torch.Tensor] = None,
670
+ logical_replica_count: Optional[torch.Tensor] = None,
671
):
672
+ if enable_eplb:
673
+ raise NotImplementedError(
674
+ "EPLB not supported for `ModelOptNvFp4FusedMoE` yet.")
675
+
676
if self.use_marlin:
677
topk_weights, topk_ids = FusedMoE.select_experts(
678
hidden_states=x,
vllm/model_executor/layers/quantization/moe_wna16.py
+8
@@ -297,7 +297,15 @@ class MoeWNA16Method(FusedMoEMethodBase):
297
e_score_correction_bias: Optional[torch.Tensor] = None,
298
apply_router_weight_on_input: bool = False,
299
activation: str = "silu",
300
+ enable_eplb: bool = False,
301
+ expert_load_view: Optional[torch.Tensor] = None,
302
+ logical_to_physical_map: Optional[torch.Tensor] = None,
303
+ logical_replica_count: Optional[torch.Tensor] = None,
304
) -> torch.Tensor:
305
+ if enable_eplb:
306
+ raise NotImplementedError(
307
+ "EPLB not supported for `MoeWNA16Method` yet.")
308
+
309
from vllm.model_executor.layers.fused_moe import fused_experts
310
assert activation == "silu", "Only SiLU activation is supported."
311
topk_weights, topk_ids = FusedMoE.select_experts(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+8
@@ -205,7 +205,15 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
205
e_score_correction_bias: Optional[torch.Tensor] = None,
206
apply_router_weight_on_input: bool = False,
207
activation: str = "silu",
208
+ enable_eplb: bool = False,
209
+ expert_load_view: Optional[torch.Tensor] = None,
210
+ logical_to_physical_map: Optional[torch.Tensor] = None,
211
+ logical_replica_count: Optional[torch.Tensor] = None,
212
) -> torch.Tensor:
213
+ if enable_eplb:
214
+ raise NotImplementedError(
215
+ "EPLB not supported for `QuarkW8A8Fp8MoEMethod` yet.")
216
+
217
from vllm.model_executor.layers.fused_moe import fused_experts
218
219
topk_weights, topk_ids = FusedMoE.select_experts(
vllm/model_executor/models/deepseek_v2.py
+109
-18
@@ -23,7 +23,8 @@
23
# See the License for the specific language governing permissions and
24
# limitations under the License.
25
"""Inference-only DeepseekV2/DeepseekV3 model."""
26
-from collections.abc import Iterable
26
+import typing
27
+from collections.abc import Callable, Iterable
28
from typing import Any, Optional, Union
29
30
import torch
@@ -32,8 +33,10 @@ from transformers import PretrainedConfig
33
34
from vllm.attention import Attention
35
from vllm.compilation.decorators import support_torch_compile
35
-from vllm.config import CacheConfig, ModelConfig, VllmConfig
36
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
+from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
37
+ get_current_vllm_config)
38
+from vllm.distributed import (get_ep_group, get_pp_group,
39
+ get_tensor_model_parallel_world_size)
40
from vllm.model_executor.layers.activation import SiluAndMul
41
from vllm.model_executor.layers.fused_moe import FusedMoE
42
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -51,7 +54,7 @@ from vllm.model_executor.model_loader.weight_utils import (
54
from vllm.model_executor.sampling_metadata import SamplingMetadata
55
from vllm.sequence import IntermediateTensors
56
54
-from .interfaces import SupportsPP
57
+from .interfaces import MixtureOfExperts, SupportsPP
58
from .utils import (PPMissingLayer, is_pp_missing_parameter,
59
make_empty_intermediate_tensors_factory, make_layers,
60
maybe_prefix)
@@ -99,11 +102,17 @@ class DeepseekV2MoE(nn.Module):
102
config: PretrainedConfig,
103
quant_config: Optional[QuantizationConfig] = None,
104
prefix: str = "",
105
+ enable_eplb: bool = False,
106
):
107
super().__init__()
108
self.tp_size = get_tensor_model_parallel_world_size()
109
self.routed_scaling_factor = config.routed_scaling_factor
106
- self.n_shared_experts = config.n_shared_experts
110
+
111
+ self.ep_group = get_ep_group().device_group
112
+ self.ep_rank = self.ep_group.rank()
113
+ self.ep_size = self.ep_group.size()
114
+ self.n_routed_experts: int = config.n_routed_experts
115
+ self.n_shared_experts: int = config.n_shared_experts
116
117
if config.hidden_act != "silu":
118
raise ValueError(f"Unsupported activation: {config.hidden_act}. "
@@ -120,6 +129,22 @@ class DeepseekV2MoE(nn.Module):
129
else:
130
self.gate.e_score_correction_bias = None
131
132
+ # Load balancing settings.
133
+ vllm_config = get_current_vllm_config()
134
+ parallel_config = vllm_config.parallel_config
135
+ self.enable_eplb = enable_eplb
136
+
137
+ self.n_redundant_experts = parallel_config.num_redundant_experts
138
+ self.n_logical_experts = self.n_routed_experts
139
+ self.n_physical_experts = (self.n_logical_experts +
140
+ self.n_redundant_experts)
141
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
142
+
143
+ self.physical_expert_start = (self.ep_rank *
144
+ self.n_local_physical_experts)
145
+ self.physical_expert_end = (self.physical_expert_start +
146
+ self.n_local_physical_experts)
147
+
148
self.experts = FusedMoE(
149
num_experts=config.n_routed_experts,
150
top_k=config.num_experts_per_tok,
@@ -133,7 +158,9 @@ class DeepseekV2MoE(nn.Module):
158
topk_group=config.topk_group,
159
prefix=f"{prefix}.experts",
160
scoring_func=config.scoring_func,
136
- e_score_correction_bias=self.gate.e_score_correction_bias)
161
+ e_score_correction_bias=self.gate.e_score_correction_bias,
162
+ enable_eplb=self.enable_eplb,
163
+ num_redundant_experts=self.n_redundant_experts)
164
165
if config.n_shared_experts is not None:
166
intermediate_size = (config.moe_intermediate_size *
@@ -503,6 +530,7 @@ class DeepseekV2DecoderLayer(nn.Module):
530
model_config: ModelConfig,
531
cache_config: Optional[CacheConfig] = None,
532
quant_config: Optional[QuantizationConfig] = None,
533
+ enable_eplb: bool = False,
534
) -> None:
535
super().__init__()
536
self.hidden_size = config.hidden_size
@@ -543,6 +571,7 @@ class DeepseekV2DecoderLayer(nn.Module):
571
config=config,
572
quant_config=quant_config,
573
prefix=f"{prefix}.mlp",
574
+ enable_eplb=enable_eplb,
575
)
576
else:
577
self.mlp = DeepseekV2MLP(
@@ -615,6 +644,7 @@ class DeepseekV2Model(nn.Module):
644
model_config = vllm_config.model_config
645
cache_config = vllm_config.cache_config
646
quant_config = vllm_config.quant_config
647
+ enable_eplb = vllm_config.parallel_config.enable_eplb
648
self.config = config
649
650
self.vocab_size = config.vocab_size
@@ -636,6 +666,7 @@ class DeepseekV2Model(nn.Module):
666
model_config=model_config,
667
cache_config=cache_config,
668
quant_config=quant_config,
669
+ enable_eplb=enable_eplb,
670
),
671
prefix=f"{prefix}.layers")
672
@@ -681,7 +712,7 @@ class DeepseekV2Model(nn.Module):
712
return hidden_states
713
714
684
-class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
715
+class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
716
717
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
718
super().__init__()
@@ -700,6 +731,44 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
731
self.logits_processor = LogitsProcessor(config.vocab_size)
732
self.make_empty_intermediate_tensors = (
733
self.model.make_empty_intermediate_tensors)
734
+ self.expert_weights = []
735
+
736
+ # Set MoE hyperparameters
737
+ self.num_moe_layers = (config.num_hidden_layers -
738
+ config.first_k_dense_replace)
739
+ self.num_expert_groups = config.n_group
740
+
741
+ self.moe_layers: list[FusedMoE] = []
742
+ for layer in self.model.layers:
743
+ assert isinstance(layer, DeepseekV2DecoderLayer)
744
+ if isinstance(layer.mlp, DeepseekV2MoE):
745
+ self.moe_layers.append(layer.mlp.experts)
746
+
747
+ # Pick last one layer since the first ones may be dense layers.
748
+ example_moe = typing.cast(
749
+ DeepseekV2MoE, self.model.layers[config.num_hidden_layers - 1].mlp)
750
+ self.num_logical_experts = example_moe.n_logical_experts
751
+ self.num_physical_experts = example_moe.n_physical_experts
752
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
753
+ self.num_routed_experts = example_moe.n_routed_experts
754
+ self.num_shared_experts = example_moe.n_shared_experts
755
+ self.num_redundant_experts = example_moe.n_redundant_experts
756
+
757
+ def set_eplb_state(
758
+ self,
759
+ expert_load_view: torch.Tensor,
760
+ logical_to_physical_map: torch.Tensor,
761
+ logical_replica_count: torch.Tensor,
762
+ ) -> None:
763
+ for layer_idx, layer in enumerate(self.moe_layers):
764
+ # Register the expert weights.
765
+ self.expert_weights.append(layer.get_expert_weights())
766
+ layer.set_eplb_state(
767
+ moe_layer_idx=layer_idx,
768
+ expert_load_view=expert_load_view,
769
+ logical_to_physical_map=logical_to_physical_map,
770
+ logical_replica_count=logical_replica_count,
771
+ )
772
773
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
774
return self.model.get_input_embeddings(input_ids)
@@ -752,7 +821,8 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
821
ckpt_gate_proj_name="gate_proj",
822
ckpt_down_proj_name="down_proj",
823
ckpt_up_proj_name="up_proj",
755
- num_experts=self.config.n_routed_experts)
824
+ num_experts=self.config.n_routed_experts,
825
+ num_redundant_experts=self.num_redundant_experts)
826
827
params_dict = dict(self.named_parameters())
828
loaded_params: set[str] = set()
@@ -789,24 +859,44 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
859
weight_loader(param, loaded_weight, shard_id)
860
break
861
else:
862
+ is_expert_weight = False
863
for mapping in expert_params_mapping:
864
param_name, weight_name, expert_id, shard_id = mapping
865
if weight_name not in name:
866
continue
796
- name = name.replace(weight_name, param_name)
867
798
- if is_pp_missing_parameter(name, self):
868
+ # Anyway, this is an expert weight and should not be
869
+ # attempted to load as other weights later
870
+ is_expert_weight = True
871
+
872
+ # Do not modify `name` since the loop may continue here
873
+ # Instead, create a new variable
874
+ name_mapped = name.replace(weight_name, param_name)
875
+
876
+ if is_pp_missing_parameter(name_mapped, self):
877
continue
878
801
- param = params_dict[name]
802
- weight_loader = param.weight_loader
803
- weight_loader(param,
804
- loaded_weight,
805
- name,
806
- shard_id=shard_id,
807
- expert_id=expert_id)
808
- break
879
+ param = params_dict[name_mapped]
880
+ # We should ask the weight loader to return success or not
881
+ # here since otherwise we may skip experts with other
882
+ # available replicas.
883
+ weight_loader = typing.cast(Callable[..., bool],
884
+ param.weight_loader)
885
+ success = weight_loader(param,
886
+ loaded_weight,
887
+ name_mapped,
888
+ shard_id=shard_id,
889
+ expert_id=expert_id,
890
+ return_success=True)
891
+ if success:
892
+ break
893
else:
894
+ if is_expert_weight:
895
+ # We've checked that this is an expert weight
896
+ # However it's not mapped locally to this rank
897
+ # So we simply skip it
898
+ continue
899
+
900
# Skip loading extra bias for GPTQ models.
901
if name.endswith(".bias") and name not in params_dict:
902
continue
@@ -824,6 +914,7 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
914
default_weight_loader)
915
weight_loader(param, loaded_weight)
916
loaded_params.add(name)
917
+
918
return loaded_params
919
920
vllm/model_executor/models/interfaces.py
+68
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Iterable, MutableSequence
5
from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,
6
Union, overload, runtime_checkable)
7
@@ -426,6 +427,73 @@ def is_hybrid(
427
return isinstance(model, IsHybrid)
428
429
430
+@runtime_checkable
431
+class MixtureOfExperts(Protocol):
432
+ """
433
+ Check if the model is a mixture of experts (MoE) model.
434
+ """
435
+
436
+ expert_weights: MutableSequence[Iterable[Tensor]]
437
+ """
438
+ Expert weights saved in this rank.
439
+
440
+ The first dimension is the layer, and the second dimension is different
441
+ parameters in the layer, e.g. up/down projection weights.
442
+ """
443
+
444
+ num_moe_layers: int
445
+ """Number of MoE layers in this model."""
446
+
447
+ num_expert_groups: int
448
+ """Number of expert groups in this model."""
449
+
450
+ num_logical_experts: int
451
+ """Number of logical experts in this model."""
452
+
453
+ num_physical_experts: int
454
+ """Number of physical experts in this model."""
455
+
456
+ num_local_physical_experts: int
457
+ """Number of local physical experts in this model."""
458
+
459
+ num_routed_experts: int
460
+ """Number of routed experts in this model."""
461
+
462
+ num_shared_experts: int
463
+ """Number of shared experts in this model."""
464
+
465
+ num_redundant_experts: int
466
+ """Number of redundant experts in this model."""
467
+
468
+ def set_eplb_state(
469
+ self,
470
+ expert_load_view: Tensor,
471
+ logical_to_physical_map: Tensor,
472
+ logical_replica_count: Tensor,
473
+ ) -> None:
474
+ """
475
+ Register the EPLB state in the MoE model.
476
+
477
+ Since these are views of the actual EPLB state, any changes made by
478
+ the EPLB algorithm are automatically reflected in the model's behavior
479
+ without requiring additional method calls to set new states.
480
+
481
+ You should also collect model's `expert_weights` here instead of in
482
+ the weight loader, since after initial weight loading, further
483
+ processing like quantization may be applied to the weights.
484
+
485
+ Args:
486
+ expert_load_view: A view of the expert load metrics tensor.
487
+ logical_to_physical_map: Mapping from logical to physical experts.
488
+ logical_replica_count: Count of replicas for each logical expert.
489
+ """
490
+ ...
491
+
492
+
493
+def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
494
+ return isinstance(model, MixtureOfExperts)
495
+
496
+
497
@runtime_checkable
498
class HasNoOps(Protocol):
499
has_noops: ClassVar[Literal[True]] = True
vllm/v1/worker/gpu_model_runner.py
+61
-4
@@ -21,6 +21,7 @@ from vllm.attention.layer import Attention
21
from vllm.compilation.counter import compilation_counter
22
from vllm.config import (CompilationLevel, VllmConfig,
23
get_layers_from_vllm_config)
24
+from vllm.distributed.eplb.eplb_state import EplbState
25
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
26
has_kv_transfer_group)
27
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
@@ -33,7 +34,8 @@ from vllm.logger import init_logger
34
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
35
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
36
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
36
-from vllm.model_executor.models.interfaces import has_step_pooler
37
+from vllm.model_executor.models.interfaces import (has_step_pooler,
38
+ is_mixture_of_experts)
39
from vllm.multimodal import MULTIMODAL_REGISTRY
40
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
41
from vllm.multimodal.utils import group_mm_inputs_by_modality
@@ -150,6 +152,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
152
# Sampler
153
self.sampler = Sampler()
154
155
+ self.eplb_state: Optional[EplbState] = None
156
+ """
157
+ State of the expert parallelism load balancer.
158
+
159
+ Will be lazily initialized when the model is loaded.
160
+ """
161
+
162
# Lazy initializations
163
# self.model: nn.Module # Set after load_model
164
# Initialize in initialize_kv_cache
@@ -1178,6 +1187,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1187
for k, v in self.intermediate_tensors.items()
1188
})
1189
1190
+ def eplb_step(self,
1191
+ is_dummy: bool = False,
1192
+ is_profile: bool = False) -> None:
1193
+ """
1194
+ Step for the EPLB (Expert Parallelism Load Balancing) state.
1195
+ """
1196
+ if not self.parallel_config.enable_eplb:
1197
+ return
1198
+
1199
+ assert self.eplb_state is not None
1200
+ assert is_mixture_of_experts(self.model)
1201
+ self.eplb_state.step(
1202
+ self.model,
1203
+ is_dummy,
1204
+ is_profile,
1205
+ log_stats=self.parallel_config.eplb_log_balancedness,
1206
+ )
1207
+
1208
def get_dp_padding(self,
1209
num_tokens: int) -> tuple[int, Optional[torch.Tensor]]:
1210
dp_size = self.vllm_config.parallel_config.data_parallel_size
@@ -1595,6 +1622,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1622
if has_kv_transfer_group():
1623
get_kv_transfer_group().clear_connector_metadata()
1624
1625
+ self.eplb_step()
1626
+
1627
return ModelRunnerOutput(
1628
req_ids=self.input_batch.req_ids,
1629
req_id_to_index=self.input_batch.req_id_to_index,
@@ -1729,6 +1758,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1758
time_after_load - time_before_load)
1759
prepare_communication_buffer_for_model(self.model)
1760
1761
+ if is_mixture_of_experts(
1762
+ self.model) and self.parallel_config.enable_eplb:
1763
+ logger.info("EPLB is enabled for model %s.",
1764
+ self.model_config.model)
1765
+ self.eplb_state = EplbState.build(
1766
+ self.model,
1767
+ self.device,
1768
+ self.parallel_config,
1769
+ )
1770
+
1771
def save_tensorized_model(
1772
self,
1773
tensorizer_config: "TensorizerConfig",
@@ -1887,6 +1926,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1926
self,
1927
num_tokens: int,
1928
capture_attn_cudagraph: bool = False,
1929
+ skip_eplb: bool = False,
1930
+ is_profile: bool = False,
1931
) -> tuple[torch.Tensor, torch.Tensor]:
1932
1933
# Padding for DP
@@ -1983,6 +2024,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2024
assert isinstance(self.drafter, EagleProposer)
2025
self.drafter.dummy_run(num_tokens)
2026
2027
+ # This is necessary to avoid blocking DP.
2028
+ # For dummy runs, we typically skip EPLB since we don't have any real
2029
+ # requests to process.
2030
+ # However, in DP settings, there may be cases when some DP ranks do
2031
+ # not have any requests to process, so they're executing dummy batches.
2032
+ # In such cases, we still have to trigger EPLB to make sure
2033
+ # ranks execute the rearrangement in synchronization.
2034
+ if not skip_eplb:
2035
+ self.eplb_step(is_dummy=True, is_profile=is_profile)
2036
+
2037
logit_indices = np.cumsum(num_scheduled_tokens) - 1
2038
return hidden_states, hidden_states[logit_indices]
2039
@@ -2175,8 +2226,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2226
# Cache the dummy encoder outputs.
2227
self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
2228
2229
+ # Add `is_profile` here to pre-allocate communication buffers
2230
hidden_states, last_hidden_states \
2179
- = self._dummy_run(self.max_num_tokens)
2231
+ = self._dummy_run(self.max_num_tokens, is_profile=True)
2232
if get_pp_group().is_last_rank:
2233
if self.is_pooling_model:
2234
output = self._dummy_pooler_run(hidden_states)
@@ -2210,10 +2262,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2262
for num_tokens in tqdm(reversed(self.cudagraph_batch_sizes),
2263
desc="Capturing CUDA graphs",
2264
total=len(self.cudagraph_batch_sizes)):
2265
+ # We skip EPLB here since we don't want to record dummy metrics
2266
for _ in range(
2267
self.compilation_config.cudagraph_num_of_warmups):
2215
- self._dummy_run(num_tokens, capture_attn_cudagraph=full_cg)
2216
- self._dummy_run(num_tokens, capture_attn_cudagraph=full_cg)
2268
+ self._dummy_run(num_tokens,
2269
+ capture_attn_cudagraph=full_cg,
2270
+ skip_eplb=True)
2271
+ self._dummy_run(num_tokens,
2272
+ capture_attn_cudagraph=full_cg,
2273
+ skip_eplb=True)
2274
2275
end_time = time.perf_counter()
2276
end_free_gpu_memory = torch.cuda.mem_get_info()[0]
vllm/v1/worker/gpu_worker.py
+7
-2
@@ -259,9 +259,10 @@ class Worker(WorkerBase):
259
x for x in warmup_sizes if x not in
260
self.vllm_config.compilation_config.cudagraph_capture_sizes
261
]
262
+ # We skip EPLB here since we don't want to record dummy metrics
263
for size in sorted(warmup_sizes, reverse=True):
264
logger.info("Compile and warming up model for size %d", size)
264
- self.model_runner._dummy_run(size)
265
+ self.model_runner._dummy_run(size, skip_eplb=True)
266
if not self.model_config.enforce_eager:
267
self.model_runner.capture_model()
268
@@ -274,8 +275,12 @@ class Worker(WorkerBase):
275
max_num_reqs = min(self.scheduler_config.max_num_seqs,
276
self.scheduler_config.max_num_batched_tokens)
277
278
+ # We skip EPLB here since we don't want to record dummy metrics
279
hidden_states, last_hidden_states = \
278
- self.model_runner._dummy_run(num_tokens=max_num_reqs)
280
+ self.model_runner._dummy_run(
281
+ num_tokens=max_num_reqs,
282
+ skip_eplb=True,
283
+ )
284
if self.model_runner.is_pooling_model:
285
self.model_runner._dummy_pooler_run(hidden_states)
286
else: