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 6, 10:45 UTC
[EPLB] Simplify EPLB rearrange by only returning one map (#36267)
Signed-off-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Sage Moore committed
Mar 18, 2026 at 17:34 UTC
c32a58cc2aa36e9fecd1410644d5b38a92f509bb
5 files changed
+196
-159
tests/distributed/test_eplb_algo.py
+63
-18
@@ -5,6 +5,7 @@ import numpy as np
5
import pytest
6
import torch
7
8
+from vllm.distributed.eplb.eplb_state import compute_logical_maps
9
from vllm.distributed.eplb.policy.default import DefaultEplbPolicy
10
11
@@ -24,9 +25,10 @@ def test_basic_rebalance():
25
num_nodes = 2
26
num_gpus = 8
27
27
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
28
+ phy2log = DefaultEplbPolicy.rebalance_experts(
29
weight, num_replicas, num_groups, num_nodes, num_gpus
30
)
31
+ log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
32
33
# Verify output shapes
34
assert phy2log.shape == (
@@ -78,9 +80,10 @@ def test_single_gpu_case():
80
num_nodes = 1
81
num_gpus = 1
82
81
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
83
+ phy2log = DefaultEplbPolicy.rebalance_experts(
84
weight, num_replicas, num_groups, num_nodes, num_gpus
85
)
86
+ log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
87
88
# Verify shapes
89
assert phy2log.shape == (1, 4)
@@ -100,9 +103,10 @@ def test_equal_weights():
103
num_nodes = 2
104
num_gpus = 4
105
103
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
106
+ phy2log = DefaultEplbPolicy.rebalance_experts(
107
weight, num_replicas, num_groups, num_nodes, num_gpus
108
)
109
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
110
111
# Verify shapes
112
assert phy2log.shape == (1, 8)
@@ -123,9 +127,10 @@ def test_extreme_weight_imbalance():
127
num_nodes = 2
128
num_gpus = 4
129
126
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
130
+ phy2log = DefaultEplbPolicy.rebalance_experts(
131
weight, num_replicas, num_groups, num_nodes, num_gpus
132
)
133
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
134
135
# Verify shapes
136
assert phy2log.shape == (1, 12)
@@ -151,9 +156,10 @@ def test_multiple_layers():
156
num_nodes = 2
157
num_gpus = 4
158
154
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
159
+ phy2log = DefaultEplbPolicy.rebalance_experts(
160
weight, num_replicas, num_groups, num_nodes, num_gpus
161
)
162
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
163
164
# Verify shapes
165
assert phy2log.shape == (3, 8)
@@ -176,7 +182,8 @@ def test_parameter_validation():
182
# Test non-divisible case - this should handle normally without throwing
183
# errors because the function will fall back to global load balancing
184
# strategy
179
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
185
+ phy2log = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
186
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
187
assert phy2log.shape == (1, 8)
188
assert logcnt.shape == (1, 4)
189
@@ -198,9 +205,10 @@ def test_small_scale_hierarchical():
205
num_nodes = 2 # 2 nodes
206
num_gpus = 4 # 4 GPUs
207
201
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
208
+ phy2log = DefaultEplbPolicy.rebalance_experts(
209
weight, num_replicas, num_groups, num_nodes, num_gpus
210
)
211
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
212
213
# Verify basic constraints
214
assert phy2log.shape == (1, 12)
@@ -225,9 +233,10 @@ def test_global_load_balance_fallback():
233
num_nodes = 2
234
num_gpus = 4
235
228
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
236
+ phy2log = DefaultEplbPolicy.rebalance_experts(
237
weight, num_replicas, num_groups, num_nodes, num_gpus
238
)
239
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
240
241
# Should work normally, just using global load balancing strategy
242
assert phy2log.shape == (1, 8)
@@ -247,9 +256,10 @@ def test_device_compatibility(device):
256
num_nodes = 1
257
num_gpus = 2
258
250
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
259
+ phy2log = DefaultEplbPolicy.rebalance_experts(
260
weight, num_replicas, num_groups, num_nodes, num_gpus
261
)
262
+ _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
263
264
# Function will convert to CPU internally, but should handle different
265
# device inputs normally
@@ -264,9 +274,8 @@ def test_additional_cases():
274
weight1 = torch.tensor(
275
[[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
276
)
267
- phy2log1, log2phy1, logcnt1 = DefaultEplbPolicy.rebalance_experts(
268
- weight1, 24, 8, 4, 8
269
- )
277
+ phy2log1 = DefaultEplbPolicy.rebalance_experts(weight1, 24, 8, 4, 8)
278
+ _, logcnt1 = compute_logical_maps(phy2log1, weight1.shape[-1])
279
280
assert phy2log1.shape == (1, 24)
281
assert logcnt1.shape == (1, 16)
@@ -279,9 +288,8 @@ def test_additional_cases():
288
[12, 25, 50, 100, 150, 200], # Increasing weights
289
]
290
)
282
- phy2log2, log2phy2, logcnt2 = DefaultEplbPolicy.rebalance_experts(
283
- weight2, 10, 3, 1, 2
284
- )
291
+ phy2log2 = DefaultEplbPolicy.rebalance_experts(weight2, 10, 3, 1, 2)
292
+ _, logcnt2 = compute_logical_maps(phy2log2, weight2.shape[-1])
293
294
assert phy2log2.shape == (2, 10)
295
assert logcnt2.shape == (2, 6)
@@ -292,6 +300,42 @@ def test_additional_cases():
300
assert logcnt2[layer, max_weight_idx] >= 2
301
302
303
+def test_compute_logical_maps_with_negative_indices():
304
+ """
305
+ Test that compute_logical_maps correctly handles physical slots containing
306
+ -1 (unused slots).
307
+ """
308
+ # 2 layers, 6 physical slots, 4 logical experts.
309
+ # Slots 2 and 5 are unused (-1).
310
+ phy2log = torch.tensor(
311
+ [
312
+ [0, 1, -1, 2, 3, -1],
313
+ [3, -1, 2, 1, 0, -1],
314
+ ]
315
+ )
316
+ num_layers = 2
317
+ num_logical_experts = 4
318
+
319
+ log2phy, logcnt = compute_logical_maps(phy2log, num_logical_experts)
320
+
321
+ assert logcnt.shape == (num_layers, num_logical_experts)
322
+ assert log2phy.shape == (num_layers, num_logical_experts, 1)
323
+
324
+ expected_logcnt = torch.ones(num_layers, num_logical_experts, dtype=phy2log.dtype)
325
+ assert torch.all(logcnt == expected_logcnt), (
326
+ f"Expected that all replica counts == 1, got {logcnt}"
327
+ )
328
+
329
+ assert torch.all(log2phy >= 0), (
330
+ "log2phy should only contain valid physical indices, not -1"
331
+ )
332
+
333
+ assert log2phy[0, 0, 0] == 0
334
+ assert log2phy[0, 1, 0] == 1
335
+ assert log2phy[0, 2, 0] == 3
336
+ assert log2phy[0, 3, 0] == 4
337
+
338
+
339
if __name__ == "__main__":
340
weight = torch.tensor(
341
[
@@ -305,7 +349,7 @@ if __name__ == "__main__":
349
num_nodes = 2
350
num_gpus = 8
351
308
- phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
352
+ phy2log = DefaultEplbPolicy.rebalance_experts(
353
weight, num_replicas, num_groups, num_nodes, num_gpus
354
)
355
print(phy2log)
@@ -434,9 +478,10 @@ def test_preserve_intragpu_slots(
478
"""Experts that stay on a GPU keep their old slots; incoming not lost."""
479
phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(new_phy2log)
480
437
- post_phy2log, post_phy_replicas_idx = DefaultEplbPolicy.preserve_intragpu_slots(
438
- new_phy2log, phy_replicas_idx, num_ranks, old_phy2log
481
+ post_phy2log = DefaultEplbPolicy.preserve_intragpu_slots(
482
+ new_phy2log, num_ranks, old_phy2log
483
)
484
+ post_phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(post_phy2log)
485
486
# Shapes preserved
487
assert post_phy2log.shape == new_phy2log.shape
vllm/distributed/eplb/async_worker.py
+1
-15
@@ -73,11 +73,7 @@ def run_rebalance_experts(
73
# Move the global expert load window to CPU for computation.
74
global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
75
# Compute new expert mappings for the model
76
- (
77
- new_physical_to_logical_map,
78
- new_logical_to_physical_map,
79
- new_logical_replica_count,
80
- ) = eplb_state.policy.rebalance_experts(
76
+ new_physical_to_logical_map = eplb_state.policy.rebalance_experts(
77
global_expert_load_window,
78
eplb_stats.num_replicas,
79
eplb_stats.num_groups,
@@ -89,16 +85,6 @@ def run_rebalance_experts(
85
86
model_state.new_physical_to_logical_map = new_physical_to_logical_map
87
92
- max_slots = model_state.logical_to_physical_map.shape[-1]
93
- padded_logical = torch.nn.functional.pad(
94
- new_logical_to_physical_map,
95
- (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
96
- value=-1,
97
- ).to(model_state.logical_to_physical_map.device)
98
- new_replica = new_logical_replica_count.to(model_state.logical_replica_count.device)
99
- model_state.new_logical_to_physical_map = padded_logical
100
- model_state.new_logical_replica_count = new_replica
101
-
88
89
async def transfer_run_periodically(
90
state: "EplbState",
vllm/distributed/eplb/eplb_state.py
+113
-59
@@ -235,16 +235,6 @@ class EplbModelState:
235
intermediate variable between `move_to_buffer` and `move_to_workspace`.
236
the size is same as physical_to_logical_map
237
"""
238
- new_logical_to_physical_map: torch.Tensor | None = None
239
- """
240
- intermediate variable between `move_to_buffer` and `move_to_workspace`.
241
- the size is same as logical_to_physical_map
242
- """
243
- new_logical_replica_count: torch.Tensor | None = None
244
- """
245
- intermediate variable between `move_to_buffer` and `move_to_workspace`.
246
- the size is same as logical_replica_count
247
- """
238
239
240
class EplbState:
@@ -508,8 +498,6 @@ class EplbState:
498
),
499
cuda_device_index=self.cuda_device_index,
500
new_physical_to_logical_map=None,
511
- new_logical_to_physical_map=None,
512
- new_logical_replica_count=None,
501
)
502
self.model_states[model_config.compute_hash()] = model_state
503
self.num_valid_physical_experts = model.num_physical_experts
@@ -738,17 +726,20 @@ class EplbState:
726
):
727
if not self.is_async or is_profile:
728
# Get new expert mappings for the model
741
- (
742
- new_physical_to_logical_map,
743
- new_logical_to_physical_map,
744
- new_logical_replica_count,
745
- ) = self.policy.rebalance_experts(
746
- global_expert_load_window,
729
+ new_physical_to_logical_map = self.policy.rebalance_experts(
730
+ global_expert_load_window.cpu(),
731
num_replicas,
732
num_groups,
733
num_nodes,
734
num_gpus,
751
- eplb_model_state.physical_to_logical_map,
735
+ eplb_model_state.physical_to_logical_map.cpu(),
736
+ )
737
+
738
+ num_logical_experts = global_expert_load_window.shape[-1]
739
+ (new_logical_to_physical_map, new_logical_replica_count) = (
740
+ compute_logical_maps(
741
+ new_physical_to_logical_map, num_logical_experts
742
+ )
743
)
744
745
# Update expert weights
@@ -847,11 +838,7 @@ class EplbState:
838
def _update_layer_mapping_from_new(
839
self, model_state: EplbModelState, layer: int
840
) -> None:
850
- if (
851
- model_state.new_physical_to_logical_map is None
852
- or model_state.new_logical_to_physical_map is None
853
- or model_state.new_logical_replica_count is None
854
- ):
841
+ if model_state.new_physical_to_logical_map is None:
842
return
843
844
target_device = model_state.physical_to_logical_map.device
@@ -865,19 +852,23 @@ class EplbState:
852
new_physical[layer].to(target_device, non_blocking=True)
853
)
854
855
+ num_logical_experts = model_state.logical_to_physical_map.shape[1]
856
+ new_logical, new_replica_count = compute_logical_maps(
857
+ new_physical[layer], num_logical_experts
858
+ )
859
+
860
logical_device = model_state.logical_to_physical_map.device
869
- new_logical = model_state.new_logical_to_physical_map[layer].to(logical_device)
861
max_slots = model_state.logical_to_physical_map.shape[-1]
862
slot_delta = max_slots - new_logical.shape[-1]
863
if slot_delta > 0:
864
new_logical = torch.nn.functional.pad(
865
new_logical, (0, slot_delta), value=-1
866
)
876
- model_state.logical_to_physical_map[layer].copy_(new_logical)
867
+ model_state.logical_to_physical_map[layer].copy_(new_logical.to(logical_device))
868
869
replica_device = model_state.logical_replica_count.device
870
model_state.logical_replica_count[layer].copy_(
880
- model_state.new_logical_replica_count[layer].to(replica_device)
871
+ new_replica_count.to(replica_device)
872
)
873
874
def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
@@ -966,7 +957,7 @@ class EplbState:
957
transferred_layer,
958
)
959
if model_state.layer_to_transfer >= model_state.model.num_moe_layers:
969
- self.post_eplb(model_state, is_profile)
960
+ self.post_eplb(model_state)
961
model_state.rebalanced = False
962
model_state.layer_to_transfer = 0
963
model_state.pending_global_ready_check = False
@@ -987,14 +978,9 @@ class EplbState:
978
str(e),
979
)
980
990
- def post_eplb(self, model_state: EplbModelState, is_profile: bool = False) -> None:
981
+ def post_eplb(self, model_state: EplbModelState) -> None:
982
assert model_state.new_physical_to_logical_map is not None
992
- assert model_state.new_logical_to_physical_map is not None
993
- assert model_state.new_logical_replica_count is not None
994
-
983
model_state.new_physical_to_logical_map = None
996
- model_state.new_logical_to_physical_map = None
997
- model_state.new_logical_replica_count = None
984
985
def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
986
"""
@@ -1052,39 +1038,28 @@ class EplbState:
1038
model_config=model_config,
1039
)
1040
eplb_state.num_valid_physical_experts = num_valid_physical_experts
1055
- num_moe_layers = expanded_physical_to_logical.shape[0]
1056
- num_physical_experts = expanded_physical_to_logical.shape[1]
1041
eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
1042
eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)
1043
1060
- logical_to_physical_map = torch.full(
1044
+ (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps(
1045
+ expanded_physical_to_logical.cpu(), model.num_logical_experts
1046
+ )
1047
+
1048
+ max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1]
1049
+ num_replicas = logical_to_physical_map_cpu.shape[-1]
1050
+ logical_to_physical_map = torch.nn.functional.pad(
1051
+ logical_to_physical_map_cpu,
1052
(
1062
- num_moe_layers,
1063
- model.num_logical_experts,
1064
- eplb_model_state.logical_to_physical_map.shape[2],
1053
+ 0,
1054
+ max_num_replicas - num_replicas,
1055
),
1066
- -1,
1067
- dtype=torch.int64,
1068
- )
1069
- logical_replica_count = torch.zeros(
1070
- (num_moe_layers, model.num_logical_experts),
1071
- dtype=torch.int64,
1072
- )
1073
- expanded_physical_to_logical_numpy = expanded_physical_to_logical.cpu().numpy()
1074
- for layer_idx in range(num_moe_layers):
1075
- for phys_idx in range(num_physical_experts):
1076
- logical_idx = expanded_physical_to_logical_numpy[layer_idx, phys_idx]
1077
- if logical_idx >= 0:
1078
- replica_idx = logical_replica_count[layer_idx, logical_idx]
1079
- logical_to_physical_map[layer_idx, logical_idx, replica_idx] = (
1080
- phys_idx
1081
- )
1082
- logical_replica_count[layer_idx, logical_idx] += 1
1056
+ value=-1,
1057
+ ).to(device)
1058
+ logical_replica_count = logical_replica_count_cpu.to(device)
1059
1084
- logical_to_physical_map = logical_to_physical_map.to(device)
1085
- logical_replica_count = logical_replica_count.to(device)
1060
eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)
1061
eplb_model_state.logical_replica_count.copy_(logical_replica_count)
1062
+
1063
return eplb_state
1064
1065
@@ -1132,3 +1107,82 @@ def _node_count_with_rank_mapping(
1107
node_assignment[other_rank] = next_node_id
1108
1109
return next_node_id
1110
+
1111
+
1112
+def compute_logical_maps(
1113
+ physical_to_logical_map: torch.Tensor,
1114
+ num_logical_experts: int,
1115
+) -> tuple[torch.Tensor, torch.Tensor]:
1116
+ """
1117
+ Derive logical_to_physical_map and logical_replica_count from
1118
+ physical_to_logical_map.
1119
+
1120
+ Args:
1121
+ physical_to_logical_map: [num_layers, num_physical_experts], logical
1122
+ expert index for each physical expert slot
1123
+ num_logical_experts: total number of logical experts
1124
+
1125
+ Returns:
1126
+ logical_to_physical_map: [num_layers, num_logical_experts, max_replicas],
1127
+ physical slots per logical expert; -1 where unused
1128
+ logical_replica_count: [num_layers, num_logical_experts], number of
1129
+ physical replicas per logical expert
1130
+ """
1131
+ device = physical_to_logical_map.device
1132
+ assert physical_to_logical_map.device.type == "cpu"
1133
+
1134
+ dtype = physical_to_logical_map.dtype
1135
+
1136
+ # If computing maps for a single layer, unsqueeze a single element layer dimension
1137
+ per_layer = physical_to_logical_map.dim() == 1
1138
+ physical_to_logical_map_view = physical_to_logical_map
1139
+ if per_layer:
1140
+ physical_to_logical_map_view = physical_to_logical_map.unsqueeze(0)
1141
+ assert len(physical_to_logical_map_view.shape) == 2
1142
+ num_layers, num_physical = physical_to_logical_map_view.shape
1143
+
1144
+ valid_mask = physical_to_logical_map_view >= 0
1145
+ logical_replica_count = torch.zeros(
1146
+ num_layers,
1147
+ num_logical_experts,
1148
+ dtype=dtype,
1149
+ device=device,
1150
+ )
1151
+ logical_replica_count.scatter_add_(
1152
+ 1,
1153
+ physical_to_logical_map_view.clamp(min=0),
1154
+ valid_mask.to(dtype),
1155
+ )
1156
+
1157
+ max_replicas = int(logical_replica_count.max().item())
1158
+ logical_to_physical_map_out = torch.full(
1159
+ (num_layers, num_logical_experts, max_replicas),
1160
+ -1,
1161
+ dtype=dtype,
1162
+ device=device,
1163
+ )
1164
+
1165
+ running_count = torch.zeros_like(logical_replica_count)
1166
+ layer_indices = torch.arange(num_layers, device=device)
1167
+ for phys_idx in range(num_physical):
1168
+ # Logical expert at physical slot phys_idx for each layer
1169
+ logical_expert_ids = physical_to_logical_map_view[:, phys_idx] # [num_layers]
1170
+
1171
+ # Scale up will set the logical expert ids to -1 for all new physical experts.
1172
+ # Only consider "valid" experts when setting up the logical_to_physical map.
1173
+ valid_expert_mask = logical_expert_ids >= 0
1174
+ if not valid_expert_mask.any():
1175
+ continue
1176
+ valid_layers = layer_indices[valid_expert_mask]
1177
+ valid_experts = logical_expert_ids[valid_expert_mask]
1178
+
1179
+ # Use the current running count as the replica index, then increment it.
1180
+ replica_idx = running_count[valid_layers, valid_experts]
1181
+ logical_to_physical_map_out[valid_layers, valid_experts, replica_idx] = phys_idx
1182
+ running_count[valid_layers, valid_experts] += 1
1183
+
1184
+ # If computing maps for a single layer, squeeze out the extra layer dimension
1185
+ # before returning
1186
+ if per_layer:
1187
+ return logical_to_physical_map_out.squeeze(0), logical_replica_count.squeeze(0)
1188
+ return logical_to_physical_map_out, logical_replica_count
vllm/distributed/eplb/policy/abstract.py
+1
-5
@@ -17,7 +17,7 @@ class AbstractEplbPolicy(ABC):
17
num_nodes: int,
18
num_ranks: int,
19
old_global_expert_indices: torch.Tensor | None = None,
20
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
20
+ ) -> torch.Tensor:
21
"""
22
Entry point for expert-parallelism load balancer.
23
@@ -35,9 +35,5 @@ class AbstractEplbPolicy(ABC):
35
Returns:
36
physical_to_logical_map: [layers, num_replicas], the expert
37
index of each replica
38
- logical_to_physical_map: [layers, num_logical_experts, X],
39
- the replica indices for each expert
40
- expert_count: [layers, num_logical_experts], number of
41
- physical replicas for each logical expert
38
"""
39
raise NotImplementedError
vllm/distributed/eplb/policy/default.py
+18
-62
@@ -75,7 +75,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
75
@classmethod
76
def replicate_experts(
77
cls, weight: np.ndarray, num_phy: int
78
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
78
+ ) -> tuple[np.ndarray, np.ndarray]:
79
"""
80
Replicate `num_log` experts to `num_phy` replicas, such that the maximum
81
load of all replicas is minimized.
@@ -86,22 +86,19 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
86
87
Returns:
88
phy2log: [X, num_phy], logical expert id of each physical expert
89
- replica_idx: [X, num_phy], the index of the replica for each logical expert
89
logcnt: [X, num_log], number of replicas for each logical expert
90
"""
91
n, num_log = weight.shape
92
num_redundant = num_phy - num_log
93
assert num_redundant >= 0
94
phy2log = np.tile(np.arange(num_phy, dtype=np.int64), (n, 1))
96
- replica_idx = np.zeros((n, num_phy), dtype=np.int64)
95
logcnt = np.ones((n, num_log), dtype=np.int64)
96
arangen = np.arange(n, dtype=np.int64)
97
for i in range(num_log, num_phy):
98
redundant_indices = np.argmax(weight / logcnt, axis=-1)
99
phy2log[:, i] = redundant_indices
102
- replica_idx[:, i] = logcnt[arangen, redundant_indices]
100
logcnt[arangen, redundant_indices] += 1
104
- return phy2log, replica_idx, logcnt
101
+ return phy2log, logcnt
102
103
@classmethod
104
def rebalance_experts_hierarchical(
@@ -111,7 +108,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
108
num_groups: int,
109
num_nodes: int,
110
num_gpus: int,
114
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
111
+ ) -> np.ndarray:
112
"""
113
Parameters:
114
weight: [num_moe_layers, num_logical_experts]
@@ -124,10 +121,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
121
Returns:
122
phy2log: [layers, num_replicas], the expert
123
index of each replica
127
- pphy_replicas_idx: [layers, num_logical_experts, X],
128
- the replica indices for each expert
129
- logcnt: [layers, num_logical_experts], number of
130
- physical replicas for each logical expert
124
"""
125
num_layers, num_logical_experts = weight.shape
126
assert num_logical_experts % num_groups == 0
@@ -167,7 +160,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
160
tokens_per_mlog = np.take_along_axis(weight, mlog2log, axis=1).reshape(
161
-1, num_logical_experts // num_nodes
162
)
170
- phy2mlog, replicas_idx, mlogcnt = cls.replicate_experts(
163
+ phy2mlog, mlogcnt = cls.replicate_experts(
164
tokens_per_mlog, num_physical_experts // num_nodes
165
)
166
@@ -193,22 +186,15 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
186
).reshape(num_layers, -1)
187
# Map node-local logical indices back to global logical expert ids.
188
pphy2log = np.take_along_axis(mlog2log, pphy2mlog, axis=1)
196
- # Reorder replica ranks to the post-packing physical ordering.
197
- pphy_replicas_idx = np.take_along_axis(replicas_idx, pphy2phy, axis=1).reshape(
198
- num_layers, -1
199
- )
200
- # Convert replica counts back to the original logical ordering.
201
- logcnt = np.take_along_axis(mlogcnt.reshape(num_layers, -1), log2mlog, axis=1)
202
- return pphy2log, pphy_replicas_idx, logcnt
189
+ return pphy2log
190
191
@classmethod
192
def preserve_intragpu_slots(
193
cls,
194
phy2log: np.ndarray,
208
- phy_replicas_idx: np.ndarray,
195
num_ranks: int,
196
old_phy2log: np.ndarray,
211
- ) -> tuple[np.ndarray, np.ndarray]:
197
+ ) -> np.ndarray:
198
"""
199
Reorder the new mapping per GPU so that experts that remain on the same GPU
200
keep their previous slot positions when possible. Incoming experts to that GPU
@@ -218,14 +204,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
204
"""
205
num_phy_experts = phy2log.shape[1]
206
if num_ranks <= 0 or num_phy_experts % num_ranks != 0:
221
- return phy2log, phy_replicas_idx
207
+ return phy2log
208
209
# Move to CPU and convert to NumPy for processing
210
slots_per_gpu = num_phy_experts // num_ranks
211
num_layers = phy2log.shape[0]
212
213
post_phy2log = phy2log.copy()
228
- post_phy_replicas_idx = phy_replicas_idx.copy()
214
215
for gpu_idx in range(num_ranks):
216
start = gpu_idx * slots_per_gpu
@@ -233,7 +218,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
218
# Experts across all layers for this GPU
219
old_local = old_phy2log[:, start:end] # [layers, slots]
220
new_local = phy2log[:, start:end] # [layers, slots]
236
- new_ridx = phy_replicas_idx[:, start:end] # [layers, slots]
221
222
used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool)
223
preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool)
@@ -253,9 +237,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
237
post_phy2log[layer_indices, start + slot_idx] = new_local[
238
layer_indices, matched_new_positions
239
]
256
- post_phy_replicas_idx[layer_indices, start + slot_idx] = new_ridx[
257
- layer_indices, matched_new_positions
258
- ]
240
used_new_indices[layer_indices, matched_new_positions] = True
241
preserved_positions[layer_indices, slot_idx] = True
242
@@ -287,11 +268,8 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
268
post_phy2log[layer_idx, start + dst_pos] = new_local[
269
layer_idx, src_pos
270
]
290
- post_phy_replicas_idx[layer_idx, start + dst_pos] = new_ridx[
291
- layer_idx, src_pos
292
- ]
271
294
- return post_phy2log, post_phy_replicas_idx
272
+ return post_phy2log
273
274
@classmethod
275
def rebalance_experts(
@@ -302,7 +280,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
280
num_nodes: int,
281
num_ranks: int,
282
old_global_expert_indices: torch.Tensor | None = None,
305
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
283
+ ) -> torch.Tensor:
284
"""
285
Entry point for expert-parallelism load balancer.
286
@@ -321,13 +299,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
299
Returns:
300
phy2log: [layers, num_replicas], the expert
301
index of each replica
324
- log2phy: [layers, num_logical_experts, X],
325
- the replica indices for each expert
326
- logcnt: [layers, num_logical_experts], number of
327
- physical replicas for each logical expert
302
"""
329
- device = weight.device
330
- num_layers, num_logical_experts = weight.shape
303
weight_np = weight.float().cpu().numpy()
304
old_phy2log_np = (
305
old_global_expert_indices.cpu().numpy()
@@ -337,17 +309,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
309
310
if num_groups % num_nodes == 0:
311
# use hierarchical load-balance policy
340
- phy2log_np, phy_replicas_idx_np, logcnt_np = (
341
- cls.rebalance_experts_hierarchical(
342
- weight_np, num_replicas, num_groups, num_nodes, num_ranks
343
- )
312
+ phy2log_np = cls.rebalance_experts_hierarchical(
313
+ weight_np, num_replicas, num_groups, num_nodes, num_ranks
314
)
315
else:
316
# use global load-balance policy
347
- phy2log_np, phy_replicas_idx_np, logcnt_np = (
348
- cls.rebalance_experts_hierarchical(
349
- weight_np, num_replicas, 1, 1, num_ranks
350
- )
317
+ phy2log_np = cls.rebalance_experts_hierarchical(
318
+ weight_np, num_replicas, 1, 1, num_ranks
319
)
320
321
# Optional postprocessing to preserve slots for experts moving
@@ -355,22 +323,10 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
323
# Only apply when the number of GPUs and slots per GPU remain unchanged.
324
# Helps to avoid unnecessary weight copying when experts move
325
# within the same GPU.
358
- if old_global_expert_indices is not None:
359
- phy2log_np, phy_replicas_idx_np = cls.preserve_intragpu_slots(
360
- phy2log_np, phy_replicas_idx_np, num_ranks, old_phy2log_np
326
+ if old_phy2log_np is not None:
327
+ phy2log_np = cls.preserve_intragpu_slots(
328
+ phy2log_np, num_ranks, old_phy2log_np
329
)
362
- num_redundant_experts = num_replicas - num_logical_experts
363
- maxlogcnt = num_redundant_experts + 1
364
- log2phy_np = np.full(
365
- (num_layers, num_logical_experts, maxlogcnt), -1, dtype=np.int64
366
- )
367
- layer_indices = np.arange(num_layers)[:, None]
368
- replica_indices = np.tile(
369
- np.arange(num_replicas, dtype=np.int64), (num_layers, 1)
370
- )
371
- log2phy_np[layer_indices, phy2log_np, phy_replicas_idx_np] = replica_indices
330
373
- phy2log = torch.from_numpy(phy2log_np).to(device)
374
- log2phy = torch.from_numpy(log2phy_np).to(device)
375
- logcnt = torch.from_numpy(logcnt_np).to(device)
376
- return phy2log, log2phy, logcnt
331
+ phy2log = torch.from_numpy(phy2log_np)
332
+ return phy2log