| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 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 numpy as np |
| 15 | import torch |
| 16 | |
| 17 | from .abstract import AbstractEplbPolicy |
| 18 | |
| 19 | |
| 20 | class DefaultEplbPolicy(AbstractEplbPolicy): |
| 21 | @classmethod |
| 22 | def balanced_packing( |
| 23 | cls, weight: np.ndarray, num_packs: int |
| 24 | ) -> tuple[np.ndarray, np.ndarray]: |
| 25 | """Pack n weighted objects to m packs, such that each bin contains exactly |
| 26 | n/m objects and the weights of all packs are as balanced as possible. |
| 27 | |
| 28 | Args: |
| 29 | weight: [X, n], the weight of each item |
| 30 | num_packs: number of packs |
| 31 | |
| 32 | Returns: |
| 33 | pack_index: [X, n], the pack index of each item |
| 34 | rank_in_pack: [X, n], the rank of the item in the pack |
| 35 | |
| 36 | """ |
| 37 | num_layers, num_groups = weight.shape |
| 38 | assert num_groups % num_packs == 0 |
| 39 | groups_per_pack = num_groups // num_packs |
| 40 | |
| 41 | if groups_per_pack == 1: |
| 42 | pack_index = np.tile(np.arange(num_groups, dtype=np.int64), (num_layers, 1)) |
| 43 | rank_in_pack = np.zeros_like(pack_index, dtype=np.int64) |
| 44 | return pack_index, rank_in_pack |
| 45 | |
| 46 | # Sort and get indices in descending order |
| 47 | indices = np.argsort(-weight, axis=-1) |
| 48 | |
| 49 | pack_index = np.full((num_layers, num_groups), -1, dtype=np.int64) |
| 50 | rank_in_pack = np.full((num_layers, num_groups), -1, dtype=np.int64) |
| 51 | |
| 52 | pack_weights = np.zeros((num_layers, num_packs), dtype=np.float64) |
| 53 | pack_items = np.zeros((num_layers, num_packs), dtype=np.int64) |
| 54 | |
| 55 | # Run the packing algorithm |
| 56 | for layer_idx in range(num_layers): |
| 57 | weights_row = pack_weights[layer_idx] |
| 58 | items_row = pack_items[layer_idx] |
| 59 | |
| 60 | for group in indices[layer_idx]: |
| 61 | # Pick the lightest pack; full packs are masked out by inf. |
| 62 | pack = int(np.argmin(weights_row)) |
| 63 | |
| 64 | pack_index[layer_idx, group] = pack |
| 65 | rank_in_pack[layer_idx, group] = items_row[pack] |
| 66 | weights_row[pack] += weight[layer_idx, group] |
| 67 | items_row[pack] += 1 |
| 68 | if items_row[pack] == groups_per_pack: |
| 69 | # Mark as unavailable for future selections. |
| 70 | weights_row[pack] = np.inf |
| 71 | |
| 72 | return pack_index, rank_in_pack |
| 73 | |
| 74 | @classmethod |
| 75 | def replicate_experts( |
| 76 | cls, weight: np.ndarray, num_phy: int |
| 77 | ) -> tuple[np.ndarray, np.ndarray]: |
| 78 | """Replicate `num_log` experts to `num_phy` replicas, such that the maximum |
| 79 | load of all replicas is minimized. |
| 80 | |
| 81 | Args: |
| 82 | weight: [X, num_log] |
| 83 | num_phy: total number of experts after replication |
| 84 | |
| 85 | Returns: |
| 86 | phy2log: [X, num_phy], logical expert id of each physical expert |
| 87 | logcnt: [X, num_log], number of replicas for each logical expert |
| 88 | |
| 89 | """ |
| 90 | n, num_log = weight.shape |
| 91 | num_redundant = num_phy - num_log |
| 92 | assert num_redundant >= 0 |
| 93 | phy2log = np.tile(np.arange(num_phy, dtype=np.int64), (n, 1)) |
| 94 | logcnt = np.ones((n, num_log), dtype=np.int64) |
| 95 | arangen = np.arange(n, dtype=np.int64) |
| 96 | for i in range(num_log, num_phy): |
| 97 | redundant_indices = np.argmax(weight / logcnt, axis=-1) |
| 98 | phy2log[:, i] = redundant_indices |
| 99 | logcnt[arangen, redundant_indices] += 1 |
| 100 | return phy2log, logcnt |
| 101 | |
| 102 | @classmethod |
| 103 | def rebalance_experts_hierarchical( |
| 104 | cls, |
| 105 | weight: np.ndarray, |
| 106 | num_physical_experts: int, |
| 107 | num_groups: int, |
| 108 | num_nodes: int, |
| 109 | num_gpus: int, |
| 110 | ) -> np.ndarray: |
| 111 | """Parameters |
| 112 | weight: [num_moe_layers, num_logical_experts] |
| 113 | num_physical_experts: number of physical experts after replication |
| 114 | num_groups: number of expert groups |
| 115 | num_nodes: number of server nodes, where the intra-node network |
| 116 | (e.g, NVLink) is faster |
| 117 | num_gpus: number of GPUs, must be a multiple of `num_nodes` |
| 118 | |
| 119 | Returns: |
| 120 | phy2log: [layers, num_replicas], the expert |
| 121 | index of each replica |
| 122 | |
| 123 | """ |
| 124 | num_layers, num_logical_experts = weight.shape |
| 125 | assert num_logical_experts % num_groups == 0 |
| 126 | group_size = num_logical_experts // num_groups |
| 127 | assert num_groups % num_nodes == 0 |
| 128 | groups_per_node = num_groups // num_nodes |
| 129 | assert num_gpus % num_nodes == 0 |
| 130 | assert num_physical_experts % num_gpus == 0 |
| 131 | phy_experts_per_gpu = num_physical_experts // num_gpus |
| 132 | |
| 133 | def inverse(perm: np.ndarray) -> np.ndarray: |
| 134 | inv = np.empty_like(perm) |
| 135 | row_idx = np.arange(perm.shape[0])[:, None] |
| 136 | col_idx = np.arange(perm.shape[1], dtype=np.int64) |
| 137 | inv[row_idx, perm] = col_idx |
| 138 | return inv |
| 139 | |
| 140 | # Step 1: pack groups to nodes |
| 141 | tokens_per_group = weight.reshape(num_layers, num_groups, group_size).sum( |
| 142 | axis=-1 |
| 143 | ) |
| 144 | group_pack_index, group_rank_in_pack = cls.balanced_packing( |
| 145 | tokens_per_group, num_nodes |
| 146 | ) |
| 147 | # Map each logical expert into a node-local ordering based on packed groups. |
| 148 | log2mlog = ( |
| 149 | ( |
| 150 | (group_pack_index * groups_per_node + group_rank_in_pack)[..., None] |
| 151 | * group_size |
| 152 | ) |
| 153 | + np.arange(group_size, dtype=np.int64) |
| 154 | ).reshape(num_layers, num_logical_experts) |
| 155 | mlog2log = inverse(log2mlog) |
| 156 | |
| 157 | # Step 2: construct redundant experts within nodes |
| 158 | # Reorder weights into the node-local layout so replication is done per node. |
| 159 | tokens_per_mlog = np.take_along_axis(weight, mlog2log, axis=1).reshape( |
| 160 | -1, num_logical_experts // num_nodes |
| 161 | ) |
| 162 | phy2mlog, mlogcnt = cls.replicate_experts( |
| 163 | tokens_per_mlog, num_physical_experts // num_nodes |
| 164 | ) |
| 165 | |
| 166 | # Step 3: pack physical_experts to GPUs |
| 167 | # Effective per-physical load = logical load divided by replica count. |
| 168 | tokens_per_phy = np.take_along_axis(tokens_per_mlog / mlogcnt, phy2mlog, axis=1) |
| 169 | pack_index, rank_in_pack = cls.balanced_packing( |
| 170 | tokens_per_phy, num_gpus // num_nodes |
| 171 | ) |
| 172 | phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack |
| 173 | pphy2phy = inverse(phy2pphy) |
| 174 | |
| 175 | # Reorder node-local logical indices into the post-packing physical order. |
| 176 | pphy2mlog = np.take_along_axis(phy2mlog, pphy2phy, axis=1) |
| 177 | pphy2mlog = ( |
| 178 | pphy2mlog.reshape(num_layers, num_nodes, -1) |
| 179 | + np.arange( |
| 180 | 0, |
| 181 | num_logical_experts, |
| 182 | num_logical_experts // num_nodes, |
| 183 | dtype=np.int64, |
| 184 | )[None, :, None] |
| 185 | ).reshape(num_layers, -1) |
| 186 | # Map node-local logical indices back to global logical expert ids. |
| 187 | pphy2log = np.take_along_axis(mlog2log, pphy2mlog, axis=1) |
| 188 | return pphy2log |
| 189 | |
| 190 | @classmethod |
| 191 | def preserve_intragpu_slots( |
| 192 | cls, |
| 193 | phy2log: np.ndarray, |
| 194 | num_ranks: int, |
| 195 | old_phy2log: np.ndarray, |
| 196 | ) -> np.ndarray: |
| 197 | """Reorder the new mapping per GPU so that experts that remain on the same GPU |
| 198 | keep their previous slot positions when possible. Incoming experts to that GPU |
| 199 | fill any remaining available slots. This is applied only when the number of GPUs |
| 200 | is unchanged and the slots per GPU remain the same between |
| 201 | the old and new mappings. |
| 202 | """ |
| 203 | num_phy_experts = phy2log.shape[1] |
| 204 | if num_ranks <= 0 or num_phy_experts % num_ranks != 0: |
| 205 | return phy2log |
| 206 | |
| 207 | # Move to CPU and convert to NumPy for processing |
| 208 | slots_per_gpu = num_phy_experts // num_ranks |
| 209 | num_layers = phy2log.shape[0] |
| 210 | |
| 211 | post_phy2log = phy2log.copy() |
| 212 | |
| 213 | for gpu_idx in range(num_ranks): |
| 214 | start = gpu_idx * slots_per_gpu |
| 215 | end = start + slots_per_gpu |
| 216 | # Experts across all layers for this GPU |
| 217 | old_local = old_phy2log[:, start:end] # [layers, slots] |
| 218 | new_local = phy2log[:, start:end] # [layers, slots] |
| 219 | |
| 220 | used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool) |
| 221 | preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool) |
| 222 | |
| 223 | # First pass: preserve same-logical experts in their previous slots |
| 224 | for slot_idx in range(slots_per_gpu): |
| 225 | # matches: [layers, slots], True where new local experts have |
| 226 | # the same logical value as the old from 'slot_idx' and not checked yet |
| 227 | matches = (new_local == old_local[:, slot_idx][:, None]) & ( |
| 228 | ~used_new_indices |
| 229 | ) |
| 230 | has_any = matches.any(axis=1) |
| 231 | if np.any(has_any): |
| 232 | first_idx = np.argmax(matches, axis=1) |
| 233 | layer_indices = np.nonzero(has_any)[0] |
| 234 | matched_new_positions = first_idx[layer_indices] |
| 235 | post_phy2log[layer_indices, start + slot_idx] = new_local[ |
| 236 | layer_indices, matched_new_positions |
| 237 | ] |
| 238 | used_new_indices[layer_indices, matched_new_positions] = True |
| 239 | preserved_positions[layer_indices, slot_idx] = True |
| 240 | |
| 241 | # Second pass: fill remaining slots with remaining new experts |
| 242 | remaining_mask = ~used_new_indices # [layers, slots] |
| 243 | fill_mask = ~preserved_positions # [layers, slots] |
| 244 | if remaining_mask.any() and fill_mask.any(): |
| 245 | idx_base = np.tile(np.arange(slots_per_gpu), (num_layers, 1)) |
| 246 | # Sentinel value for unavailable positions. |
| 247 | large = slots_per_gpu + 1 |
| 248 | # Priorities: keep original index for available spots, set sentinel |
| 249 | # for unavailable; lower is earlier. |
| 250 | remaining_priority = np.where(remaining_mask, idx_base, large) |
| 251 | fill_priority = np.where(fill_mask, idx_base, large) |
| 252 | # Sort to get ordered indices of available src/dst positions per layer. |
| 253 | remaining_indices = np.argsort(remaining_priority, axis=1) |
| 254 | fill_indices = np.argsort(fill_priority, axis=1) |
| 255 | # Fill count per layer (cannot exceed either side). |
| 256 | remaining_counts = remaining_mask.sum(axis=1) |
| 257 | fill_counts = fill_mask.sum(axis=1) |
| 258 | take_counts = np.minimum(remaining_counts, fill_counts) |
| 259 | # Assign remaining new experts to remaining slots per layer. |
| 260 | for layer_idx in range(num_layers): |
| 261 | k = int(take_counts[layer_idx]) |
| 262 | if k <= 0: |
| 263 | continue |
| 264 | src_pos = remaining_indices[layer_idx, :k] |
| 265 | dst_pos = fill_indices[layer_idx, :k] |
| 266 | post_phy2log[layer_idx, start + dst_pos] = new_local[ |
| 267 | layer_idx, src_pos |
| 268 | ] |
| 269 | |
| 270 | return post_phy2log |
| 271 | |
| 272 | @classmethod |
| 273 | def rebalance_experts( |
| 274 | cls, |
| 275 | weight: torch.Tensor, |
| 276 | num_replicas: int, |
| 277 | num_groups: int, |
| 278 | num_nodes: int, |
| 279 | num_ranks: int, |
| 280 | old_global_expert_indices: torch.Tensor | None = None, |
| 281 | ) -> torch.Tensor: |
| 282 | """Entry point for expert-parallelism load balancer. |
| 283 | |
| 284 | Args: |
| 285 | weight: [layers, num_logical_experts], the load statistics for all |
| 286 | logical experts |
| 287 | num_replicas: number of physical experts, must be a multiple of |
| 288 | `num_gpus` |
| 289 | num_groups: number of expert groups |
| 290 | num_nodes: number of server nodes, where the intra-node network |
| 291 | (e.g, NVLink) is faster |
| 292 | num_ranks: number of ranks, must be a multiple of `num_nodes` |
| 293 | old_global_expert_indices: [layers, num_logical_experts], the old global |
| 294 | expert indices. Used to avoid unnecessary weight copying |
| 295 | for experts moving within one rank. |
| 296 | |
| 297 | Returns: |
| 298 | ------- |
| 299 | phy2log: [layers, num_replicas], the expert |
| 300 | index of each replica |
| 301 | |
| 302 | """ |
| 303 | weight_np = weight.float().cpu().numpy() |
| 304 | old_phy2log_np = ( |
| 305 | old_global_expert_indices.cpu().numpy() |
| 306 | if old_global_expert_indices is not None |
| 307 | else None |
| 308 | ) |
| 309 | |
| 310 | if num_groups % num_nodes == 0: |
| 311 | # use hierarchical load-balance policy |
| 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 |
| 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 |
| 322 | # within the same GPU |
| 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. |
| 326 | if old_phy2log_np is not None: |
| 327 | phy2log_np = cls.preserve_intragpu_slots( |
| 328 | phy2log_np, num_ranks, old_phy2log_np |
| 329 | ) |
| 330 | |
| 331 | phy2log = torch.from_numpy(phy2log_np) |
| 332 | return phy2log |