| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import numpy as np |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.distributed.eplb.eplb_state import ( |
| 9 | _compute_eplb_load_stats, # pyright: ignore[reportPrivateUsage] |
| 10 | compute_logical_maps, |
| 11 | ) |
| 12 | from vllm.distributed.eplb.policy.default import DefaultEplbPolicy |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | |
| 16 | def test_eplb_load_stats_reduce_across_ranks(): |
| 17 | num_tokens_per_rank = torch.tensor( |
| 18 | [ |
| 19 | [100, 0, 0, 0], |
| 20 | [100, 0, 0, 0], |
| 21 | ], |
| 22 | dtype=torch.float32, |
| 23 | ) |
| 24 | |
| 25 | avg_tokens, max_tokens = _compute_eplb_load_stats(num_tokens_per_rank) |
| 26 | |
| 27 | assert avg_tokens.item() == 50 |
| 28 | assert max_tokens.item() == 200 |
| 29 | assert (avg_tokens / max_tokens).item() == 0.25 |
| 30 | |
| 31 | |
| 32 | def test_basic_rebalance(): |
| 33 | """Test basic rebalancing functionality.""" |
| 34 | # Example from https://github.com/deepseek-ai/eplb |
| 35 | weight = torch.tensor( |
| 36 | [ |
| 37 | [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86], |
| 38 | [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27], |
| 39 | ] |
| 40 | ) |
| 41 | |
| 42 | num_layers = weight.shape[0] |
| 43 | num_replicas = 16 |
| 44 | num_groups = 4 |
| 45 | num_nodes = 2 |
| 46 | num_gpus = 8 |
| 47 | |
| 48 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 49 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 50 | ) |
| 51 | log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 52 | |
| 53 | # Verify output shapes |
| 54 | assert phy2log.shape == ( |
| 55 | 2, |
| 56 | 16, |
| 57 | ), f"Expected `phy2log` shape (2, 16), got {phy2log.shape}" |
| 58 | assert log2phy.shape[0] == 2, ( |
| 59 | f"Expected `log2phy` first dimension 2, got {log2phy.shape[0]}" |
| 60 | ) |
| 61 | assert log2phy.shape[1] == 12, ( |
| 62 | f"Expected `log2phy` second dimension 12, got {log2phy.shape[1]}" |
| 63 | ) |
| 64 | assert logcnt.shape == ( |
| 65 | 2, |
| 66 | 12, |
| 67 | ), f"Expected `logcnt` shape (2, 12), got {logcnt.shape}" |
| 68 | |
| 69 | # Verify physical to logical expert mapping range is correct |
| 70 | assert torch.all(phy2log >= 0) and torch.all(phy2log < 12), ( |
| 71 | "Physical to logical mapping should be in range [0, 12)" |
| 72 | ) |
| 73 | |
| 74 | # Verify expert count reasonableness |
| 75 | assert torch.all(logcnt >= 1), "Each logical expert should have at least 1 replica" |
| 76 | assert torch.sum(logcnt, dim=1).sum() == num_replicas * num_layers, ( |
| 77 | f"Total replicas should be {num_replicas * num_layers}" |
| 78 | ) |
| 79 | |
| 80 | # Verify expected output |
| 81 | expected_phy2log = torch.tensor( |
| 82 | [ |
| 83 | [5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1], |
| 84 | [7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1], |
| 85 | ] |
| 86 | ) |
| 87 | assert torch.all(phy2log == expected_phy2log) |
| 88 | |
| 89 | expected_logcnt = torch.tensor( |
| 90 | [[1, 2, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1], [1, 2, 1, 1, 1, 2, 2, 1, 2, 1, 1, 1]] |
| 91 | ) |
| 92 | assert torch.all(logcnt == expected_logcnt) |
| 93 | |
| 94 | |
| 95 | def test_single_gpu_case(): |
| 96 | """Test single GPU case.""" |
| 97 | weight = torch.tensor([[10, 20, 30, 40]]) |
| 98 | num_replicas = 4 |
| 99 | num_groups = 1 |
| 100 | num_nodes = 1 |
| 101 | num_gpus = 1 |
| 102 | |
| 103 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 104 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 105 | ) |
| 106 | log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 107 | |
| 108 | # Verify shapes |
| 109 | assert phy2log.shape == (1, 4) |
| 110 | assert log2phy.shape[0] == 1 |
| 111 | assert log2phy.shape[1] == 4 |
| 112 | assert logcnt.shape == (1, 4) |
| 113 | |
| 114 | # Verify all logical experts are mapped |
| 115 | assert set(phy2log[0].tolist()) == {0, 1, 2, 3} |
| 116 | |
| 117 | |
| 118 | def test_equal_weights(): |
| 119 | """Test case with equal weights.""" |
| 120 | weight = torch.tensor([[50, 50, 50, 50, 50, 50, 50, 50]]) |
| 121 | num_replicas = 8 |
| 122 | num_groups = 2 |
| 123 | num_nodes = 2 |
| 124 | num_gpus = 4 |
| 125 | |
| 126 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 127 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 128 | ) |
| 129 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 130 | |
| 131 | # Verify shapes |
| 132 | assert phy2log.shape == (1, 8) |
| 133 | assert logcnt.shape == (1, 8) |
| 134 | |
| 135 | # With equal weights, each expert should have exactly one replica |
| 136 | assert torch.all(logcnt == 1), ( |
| 137 | "With equal weights and no replication, " |
| 138 | "each expert should have exactly 1 replica" |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | def test_extreme_weight_imbalance(): |
| 143 | """Test extreme weight imbalance case.""" |
| 144 | weight = torch.tensor([[1000, 1, 1, 1, 1, 1, 1, 1]]) |
| 145 | num_replicas = 12 |
| 146 | num_groups = 2 |
| 147 | num_nodes = 2 |
| 148 | num_gpus = 4 |
| 149 | |
| 150 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 151 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 152 | ) |
| 153 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 154 | |
| 155 | # Verify shapes |
| 156 | assert phy2log.shape == (1, 12) |
| 157 | assert logcnt.shape == (1, 8) |
| 158 | |
| 159 | # Expert with highest weight (index 0) should have more replicas |
| 160 | assert logcnt[0, 0] > logcnt[0, 1], ( |
| 161 | "Expert with highest weight should have more replicas" |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | def test_multiple_layers(): |
| 166 | """Test multiple layers case.""" |
| 167 | weight = torch.tensor( |
| 168 | [ |
| 169 | [10, 20, 30, 40, 50, 60], # First layer |
| 170 | [60, 50, 40, 30, 20, 10], # Second layer (opposite weight pattern) |
| 171 | [25, 25, 25, 25, 25, 25], # Third layer (equal weights) |
| 172 | ] |
| 173 | ) |
| 174 | num_replicas = 8 |
| 175 | num_groups = 2 |
| 176 | num_nodes = 2 |
| 177 | num_gpus = 4 |
| 178 | |
| 179 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 180 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 181 | ) |
| 182 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 183 | |
| 184 | # Verify shapes |
| 185 | assert phy2log.shape == (3, 8) |
| 186 | assert logcnt.shape == (3, 6) |
| 187 | |
| 188 | # Verify expert allocation is reasonable for each layer |
| 189 | for layer in range(3): |
| 190 | assert torch.all(phy2log[layer] >= 0) and torch.all(phy2log[layer] < 6), ( |
| 191 | f"Layer {layer} physical to logical mappingshould be in range [0, 6)" |
| 192 | ) |
| 193 | assert torch.sum(logcnt[layer]) == num_replicas, ( |
| 194 | f"Layer {layer} total replicas should be {num_replicas}" |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | def test_parameter_validation(): |
| 199 | """Test parameter validation.""" |
| 200 | weight = torch.tensor([[10, 20, 30, 40]]) |
| 201 | |
| 202 | # Test non-divisible case - this should handle normally without throwing |
| 203 | # errors because the function will fall back to global load balancing |
| 204 | # strategy |
| 205 | phy2log = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4) |
| 206 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 207 | assert phy2log.shape == (1, 8) |
| 208 | assert logcnt.shape == (1, 4) |
| 209 | |
| 210 | # Test cases that will actually cause errors: |
| 211 | # num_physical_experts not divisible by num_gpus |
| 212 | with pytest.raises(AssertionError): |
| 213 | DefaultEplbPolicy.rebalance_experts(weight, 7, 2, 2, 4) # 7 not divisible by 4 |
| 214 | |
| 215 | |
| 216 | def test_small_scale_hierarchical(): |
| 217 | """Test small-scale hierarchical load balancing.""" |
| 218 | weight = torch.tensor( |
| 219 | [ |
| 220 | [100, 50, 200, 75, 150, 25, 300, 80], # 8 experts |
| 221 | ] |
| 222 | ) |
| 223 | num_replicas = 12 |
| 224 | num_groups = 4 # 4 groups, 2 experts each |
| 225 | num_nodes = 2 # 2 nodes |
| 226 | num_gpus = 4 # 4 GPUs |
| 227 | |
| 228 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 229 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 230 | ) |
| 231 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 232 | |
| 233 | # Verify basic constraints |
| 234 | assert phy2log.shape == (1, 12) |
| 235 | assert logcnt.shape == (1, 8) |
| 236 | assert torch.sum(logcnt) == num_replicas |
| 237 | assert torch.all(logcnt >= 1) |
| 238 | |
| 239 | # Expert with highest weight should have more replicas |
| 240 | max_weight_expert = torch.argmax(weight[0]) |
| 241 | assert logcnt[0, max_weight_expert] >= 2, ( |
| 242 | "Highest weight expert should have multiple replicas" |
| 243 | ) |
| 244 | |
| 245 | |
| 246 | def test_global_load_balance_fallback(): |
| 247 | """Test global load balancing fallback case.""" |
| 248 | # When num_groups % num_nodes != 0, should fall back to global load |
| 249 | # balancing |
| 250 | weight = torch.tensor([[10, 20, 30, 40, 50, 60]]) |
| 251 | num_replicas = 8 |
| 252 | num_groups = 3 # Cannot be divided evenly by num_nodes=2 |
| 253 | num_nodes = 2 |
| 254 | num_gpus = 4 |
| 255 | |
| 256 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 257 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 258 | ) |
| 259 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 260 | |
| 261 | # Should work normally, just using global load balancing strategy |
| 262 | assert phy2log.shape == (1, 8) |
| 263 | assert logcnt.shape == (1, 6) |
| 264 | assert torch.sum(logcnt) == num_replicas |
| 265 | |
| 266 | |
| 267 | @pytest.mark.parametrize("device", ["cpu", "cuda", "xpu"]) |
| 268 | def test_device_compatibility(device): |
| 269 | """Test device compatibility.""" |
| 270 | if device == "cuda" and not torch.cuda.is_available(): |
| 271 | pytest.skip("CUDA not available") |
| 272 | if device == "xpu" and not current_platform.is_xpu(): |
| 273 | pytest.skip("XPU not available") |
| 274 | |
| 275 | weight = torch.tensor([[10, 20, 30, 40]], device=device) |
| 276 | num_replicas = 6 |
| 277 | num_groups = 2 |
| 278 | num_nodes = 1 |
| 279 | num_gpus = 2 |
| 280 | |
| 281 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 282 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 283 | ) |
| 284 | _, logcnt = compute_logical_maps(phy2log, weight.shape[-1]) |
| 285 | |
| 286 | # Function will convert to CPU internally, but should handle different |
| 287 | # device inputs normally |
| 288 | assert phy2log.shape == (1, 6) |
| 289 | assert logcnt.shape == (1, 4) |
| 290 | |
| 291 | |
| 292 | def test_additional_cases(): |
| 293 | """Test more edge cases and different parameter combinations.""" |
| 294 | # Test case 1: Large-scale distributed setup |
| 295 | weight1 = torch.tensor( |
| 296 | [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]] |
| 297 | ) |
| 298 | phy2log1 = DefaultEplbPolicy.rebalance_experts(weight1, 24, 8, 4, 8) |
| 299 | _, logcnt1 = compute_logical_maps(phy2log1, weight1.shape[-1]) |
| 300 | |
| 301 | assert phy2log1.shape == (1, 24) |
| 302 | assert logcnt1.shape == (1, 16) |
| 303 | assert torch.sum(logcnt1) == 24 |
| 304 | |
| 305 | # Test case 2: Different weight distributions |
| 306 | weight2 = torch.tensor( |
| 307 | [ |
| 308 | [200, 150, 100, 50, 25, 12], # Decreasing weights |
| 309 | [12, 25, 50, 100, 150, 200], # Increasing weights |
| 310 | ] |
| 311 | ) |
| 312 | phy2log2 = DefaultEplbPolicy.rebalance_experts(weight2, 10, 3, 1, 2) |
| 313 | _, logcnt2 = compute_logical_maps(phy2log2, weight2.shape[-1]) |
| 314 | |
| 315 | assert phy2log2.shape == (2, 10) |
| 316 | assert logcnt2.shape == (2, 6) |
| 317 | |
| 318 | # Verify high-weight experts have more replicas |
| 319 | for layer in range(2): |
| 320 | max_weight_idx = torch.argmax(weight2[layer]) |
| 321 | assert logcnt2[layer, max_weight_idx] >= 2 |
| 322 | |
| 323 | |
| 324 | def test_compute_logical_maps_with_negative_indices(): |
| 325 | """Test that compute_logical_maps correctly handles physical slots containing |
| 326 | -1 (unused slots). |
| 327 | """ |
| 328 | # 2 layers, 6 physical slots, 4 logical experts. |
| 329 | # Slots 2 and 5 are unused (-1). |
| 330 | phy2log = torch.tensor( |
| 331 | [ |
| 332 | [0, 1, -1, 2, 3, -1], |
| 333 | [3, -1, 2, 1, 0, -1], |
| 334 | ] |
| 335 | ) |
| 336 | num_layers = 2 |
| 337 | num_logical_experts = 4 |
| 338 | |
| 339 | log2phy, logcnt = compute_logical_maps(phy2log, num_logical_experts) |
| 340 | |
| 341 | assert logcnt.shape == (num_layers, num_logical_experts) |
| 342 | assert log2phy.shape == (num_layers, num_logical_experts, 1) |
| 343 | |
| 344 | expected_logcnt = torch.ones(num_layers, num_logical_experts, dtype=phy2log.dtype) |
| 345 | assert torch.all(logcnt == expected_logcnt), ( |
| 346 | f"Expected that all replica counts == 1, got {logcnt}" |
| 347 | ) |
| 348 | |
| 349 | assert torch.all(log2phy >= 0), ( |
| 350 | "log2phy should only contain valid physical indices, not -1" |
| 351 | ) |
| 352 | |
| 353 | assert log2phy[0, 0, 0] == 0 |
| 354 | assert log2phy[0, 1, 0] == 1 |
| 355 | assert log2phy[0, 2, 0] == 3 |
| 356 | assert log2phy[0, 3, 0] == 4 |
| 357 | |
| 358 | |
| 359 | if __name__ == "__main__": |
| 360 | weight = torch.tensor( |
| 361 | [ |
| 362 | [90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86], |
| 363 | [20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27], |
| 364 | ] |
| 365 | ) |
| 366 | |
| 367 | num_replicas = 16 |
| 368 | num_groups = 4 |
| 369 | num_nodes = 2 |
| 370 | num_gpus = 8 |
| 371 | |
| 372 | phy2log = DefaultEplbPolicy.rebalance_experts( |
| 373 | weight, num_replicas, num_groups, num_nodes, num_gpus |
| 374 | ) |
| 375 | print(phy2log) |
| 376 | |
| 377 | test_basic_rebalance() |
| 378 | |
| 379 | |
| 380 | def _make_phy_replicas_idx_from_phy2log(phy2log: np.ndarray) -> np.ndarray: |
| 381 | """Create replicas indices mapping from phy2log.""" |
| 382 | pr = np.zeros_like(phy2log, dtype=np.int64) |
| 383 | for layer in range(phy2log.shape[0]): |
| 384 | seen: dict[int, int] = {} |
| 385 | row = phy2log[layer].tolist() |
| 386 | for i, expert in enumerate(row): |
| 387 | r = seen.get(expert, 0) |
| 388 | pr[layer, i] = r |
| 389 | seen[expert] = r + 1 |
| 390 | return pr |
| 391 | |
| 392 | |
| 393 | def _validate_intragpu_rearrangement( |
| 394 | old_global_expert_indices: np.ndarray, |
| 395 | new_phy2log: np.ndarray, |
| 396 | new_phy_replicas_idx: np.ndarray, |
| 397 | post_phy2log: np.ndarray, |
| 398 | post_phy_replicas_idx: np.ndarray, |
| 399 | num_ranks: int, |
| 400 | slots_per_gpu: int, |
| 401 | ): |
| 402 | # Per-GPU checks |
| 403 | for gpu_idx in range(num_ranks): |
| 404 | start = gpu_idx * slots_per_gpu |
| 405 | end = start + slots_per_gpu |
| 406 | old_seg = old_global_expert_indices[0, start:end] |
| 407 | new_seg = new_phy2log[0, start:end] |
| 408 | new_rnk = new_phy_replicas_idx[0, start:end] |
| 409 | post_seg = post_phy2log[0, start:end] |
| 410 | post_rnk = post_phy_replicas_idx[0, start:end] |
| 411 | |
| 412 | # Pairwise equality for (expert, rank) pairs to ensure nothing is lost |
| 413 | def sorted_pairs(seg, rnk): |
| 414 | pairs = list(zip(seg.tolist(), rnk.tolist())) |
| 415 | pairs.sort() |
| 416 | return pairs |
| 417 | |
| 418 | assert sorted_pairs(post_seg, post_rnk) == sorted_pairs(new_seg, new_rnk), ( |
| 419 | f"Per-GPU pairs of (expert,rank) must match new mapping for GPU {gpu_idx}" |
| 420 | ) |
| 421 | |
| 422 | # For experts that remain on the same GPU, the old slot is preserved |
| 423 | # for at least one occurrence; rank at that slot must be valid for that expert |
| 424 | old_list = old_seg.tolist() |
| 425 | new_list = new_seg.tolist() |
| 426 | post_list = post_seg.tolist() |
| 427 | remained = set(old_list) & set(new_list) |
| 428 | new_ranks_for_expert: dict[int, list[int]] = {} |
| 429 | for v, r in zip(new_list, new_rnk.tolist()): |
| 430 | new_ranks_for_expert.setdefault(v, []).append(r) |
| 431 | for expert in remained: |
| 432 | old_pos = old_list.index(expert) |
| 433 | assert post_list[old_pos] == expert, ( |
| 434 | f"Expert {expert} on GPU {gpu_idx} should stay at old slot {old_pos}" |
| 435 | ) |
| 436 | # Rank at preserved slot must be one of the ranks |
| 437 | # the expert has in new mapping |
| 438 | assert post_rnk.tolist()[old_pos] in new_ranks_for_expert[expert], ( |
| 439 | f"Rank for expert {expert} at preserved slot on GPU {gpu_idx} " |
| 440 | "must come from new mapping" |
| 441 | ) |
| 442 | |
| 443 | |
| 444 | @pytest.mark.parametrize( |
| 445 | "num_ranks, slots_per_gpu, old_phy2log, new_phy2log", |
| 446 | [ |
| 447 | pytest.param( |
| 448 | # Setup: 2 GPUs, 4 slots each, 1 layer |
| 449 | # Old mapping: GPU0 -> [0,1,2,3], GPU1 -> [4,5,6,7] |
| 450 | # New mapping shuffles within GPU0 and brings 4,5 into GPU0. |
| 451 | # GPU0 new -> [1,5,0,4]; GPU1 new -> [6,2,7,3] |
| 452 | 2, |
| 453 | 4, |
| 454 | np.array([[0, 1, 2, 3, 4, 5, 6, 7]]), |
| 455 | np.array([[1, 5, 0, 4, 6, 2, 7, 3]]), |
| 456 | id="simple", |
| 457 | ), |
| 458 | pytest.param( |
| 459 | # Setup: 2 GPUs, 5 slots each (total 10 physical experts), 1 layer |
| 460 | # Old mapping: |
| 461 | # GPU0 -> [0, 1, 0, 2, 3] (expert 0 duplicated) |
| 462 | # GPU1 -> [4, 5, 6, 1, 2] |
| 463 | # New mapping reorders within GPUs and moves some experts across GPUs, |
| 464 | # while still including duplicates: |
| 465 | # GPU0 new -> [0, 5, 4, 0, 1] (expert 0 duplicated, 4/5 incoming) |
| 466 | # GPU1 new -> [6, 2, 3, 2, 1] (expert 2 duplicated) |
| 467 | 2, |
| 468 | 5, |
| 469 | np.array([[0, 1, 0, 2, 3, 4, 5, 6, 1, 2]]), |
| 470 | np.array([[0, 5, 4, 0, 1, 6, 2, 3, 2, 1]]), |
| 471 | id="duplicates", |
| 472 | ), |
| 473 | pytest.param( |
| 474 | # Setup: 3 GPUs, 4 slots each (total 12 physical experts), 1 layer |
| 475 | # Old mapping: |
| 476 | # GPU0 -> [0, 1, 2, 3] |
| 477 | # GPU1 -> [0, 1, 2, 3] |
| 478 | # GPU2 -> [0, 1, 2, 3] |
| 479 | # New mapping decides to use one expert on 2 GPUs and shuffles |
| 480 | # experts on the third GPU, |
| 481 | # GPU0 new -> [0, 0, 0, 0] |
| 482 | # GPU1 new -> [0, 0, 0, 0] |
| 483 | # GPU2 new -> [1, 2, 3, 0] |
| 484 | 3, |
| 485 | 4, |
| 486 | np.array([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]), |
| 487 | np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 0]]), |
| 488 | id="skewed_expert", |
| 489 | ), |
| 490 | ], |
| 491 | ) |
| 492 | def test_preserve_intragpu_slots( |
| 493 | num_ranks: int, |
| 494 | slots_per_gpu: int, |
| 495 | old_phy2log: torch.Tensor, |
| 496 | new_phy2log: torch.Tensor, |
| 497 | ): |
| 498 | """Experts that stay on a GPU keep their old slots; incoming not lost.""" |
| 499 | phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(new_phy2log) |
| 500 | |
| 501 | post_phy2log = DefaultEplbPolicy.preserve_intragpu_slots( |
| 502 | new_phy2log, num_ranks, old_phy2log |
| 503 | ) |
| 504 | post_phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(post_phy2log) |
| 505 | |
| 506 | # Shapes preserved |
| 507 | assert post_phy2log.shape == new_phy2log.shape |
| 508 | assert post_phy_replicas_idx.shape == phy_replicas_idx.shape |
| 509 | |
| 510 | _validate_intragpu_rearrangement( |
| 511 | old_phy2log, |
| 512 | new_phy2log, |
| 513 | phy_replicas_idx, |
| 514 | post_phy2log, |
| 515 | post_phy_replicas_idx, |
| 516 | num_ranks, |
| 517 | slots_per_gpu, |
| 518 | ) |