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 12, 01:00 UTC
Abstract eplb algo (#26471)
Signed-off-by: Che Ruan <cr623@ic.ac.uk> Signed-off-by: mengxingkongzhouhan <117415539+mengxingkongzhouhan@users.noreply.github.com> Signed-off-by: Mercykid-bash <ruanche0218@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Che Ruan <cr623@ic.ac.uk> Co-authored-by: mengxingkongzhouhan <117415539+mengxingkongzhouhan@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Mercykid-bash committed
Dec 5, 2025 at 03:09 UTC
1119f6e47abd8a56d2279520ccf1721ac5176f66
8 files changed
+364
-285
tests/distributed/test_eplb_algo.py
+18
-14
@@ -4,7 +4,7 @@
4
import pytest
5
import torch
6
7
-from vllm.distributed.eplb.rebalance_algo import rebalance_experts
7
+from vllm.distributed.eplb.policy.default import DefaultEplbPolicy
8
9
10
def test_basic_rebalance():
@@ -23,7 +23,7 @@ def test_basic_rebalance():
23
num_nodes = 2
24
num_gpus = 8
25
26
- phy2log, log2phy, logcnt = rebalance_experts(
26
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
27
weight, num_replicas, num_groups, num_nodes, num_gpus
28
)
29
@@ -77,7 +77,7 @@ def test_single_gpu_case():
77
num_nodes = 1
78
num_gpus = 1
79
80
- phy2log, log2phy, logcnt = rebalance_experts(
80
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
81
weight, num_replicas, num_groups, num_nodes, num_gpus
82
)
83
@@ -99,7 +99,7 @@ def test_equal_weights():
99
num_nodes = 2
100
num_gpus = 4
101
102
- phy2log, log2phy, logcnt = rebalance_experts(
102
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
103
weight, num_replicas, num_groups, num_nodes, num_gpus
104
)
105
@@ -122,7 +122,7 @@ def test_extreme_weight_imbalance():
122
num_nodes = 2
123
num_gpus = 4
124
125
- phy2log, log2phy, logcnt = rebalance_experts(
125
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
126
weight, num_replicas, num_groups, num_nodes, num_gpus
127
)
128
@@ -150,7 +150,7 @@ def test_multiple_layers():
150
num_nodes = 2
151
num_gpus = 4
152
153
- phy2log, log2phy, logcnt = rebalance_experts(
153
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
154
weight, num_replicas, num_groups, num_nodes, num_gpus
155
)
156
@@ -175,14 +175,14 @@ def test_parameter_validation():
175
# Test non-divisible case - this should handle normally without throwing
176
# errors because the function will fall back to global load balancing
177
# strategy
178
- phy2log, log2phy, logcnt = rebalance_experts(weight, 8, 3, 2, 4)
178
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
179
assert phy2log.shape == (1, 8)
180
assert logcnt.shape == (1, 4)
181
182
# Test cases that will actually cause errors:
183
# num_physical_experts not divisible by num_gpus
184
with pytest.raises(AssertionError):
185
- rebalance_experts(weight, 7, 2, 2, 4) # 7 not divisible by 4
185
+ DefaultEplbPolicy.rebalance_experts(weight, 7, 2, 2, 4) # 7 not divisible by 4
186
187
188
def test_small_scale_hierarchical():
@@ -197,7 +197,7 @@ def test_small_scale_hierarchical():
197
num_nodes = 2 # 2 nodes
198
num_gpus = 4 # 4 GPUs
199
200
- phy2log, log2phy, logcnt = rebalance_experts(
200
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
201
weight, num_replicas, num_groups, num_nodes, num_gpus
202
)
203
@@ -224,7 +224,7 @@ def test_global_load_balance_fallback():
224
num_nodes = 2
225
num_gpus = 4
226
227
- phy2log, log2phy, logcnt = rebalance_experts(
227
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
228
weight, num_replicas, num_groups, num_nodes, num_gpus
229
)
230
@@ -246,7 +246,7 @@ def test_device_compatibility(device):
246
num_nodes = 1
247
num_gpus = 2
248
249
- phy2log, log2phy, logcnt = rebalance_experts(
249
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
250
weight, num_replicas, num_groups, num_nodes, num_gpus
251
)
252
@@ -263,7 +263,9 @@ def test_additional_cases():
263
weight1 = torch.tensor(
264
[[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
265
)
266
- phy2log1, log2phy1, logcnt1 = rebalance_experts(weight1, 24, 8, 4, 8)
266
+ phy2log1, log2phy1, logcnt1 = DefaultEplbPolicy.rebalance_experts(
267
+ weight1, 24, 8, 4, 8
268
+ )
269
270
assert phy2log1.shape == (1, 24)
271
assert logcnt1.shape == (1, 16)
@@ -276,7 +278,9 @@ def test_additional_cases():
278
[12, 25, 50, 100, 150, 200], # Increasing weights
279
]
280
)
279
- phy2log2, log2phy2, logcnt2 = rebalance_experts(weight2, 10, 3, 1, 2)
281
+ phy2log2, log2phy2, logcnt2 = DefaultEplbPolicy.rebalance_experts(
282
+ weight2, 10, 3, 1, 2
283
+ )
284
285
assert phy2log2.shape == (2, 10)
286
assert logcnt2.shape == (2, 6)
@@ -300,7 +304,7 @@ if __name__ == "__main__":
304
num_nodes = 2
305
num_gpus = 8
306
303
- phy2log, log2phy, logcnt = rebalance_experts(
307
+ phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
308
weight, num_replicas, num_groups, num_nodes, num_gpus
309
)
310
print(phy2log)
vllm/config/parallel.py
+4
@@ -35,6 +35,7 @@ logger = init_logger(__name__)
35
ExpertPlacementStrategy = Literal["linear", "round_robin"]
36
DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
37
DataParallelBackend = Literal["ray", "mp"]
38
+EPLBPolicyOption = Literal["default"]
39
40
41
@config
@@ -65,6 +66,9 @@ class EPLBConfig:
66
Whether to use non-blocking EPLB.
67
"""
68
69
+ policy: EPLBPolicyOption = "default"
70
+ """The policy type for expert parallel load balancing (EPLB)."""
71
+
72
73
@config
74
@dataclass
vllm/distributed/eplb/__init__.py
+1
-6
@@ -1,8 +1,3 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Expert parallelism load balancer (EPLB).
5
-"""
6
-
7
-from .eplb_state import *
8
-from .rebalance_algo import *
3
+"""Expert parallelism load balancer (EPLB)."""
vllm/distributed/eplb/eplb_state.py
+15
-5
@@ -45,7 +45,7 @@ from vllm.logger import init_logger
45
from vllm.model_executor.models.interfaces import MixtureOfExperts
46
47
from .async_worker import start_async_worker
48
-from .rebalance_algo import rebalance_experts
48
+from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy
49
from .rebalance_execute import move_from_buffer, rearrange_expert_weights_inplace
50
51
logger = init_logger(__name__)
@@ -213,18 +213,23 @@ class EplbState:
213
self.parallel_config = parallel_config
214
self.device = device
215
self.model_states: dict[str, EplbModelState] = {}
216
+ self.policy: type[AbstractEplbPolicy] = DefaultEplbPolicy
217
+ """
218
+ Selected EPLB algorithm class
219
+ """
220
+ self.expert_load_window_step: int = 0
221
"""
222
Current step in the sliding window.
223
224
Different from `expert_rearrangement_step`,
225
each EP rank may have its own `expert_load_window_step`.
226
"""
222
- self.expert_load_window_step: int = 0
227
+ self.expert_load_window_size: int = 0
228
"""
229
Size of the expert load sliding window.
230
This is a constant and is taken from the config.
231
"""
227
- self.expert_load_window_size: int = 0
232
+ self.expert_rearrangement_step: int = 0
233
"""
234
Steps after last rearrangement.
235
Will trigger a rearrangement if it exceeds the threshold.
@@ -415,6 +420,10 @@ class EplbState:
420
)
421
self.expert_rearrangement_step_interval = eplb_step_interval
422
423
+ # Set the policy based on the selected eplb algorithm type.
424
+ policy_type = self.parallel_config.eplb_config.policy
425
+ self.policy = EPLB_POLICIES[policy_type]
426
+ logger.debug("Selected EPLB policy: %d", policy_type)
427
if global_expert_load is not None:
428
ep_group = get_ep_group().device_group
429
assert global_expert_load.shape == (
@@ -441,7 +450,7 @@ class EplbState:
450
new_physical_to_logical_map,
451
new_logical_to_physical_map,
452
new_logical_replica_count,
444
- ) = rebalance_experts(
453
+ ) = self.policy.rebalance_experts(
454
global_expert_load,
455
num_replicas,
456
num_groups,
@@ -776,6 +785,7 @@ class EplbState:
785
f"{num_gpus=}, {num_nodes=}"
786
)
787
788
+ # Get new expert mappings
789
for eplb_model_state, global_expert_load_window in zip(
790
self.model_states.values(), global_expert_load_windows
791
):
@@ -784,7 +794,7 @@ class EplbState:
794
new_physical_to_logical_map,
795
new_logical_to_physical_map,
796
new_logical_replica_count,
787
- ) = rebalance_experts(
797
+ ) = self.policy.rebalance_experts(
798
global_expert_load_window,
799
num_replicas,
800
num_groups,
vllm/distributed/eplb/policy/__init__.py
new
+19
@@ -0,0 +1,19 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import get_args
4
+
5
+from vllm.config.parallel import EPLBPolicyOption
6
+
7
+from .abstract import AbstractEplbPolicy
8
+from .default import DefaultEplbPolicy
9
+
10
+EPLB_POLICIES = {"default": DefaultEplbPolicy}
11
+
12
+# Ensure that the EPLB_POLICIES keys match the EPLBPolicyOption values
13
+assert set(EPLB_POLICIES.keys()) == set(get_args(EPLBPolicyOption))
14
+
15
+__all__ = [
16
+ "AbstractEplbPolicy",
17
+ "DefaultEplbPolicy",
18
+ "EPLB_POLICIES",
19
+]
vllm/distributed/eplb/policy/abstract.py
new
+40
@@ -0,0 +1,40 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+
6
+import torch
7
+
8
+
9
+class AbstractEplbPolicy(ABC):
10
+ @classmethod
11
+ @abstractmethod
12
+ def rebalance_experts(
13
+ cls,
14
+ weight: torch.Tensor,
15
+ num_replicas: int,
16
+ num_groups: int,
17
+ num_nodes: int,
18
+ num_ranks: int,
19
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
20
+ """
21
+ Entry point for expert-parallelism load balancer.
22
+
23
+ Parameters:
24
+ weight: [layers, num_logical_experts], the load statistics
25
+ for all logical experts
26
+ num_replicas: number of physical experts, must be a multiple of
27
+ `num_ranks`
28
+ num_groups: number of expert groups
29
+ num_nodes: number of server nodes
30
+ num_ranks: number of ranks, must be a multiple of `num_nodes`
31
+
32
+ Returns:
33
+ physical_to_logical_map: [layers, num_replicas], the expert
34
+ index of each replica
35
+ logical_to_physical_map: [layers, num_logical_experts, X],
36
+ the replica indices for each expert
37
+ expert_count: [layers, num_logical_experts], number of
38
+ physical replicas for each logical expert
39
+ """
40
+ raise NotImplementedError
vllm/distributed/eplb/policy/default.py
new
+267
@@ -0,0 +1,267 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Expert parallelism load balancer (EPLB) for vLLM.
5
+
6
+This module implements the core rearrangement algorithm.
7
+
8
+The rearrangement algorithm is adapted from
9
+[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).
10
+
11
+Please find at [#12](https://github.com/deepseek-ai/EPLB/issues/12) an example
12
+on how the EPLB algorithm works.
13
+"""
14
+
15
+import numpy as np
16
+import torch
17
+
18
+from .abstract import AbstractEplbPolicy
19
+
20
+
21
+class DefaultEplbPolicy(AbstractEplbPolicy):
22
+ @classmethod
23
+ def balanced_packing(
24
+ cls, weight: torch.Tensor, num_packs: int
25
+ ) -> tuple[torch.Tensor, torch.Tensor]:
26
+ """
27
+ Pack n weighted objects to m packs, such that each bin contains exactly
28
+ n/m objects and the weights of all packs are as balanced as possible.
29
+
30
+ Parameters:
31
+ weight: [X, n], the weight of each item
32
+ num_packs: number of packs
33
+
34
+ Returns:
35
+ pack_index: [X, n], the pack index of each item
36
+ rank_in_pack: [X, n], the rank of the item in the pack
37
+ """
38
+ num_layers, num_groups = weight.shape
39
+ assert num_groups % num_packs == 0
40
+ groups_per_pack = num_groups // num_packs
41
+
42
+ device = weight.device
43
+
44
+ if groups_per_pack == 1:
45
+ pack_index = torch.arange(
46
+ weight.size(-1), dtype=torch.int64, device=device
47
+ ).expand(weight.shape)
48
+ rank_in_pack = torch.zeros_like(weight, dtype=torch.int64, device=device)
49
+ return pack_index, rank_in_pack
50
+
51
+ weight_np = weight.cpu().numpy()
52
+
53
+ # Sort and get indices in decending order
54
+ indices_np = np.argsort(-weight_np, axis=-1)
55
+
56
+ pack_index_np = np.full((num_layers, num_groups), -1, dtype=np.int64)
57
+ rank_in_pack_np = np.full((num_layers, num_groups), -1, dtype=np.int64)
58
+
59
+ # Run the packing algorithm
60
+ for i in range(num_layers):
61
+ pack_weights = [0.0] * num_packs
62
+ pack_items = [0] * num_packs
63
+
64
+ for group in indices_np[i]:
65
+ # Find a pack with capacity that has the lowest weight
66
+ pack = min(
67
+ (j for j in range(num_packs) if pack_items[j] < groups_per_pack),
68
+ key=pack_weights.__getitem__,
69
+ )
70
+
71
+ assert pack_items[pack] < groups_per_pack
72
+ pack_index_np[i, group] = pack
73
+ rank_in_pack_np[i, group] = pack_items[pack]
74
+ pack_weights[pack] += weight_np[i, group]
75
+ pack_items[pack] += 1
76
+
77
+ pack_index = torch.from_numpy(pack_index_np).to(device)
78
+ rank_in_pack = torch.from_numpy(rank_in_pack_np).to(device)
79
+
80
+ return pack_index, rank_in_pack
81
+
82
+ @classmethod
83
+ def replicate_experts(
84
+ cls, weight: torch.Tensor, num_phy: int
85
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
86
+ """
87
+ Replicate `num_log` experts to `num_phy` replicas, such that the maximum
88
+ load of all replicas is minimized.
89
+
90
+ Parameters:
91
+ weight: [X, num_log]
92
+ num_phy: total number of experts after replication
93
+
94
+ Returns:
95
+ phy2log: [X, num_phy], logical expert id of each physical expert
96
+ rank: [X, num_phy], the replica rank
97
+ logcnt: [X, num_log], number of replicas for each logical expert
98
+ """
99
+ n, num_log = weight.shape
100
+ num_redundant = num_phy - num_log
101
+ assert num_redundant >= 0
102
+ device = weight.device
103
+ phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)
104
+ rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
105
+ logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
106
+ arangen = torch.arange(n, dtype=torch.int64, device=device)
107
+ for i in range(num_log, num_phy):
108
+ redundant_indices = (weight / logcnt).max(dim=-1).indices
109
+ phy2log[:, i] = redundant_indices
110
+ rank[:, i] = logcnt[arangen, redundant_indices]
111
+ logcnt[arangen, redundant_indices] += 1
112
+ return phy2log, rank, logcnt
113
+
114
+ @classmethod
115
+ def rebalance_experts_hierarchical(
116
+ cls,
117
+ weight: torch.Tensor,
118
+ num_physical_experts: int,
119
+ num_groups: int,
120
+ num_nodes: int,
121
+ num_gpus: int,
122
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
123
+ """
124
+ Parameters:
125
+ weight: [num_moe_layers, num_logical_experts]
126
+ num_physical_experts: number of physical experts after replication
127
+ num_groups: number of expert groups
128
+ num_nodes: number of server nodes, where the intra-node network
129
+ (e.g, NVLink) is faster
130
+ num_gpus: number of GPUs, must be a multiple of `num_nodes`
131
+
132
+ Returns:
133
+ phy2log: [layers, num_replicas], the expert
134
+ index of each replica
135
+ log2phy: [layers, num_logical_experts, X],
136
+ the replica indices for each expert
137
+ logcnt: [layers, num_logical_experts], number of
138
+ physical replicas for each logical expert
139
+ """
140
+ num_layers, num_logical_experts = weight.shape
141
+ assert num_logical_experts % num_groups == 0
142
+ group_size = num_logical_experts // num_groups
143
+ assert num_groups % num_nodes == 0
144
+ groups_per_node = num_groups // num_nodes
145
+ assert num_gpus % num_nodes == 0
146
+ assert num_physical_experts % num_gpus == 0
147
+ phy_experts_per_gpu = num_physical_experts // num_gpus
148
+
149
+ def inverse(perm: torch.Tensor) -> torch.Tensor:
150
+ inv = torch.empty_like(perm)
151
+ inv.scatter_(
152
+ 1,
153
+ perm,
154
+ torch.arange(
155
+ perm.size(1), dtype=torch.int64, device=perm.device
156
+ ).expand(perm.shape),
157
+ )
158
+ return inv
159
+
160
+ # Step 1: pack groups to nodes
161
+ tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
162
+ group_pack_index, group_rank_in_pack = cls.balanced_packing(
163
+ tokens_per_group, num_nodes
164
+ )
165
+ log2mlog = (
166
+ (
167
+ (group_pack_index * groups_per_node + group_rank_in_pack) * group_size
168
+ ).unsqueeze(-1)
169
+ + torch.arange(
170
+ group_size, dtype=torch.int64, device=group_pack_index.device
171
+ )
172
+ ).flatten(-2)
173
+ mlog2log = inverse(log2mlog)
174
+
175
+ # Step 2: construct redundant experts within nodes
176
+ # [num_layers * num_nodes, num_logical_experts // num_nodes]
177
+ tokens_per_mlog = weight.gather(-1, mlog2log).view(
178
+ -1, num_logical_experts // num_nodes
179
+ )
180
+ phy2mlog, phyrank, mlogcnt = cls.replicate_experts(
181
+ tokens_per_mlog, num_physical_experts // num_nodes
182
+ )
183
+
184
+ # Step 3: pack physical_experts to GPUs
185
+ # [num_layers * num_nodes, num_physical_experts // num_nodes]
186
+ tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
187
+ pack_index, rank_in_pack = cls.balanced_packing(
188
+ tokens_per_phy, num_gpus // num_nodes
189
+ )
190
+ phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
191
+ pphy2phy = inverse(phy2pphy)
192
+
193
+ pphy2mlog = phy2mlog.gather(
194
+ -1, pphy2phy
195
+ ) # [num_layers * num_nodes, num_log_per_nodes]
196
+ pphy2mlog = (
197
+ pphy2mlog.view(num_layers, num_nodes, -1)
198
+ + torch.arange(
199
+ 0,
200
+ num_logical_experts,
201
+ num_logical_experts // num_nodes,
202
+ device=group_pack_index.device,
203
+ ).view(1, -1, 1)
204
+ ).flatten(-2)
205
+ pphy2log = mlog2log.gather(-1, pphy2mlog)
206
+ pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
207
+ logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
208
+ return pphy2log, pphyrank, logcnt
209
+
210
+ @classmethod
211
+ def rebalance_experts(
212
+ cls,
213
+ weight: torch.Tensor,
214
+ num_replicas: int,
215
+ num_groups: int,
216
+ num_nodes: int,
217
+ num_ranks: int,
218
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
219
+ """
220
+ Entry point for expert-parallelism load balancer.
221
+
222
+ Parameters:
223
+ weight: [layers, num_logical_experts], the load statistics for all
224
+ logical experts
225
+ num_replicas: number of physical experts, must be a multiple of
226
+ `num_gpus`
227
+ num_groups: number of expert groups
228
+ num_nodes: number of server nodes, where the intra-node network
229
+ (e.g, NVLink) is faster
230
+ num_ranks: number of ranks, must be a multiple of `num_nodes`
231
+
232
+ Returns:
233
+ phy2log: [layers, num_replicas], the expert
234
+ index of each replica
235
+ log2phy: [layers, num_logical_experts, X],
236
+ the replica indices for each expert
237
+ logcnt: [layers, num_logical_experts], number of
238
+ physical replicas for each logical expert
239
+ """
240
+ num_layers, num_logical_experts = weight.shape
241
+ weight = weight.float()
242
+ if num_groups % num_nodes == 0:
243
+ # use hierarchical load-balance policy
244
+ phy2log, phyrank, logcnt = cls.rebalance_experts_hierarchical(
245
+ weight, num_replicas, num_groups, num_nodes, num_ranks
246
+ )
247
+ else:
248
+ # use global load-balance policy
249
+ phy2log, phyrank, logcnt = cls.rebalance_experts_hierarchical(
250
+ weight, num_replicas, 1, 1, num_ranks
251
+ )
252
+ num_redundant_experts = num_replicas - num_logical_experts
253
+ maxlogcnt = num_redundant_experts + 1
254
+ log2phy: torch.Tensor = torch.full(
255
+ (num_layers, num_logical_experts, maxlogcnt),
256
+ -1,
257
+ dtype=torch.int64,
258
+ device=logcnt.device,
259
+ )
260
+ log2phy.view(num_layers, -1).scatter_(
261
+ -1,
262
+ phy2log * maxlogcnt + phyrank,
263
+ torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(
264
+ num_layers, -1
265
+ ),
266
+ )
267
+ return phy2log, log2phy, logcnt
vllm/distributed/eplb/rebalance_algo.py
deleted
-260
@@ -1,260 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Expert parallelism load balancer (EPLB) for vLLM.
5
-
6
-This module implements the core rearrangement algorithm.
7
-
8
-The rearrangement algorithm is adapted from
9
-[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).
10
-
11
-Please find at [#12](https://github.com/deepseek-ai/EPLB/issues/12) an example
12
-on how the EPLB algorithm works.
13
-"""
14
-
15
-import numpy as np
16
-import torch
17
-
18
-
19
-def balanced_packing(
20
- weight: torch.Tensor, num_packs: int
21
-) -> tuple[torch.Tensor, torch.Tensor]:
22
- """
23
- Pack n weighted objects to m packs, such that each bin contains exactly
24
- n/m objects and the weights of all packs are as balanced as possible.
25
-
26
- Parameters:
27
- weight: [X, n], the weight of each item
28
- num_packs: number of packs
29
-
30
- Returns:
31
- pack_index: [X, n], the pack index of each item
32
- rank_in_pack: [X, n], the rank of the item in the pack
33
- """
34
- num_layers, num_groups = weight.shape
35
- assert num_groups % num_packs == 0
36
- groups_per_pack = num_groups // num_packs
37
-
38
- device = weight.device
39
-
40
- if groups_per_pack == 1:
41
- pack_index = torch.arange(
42
- weight.size(-1), dtype=torch.int64, device=device
43
- ).expand(weight.shape)
44
- rank_in_pack = torch.zeros_like(weight, dtype=torch.int64, device=device)
45
- return pack_index, rank_in_pack
46
-
47
- weight_np = weight.cpu().numpy()
48
-
49
- # Sort and get indices in decending order
50
- indices_np = np.argsort(-weight_np, axis=-1)
51
-
52
- pack_index_np = np.full((num_layers, num_groups), -1, dtype=np.int64)
53
- rank_in_pack_np = np.full((num_layers, num_groups), -1, dtype=np.int64)
54
-
55
- # Run the packing algorithm
56
- for i in range(num_layers):
57
- pack_weights = [0.0] * num_packs
58
- pack_items = [0] * num_packs
59
-
60
- for group in indices_np[i]:
61
- # Find a pack with capacity that has the lowest weight
62
- pack = min(
63
- (j for j in range(num_packs) if pack_items[j] < groups_per_pack),
64
- key=pack_weights.__getitem__,
65
- )
66
-
67
- assert pack_items[pack] < groups_per_pack
68
- pack_index_np[i, group] = pack
69
- rank_in_pack_np[i, group] = pack_items[pack]
70
- pack_weights[pack] += weight_np[i, group]
71
- pack_items[pack] += 1
72
-
73
- pack_index = torch.from_numpy(pack_index_np).to(device)
74
- rank_in_pack = torch.from_numpy(rank_in_pack_np).to(device)
75
-
76
- return pack_index, rank_in_pack
77
-
78
-
79
-def replicate_experts(
80
- weight: torch.Tensor, num_phy: int
81
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
82
- """
83
- Replicate `num_log` experts to `num_phy` replicas, such that the maximum
84
- load of all replicas is minimized.
85
-
86
- Parameters:
87
- weight: [X, num_log]
88
- num_phy: total number of experts after replication
89
-
90
- Returns:
91
- phy2log: [X, num_phy], logical expert id of each physical expert
92
- rank: [X, num_phy], the replica rank
93
- logcnt: [X, num_log], number of replicas for each logical expert
94
- """
95
- n, num_log = weight.shape
96
- num_redundant = num_phy - num_log
97
- assert num_redundant >= 0
98
- device = weight.device
99
- phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)
100
- rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
101
- logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
102
- arangen = torch.arange(n, dtype=torch.int64, device=device)
103
- for i in range(num_log, num_phy):
104
- redundant_indices = (weight / logcnt).max(dim=-1).indices
105
- phy2log[:, i] = redundant_indices
106
- rank[:, i] = logcnt[arangen, redundant_indices]
107
- logcnt[arangen, redundant_indices] += 1
108
- return phy2log, rank, logcnt
109
-
110
-
111
-def rebalance_experts_hierarchical(
112
- weight: torch.Tensor,
113
- num_physical_experts: int,
114
- num_groups: int,
115
- num_nodes: int,
116
- num_gpus: int,
117
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
118
- """
119
- Parameters:
120
- weight: [num_moe_layers, num_logical_experts]
121
- num_physical_experts: number of physical experts after replication
122
- num_groups: number of expert groups
123
- num_nodes: number of server nodes, where the intra-node network
124
- (e.g., NVLink) is faster
125
- num_gpus: number of GPUs, must be a multiple of `num_nodes`
126
-
127
- Returns:
128
- physical_to_logical_map (torch.Tensor):
129
- [num_moe_layers, num_physical_experts]
130
- logical_to_physical_map (torch.Tensor):
131
- [num_moe_layers, num_logical_experts, X]
132
- logical_count (torch.Tensor):
133
- [num_moe_layers, num_logical_experts]
134
- """
135
- num_layers, num_logical_experts = weight.shape
136
- assert num_logical_experts % num_groups == 0
137
- group_size = num_logical_experts // num_groups
138
- assert num_groups % num_nodes == 0
139
- groups_per_node = num_groups // num_nodes
140
- assert num_gpus % num_nodes == 0
141
- assert num_physical_experts % num_gpus == 0
142
- phy_experts_per_gpu = num_physical_experts // num_gpus
143
-
144
- def inverse(perm: torch.Tensor) -> torch.Tensor:
145
- inv = torch.empty_like(perm)
146
- inv.scatter_(
147
- 1,
148
- perm,
149
- torch.arange(perm.size(1), dtype=torch.int64, device=perm.device).expand(
150
- perm.shape
151
- ),
152
- )
153
- return inv
154
-
155
- # Step 1: pack groups to nodes
156
- tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
157
- group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes)
158
- log2mlog = (
159
- (
160
- (group_pack_index * groups_per_node + group_rank_in_pack) * group_size
161
- ).unsqueeze(-1)
162
- + torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device)
163
- ).flatten(-2)
164
- mlog2log = inverse(log2mlog)
165
-
166
- # Step 2: construct redundant experts within nodes
167
- # [num_layers * num_nodes, num_logical_experts // num_nodes]
168
- tokens_per_mlog = weight.gather(-1, mlog2log).view(
169
- -1, num_logical_experts // num_nodes
170
- )
171
- phy2mlog, phyrank, mlogcnt = replicate_experts(
172
- tokens_per_mlog, num_physical_experts // num_nodes
173
- )
174
-
175
- # Step 3: pack physical_experts to GPUs
176
- # [num_layers * num_nodes, num_physical_experts // num_nodes]
177
- tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
178
- pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes)
179
- phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
180
- pphy2phy = inverse(phy2pphy)
181
-
182
- pphy2mlog = phy2mlog.gather(
183
- -1, pphy2phy
184
- ) # [num_layers * num_nodes, num_log_per_nodes]
185
- pphy2mlog = (
186
- pphy2mlog.view(num_layers, num_nodes, -1)
187
- + torch.arange(
188
- 0,
189
- num_logical_experts,
190
- num_logical_experts // num_nodes,
191
- device=group_pack_index.device,
192
- ).view(1, -1, 1)
193
- ).flatten(-2)
194
- pphy2log = mlog2log.gather(-1, pphy2mlog)
195
- pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
196
- logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
197
- return pphy2log, pphyrank, logcnt
198
-
199
-
200
-def rebalance_experts(
201
- weight: torch.Tensor,
202
- num_replicas: int,
203
- num_groups: int,
204
- num_nodes: int,
205
- num_gpus: int,
206
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
207
- """
208
- Entry point for expert-parallelism load balancer.
209
-
210
- Parameters:
211
- weight: [layers, num_logical_experts], the load statistics for all
212
- logical experts
213
- num_replicas: number of physical experts, must be a multiple of
214
- `num_gpus`
215
- num_groups: number of expert groups
216
- num_nodes: number of server nodes, where the intra-node network
217
- (e.g, NVLink) is faster
218
- num_gpus: number of GPUs, must be a multiple of `num_nodes`
219
-
220
- Returns:
221
- physical_to_logical_map:
222
- [layers, num_replicas], the expert index of each replica
223
- logical_to_physical_map:
224
- [layers, num_logical_experts, X], the replica indices for each
225
- expert
226
- expert_count:
227
- [layers, num_logical_experts], number of physical
228
- replicas for each logical expert
229
- """
230
- num_layers, num_logical_experts = weight.shape
231
- weight = weight.float()
232
- if num_groups % num_nodes == 0:
233
- # use hierarchical load-balance policy
234
- phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
235
- weight, num_replicas, num_groups, num_nodes, num_gpus
236
- )
237
- else:
238
- # use global load-balance policy
239
- phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
240
- weight, num_replicas, 1, 1, num_gpus
241
- )
242
- num_redundant_experts = num_replicas - num_logical_experts
243
- maxlogcnt = num_redundant_experts + 1
244
- log2phy: torch.Tensor = torch.full(
245
- (num_layers, num_logical_experts, maxlogcnt),
246
- -1,
247
- dtype=torch.int64,
248
- device=logcnt.device,
249
- )
250
- log2phy.view(num_layers, -1).scatter_(
251
- -1,
252
- phy2log * maxlogcnt + phyrank,
253
- torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(
254
- num_layers, -1
255
- ),
256
- )
257
- return phy2log, log2phy, logcnt
258
-
259
-
260
-__all__ = ["rebalance_experts"]