| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | import torch.distributed |
| 9 | |
| 10 | import vllm.distributed.eplb.eplb_communicator as eplb_comm |
| 11 | import vllm.utils.gpu_sync_debug as gsd |
| 12 | from vllm.config import VllmConfig, set_current_vllm_config |
| 13 | from vllm.distributed.eplb.eplb_communicator import ( |
| 14 | create_eplb_communicator, |
| 15 | has_nixl, |
| 16 | ) |
| 17 | from vllm.distributed.eplb.rebalance_execute import ( |
| 18 | move_from_buffer, |
| 19 | rearrange_expert_weights_inplace, |
| 20 | transfer_layer, |
| 21 | ) |
| 22 | from vllm.distributed.parallel_state import ( |
| 23 | ensure_model_parallel_initialized, |
| 24 | get_tp_group, |
| 25 | ) |
| 26 | from vllm.platforms import current_platform |
| 27 | |
| 28 | from .eplb_utils import distributed_run, set_env_vars_and_device |
| 29 | |
| 30 | |
| 31 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 32 | def test_gloo_receive_staging_does_not_force_gpu_sync(monkeypatch): |
| 33 | """Gloo receives must reach the GPU without an implicit pageable copy.""" |
| 34 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") |
| 35 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 36 | gsd._install_copy_checkers() |
| 37 | monkeypatch.setattr(eplb_comm, "is_local_first_rank", lambda: False) |
| 38 | |
| 39 | monkeypatch.setattr(eplb_comm, "P2POp", lambda op, tensor, peer, group: tensor) |
| 40 | |
| 41 | def receive(tensors): |
| 42 | for tensor in tensors: |
| 43 | tensor.fill_(7) |
| 44 | return [] |
| 45 | |
| 46 | monkeypatch.setattr(eplb_comm, "batch_isend_irecv", receive) |
| 47 | communicator = eplb_comm.TorchDistGlooStagedEplbCommunicator(cpu_group=None) |
| 48 | dst = torch.empty(32, device="cuda") |
| 49 | communicator.add_recv([dst], src_rank=1, expert_id=0) |
| 50 | gsd.with_gpu_sync_check(communicator.execute)() |
| 51 | torch.testing.assert_close(dst.cpu(), torch.full((32,), 7.0)) |
| 52 | |
| 53 | |
| 54 | def create_expert_indices_with_redundancy( |
| 55 | num_layers: int, |
| 56 | num_logical_experts: int, |
| 57 | total_physical_experts: int, |
| 58 | redundancy_config: list[int], # redundancy for each logical expert |
| 59 | ) -> torch.Tensor: |
| 60 | """Create expert indices with redundancy. |
| 61 | |
| 62 | Args: |
| 63 | num_layers: number of layers |
| 64 | num_logical_experts: number of logical experts |
| 65 | total_physical_experts: total number of physical experts |
| 66 | redundancy_config: redundancy for each logical expert |
| 67 | |
| 68 | Returns: |
| 69 | indices: Shape (num_layers, total_physical_experts) |
| 70 | |
| 71 | """ |
| 72 | assert sum(redundancy_config) == total_physical_experts |
| 73 | assert len(redundancy_config) == num_logical_experts |
| 74 | |
| 75 | indices = torch.zeros(num_layers, total_physical_experts, dtype=torch.long) |
| 76 | |
| 77 | for layer in range(num_layers): |
| 78 | physical_pos = 0 |
| 79 | for logical_expert_id, redundancy in enumerate(redundancy_config): |
| 80 | for _ in range(redundancy): |
| 81 | indices[layer, physical_pos] = logical_expert_id |
| 82 | physical_pos += 1 |
| 83 | |
| 84 | # Shuffle the indices at dim 1 |
| 85 | for layer in range(num_layers): |
| 86 | indices[layer] = indices[layer][torch.randperm(indices.shape[1])] |
| 87 | |
| 88 | return indices |
| 89 | |
| 90 | |
| 91 | def create_expert_weights( |
| 92 | num_layers: int, |
| 93 | num_local_experts: int, |
| 94 | hidden_sizes: list[int], |
| 95 | rank: int, |
| 96 | device: torch.device, |
| 97 | physical_to_logical_mapping: torch.Tensor, |
| 98 | ) -> list[list[torch.Tensor]]: |
| 99 | """Create fake expert weights tensor for testing. |
| 100 | |
| 101 | Use `arange` to generate predictable weights values, based on logical |
| 102 | expert ID. |
| 103 | All replicas of the same logical expert should have the same weights. |
| 104 | |
| 105 | Args: |
| 106 | physical_to_logical_mapping: Shape (num_layers, num_local_experts) |
| 107 | mapping[layer, physical_pos] = logical_expert_id |
| 108 | |
| 109 | """ |
| 110 | expert_weights = [] |
| 111 | |
| 112 | for layer in range(num_layers): |
| 113 | layer_weights = [] |
| 114 | for weight_idx, hidden_size in enumerate(hidden_sizes): |
| 115 | weight_tensor = torch.zeros( |
| 116 | num_local_experts, hidden_size, device=device, dtype=torch.float32 |
| 117 | ) |
| 118 | |
| 119 | for local_expert in range(num_local_experts): |
| 120 | # Get the logical expert ID for this physical expert |
| 121 | global_pos = rank * num_local_experts + local_expert |
| 122 | logical_expert_id = physical_to_logical_mapping[ |
| 123 | layer, global_pos |
| 124 | ].item() |
| 125 | |
| 126 | # Generate weights based on logical expert ID |
| 127 | # (so that all replicas of the same logical expert have the |
| 128 | # same weights) |
| 129 | base_value = logical_expert_id * 1000 + layer * 100 + weight_idx * 10 |
| 130 | weight_tensor[local_expert] = torch.arange( |
| 131 | base_value, |
| 132 | base_value + hidden_size, |
| 133 | device=device, |
| 134 | dtype=torch.float32, |
| 135 | ) |
| 136 | |
| 137 | layer_weights.append(weight_tensor) |
| 138 | expert_weights.append(layer_weights) |
| 139 | |
| 140 | return expert_weights |
| 141 | |
| 142 | |
| 143 | def create_redundancy_config( |
| 144 | num_logical_experts: int, |
| 145 | num_physical_experts: int, |
| 146 | ) -> list[int]: |
| 147 | """Create a redundancy configuration.""" |
| 148 | redundancy_config = [1] * num_logical_experts |
| 149 | remaining = num_physical_experts - num_logical_experts |
| 150 | # Randomly assign the remaining physical experts to the logical experts |
| 151 | for _ in range(remaining): |
| 152 | redundancy_config[random.choice(range(num_logical_experts))] += 1 |
| 153 | return redundancy_config |
| 154 | |
| 155 | |
| 156 | def verify_expert_weights_after_shuffle( |
| 157 | expert_weights: list[list[torch.Tensor]], |
| 158 | new_indices: torch.Tensor, |
| 159 | hidden_sizes: list[int], |
| 160 | ep_rank: int, |
| 161 | num_local_experts: int, |
| 162 | ) -> bool: |
| 163 | """Verify the weights after shuffling are correct.""" |
| 164 | num_layers = len(expert_weights) |
| 165 | ok = True |
| 166 | |
| 167 | for layer in range(num_layers): |
| 168 | for weight_idx, hidden_size in enumerate(hidden_sizes): |
| 169 | weight_tensor = expert_weights[layer][weight_idx] |
| 170 | |
| 171 | for local_expert in range(num_local_experts): |
| 172 | # Calculate the global expert ID for this local expert |
| 173 | global_pos = ep_rank * num_local_experts + local_expert |
| 174 | expected_logical_expert = new_indices[layer, global_pos].item() |
| 175 | |
| 176 | # Check if the weights are correct |
| 177 | actual_weights = weight_tensor[local_expert] |
| 178 | expected_base = ( |
| 179 | expected_logical_expert * 1000 + layer * 100 + weight_idx * 10 |
| 180 | ) |
| 181 | expected_weights = torch.arange( |
| 182 | expected_base, |
| 183 | expected_base + hidden_size, |
| 184 | device=actual_weights.device, |
| 185 | dtype=actual_weights.dtype, |
| 186 | ) |
| 187 | |
| 188 | if not torch.equal(actual_weights, expected_weights): |
| 189 | ok = False |
| 190 | actual_head = actual_weights[:8].detach().cpu().tolist() |
| 191 | expected_head = expected_weights[:8].detach().cpu().tolist() |
| 192 | print( |
| 193 | "verify_expert_weights_after_shuffle failed: " |
| 194 | f"rank={ep_rank}, " |
| 195 | f"layer={layer}, weight_idx={weight_idx}, " |
| 196 | f"local_expert={local_expert}, " |
| 197 | f"expected_logical_expert={expected_logical_expert}, " |
| 198 | f"actual_head={actual_head}, expected_head={expected_head}", |
| 199 | flush=True, |
| 200 | ) |
| 201 | |
| 202 | return ok |
| 203 | |
| 204 | |
| 205 | def verify_redundant_experts_have_same_weights( |
| 206 | expert_weights: list[list[torch.Tensor]], |
| 207 | indices: torch.Tensor, |
| 208 | hidden_sizes: list[int], |
| 209 | ep_rank: int, |
| 210 | world_size: int, |
| 211 | num_local_experts: int, |
| 212 | cpu_group, |
| 213 | ) -> bool: |
| 214 | """Verify that all replicas of the same logical expert have the same weights.""" |
| 215 | num_layers = len(expert_weights) |
| 216 | total_physical_experts = world_size * num_local_experts |
| 217 | # XPU: gloo does not support XPU tensors, so gather on CPU then move back. |
| 218 | # CUDA: gloo supports CUDA tensors directly, keep original device. |
| 219 | use_cpu_gather = current_platform.is_xpu() |
| 220 | |
| 221 | ok = True |
| 222 | for layer in range(num_layers): |
| 223 | # Collect weights for all physical experts for each weight matrix |
| 224 | all_weights: list[torch.Tensor] = [] |
| 225 | |
| 226 | for weight_idx, hidden_size in enumerate(hidden_sizes): |
| 227 | orig_device = expert_weights[layer][weight_idx].device |
| 228 | |
| 229 | if use_cpu_gather: |
| 230 | local_weights = expert_weights[layer][weight_idx].cpu() |
| 231 | gather_device = "cpu" |
| 232 | else: |
| 233 | local_weights = expert_weights[layer][weight_idx] |
| 234 | gather_device = orig_device |
| 235 | |
| 236 | # Create tensor to store all expert weights |
| 237 | # Shape: [total_physical_experts, hidden_size] |
| 238 | gathered_weights = torch.zeros( |
| 239 | total_physical_experts, |
| 240 | hidden_size, |
| 241 | device=gather_device, |
| 242 | dtype=local_weights.dtype, |
| 243 | ) |
| 244 | |
| 245 | # Split tensor along dim 0 into a list for all_gather |
| 246 | gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0) |
| 247 | |
| 248 | torch.distributed.all_gather( |
| 249 | list(gathered_weights_list), |
| 250 | local_weights, |
| 251 | group=cpu_group, |
| 252 | ) |
| 253 | |
| 254 | if use_cpu_gather: |
| 255 | gathered_weights = gathered_weights.to(orig_device) |
| 256 | all_weights.append(gathered_weights) |
| 257 | |
| 258 | # Verify that all replicas of the same logical expert have the same |
| 259 | # weights |
| 260 | logical_expert_weights: dict[int, dict[int, torch.Tensor]] = {} |
| 261 | |
| 262 | for physical_pos in range(total_physical_experts): |
| 263 | logical_expert_id = int(indices[layer, physical_pos].item()) |
| 264 | |
| 265 | if logical_expert_id not in logical_expert_weights: |
| 266 | # First time encountering this logical expert, save its weights |
| 267 | logical_expert_weights[logical_expert_id] = { |
| 268 | weight_idx: all_weights[weight_idx][physical_pos] |
| 269 | for weight_idx in range(len(hidden_sizes)) |
| 270 | } |
| 271 | else: |
| 272 | # Verify that current physical expert's weights match the |
| 273 | # previously saved logical expert weights |
| 274 | for weight_idx in range(len(hidden_sizes)): |
| 275 | if not torch.equal( |
| 276 | all_weights[weight_idx][physical_pos], |
| 277 | logical_expert_weights[logical_expert_id][weight_idx], |
| 278 | ): |
| 279 | ok = False |
| 280 | actual_head = ( |
| 281 | all_weights[weight_idx][physical_pos][:8] |
| 282 | .detach() |
| 283 | .cpu() |
| 284 | .tolist() |
| 285 | ) |
| 286 | reference_head = ( |
| 287 | logical_expert_weights[logical_expert_id][weight_idx][:8] |
| 288 | .detach() |
| 289 | .cpu() |
| 290 | .tolist() |
| 291 | ) |
| 292 | print( |
| 293 | "verify_redundant_experts_have_same_weights failed: " |
| 294 | f"rank={ep_rank}, " |
| 295 | f"layer={layer}, weight_idx={weight_idx}, " |
| 296 | f"logical_expert={logical_expert_id}, " |
| 297 | f"physical_pos={physical_pos}, " |
| 298 | f"actual_head={actual_head}, " |
| 299 | f"reference_head={reference_head}", |
| 300 | flush=True, |
| 301 | ) |
| 302 | |
| 303 | return ok |
| 304 | |
| 305 | |
| 306 | def assert_verification_synced(local_ok: bool, msg: str, cpu_group) -> None: |
| 307 | if current_platform.is_xpu(): |
| 308 | # XPU: default group backend (gloo) does not support XPU tensors; |
| 309 | # use CPU tensor with cpu_group instead. |
| 310 | ok_tensor = torch.tensor( |
| 311 | [1 if local_ok else 0], device="cpu", dtype=torch.int32 |
| 312 | ) |
| 313 | torch.distributed.all_reduce( |
| 314 | ok_tensor, op=torch.distributed.ReduceOp.MIN, group=cpu_group |
| 315 | ) |
| 316 | else: |
| 317 | # CUDA: nccl (default group) supports device tensors directly. |
| 318 | device = torch.accelerator.current_accelerator() |
| 319 | ok_tensor = torch.tensor( |
| 320 | [1 if local_ok else 0], device=device, dtype=torch.int32 |
| 321 | ) |
| 322 | torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN) |
| 323 | assert bool(ok_tensor.item()), msg |
| 324 | |
| 325 | |
| 326 | def create_eplb_communicator_or_raise( |
| 327 | *, group_coordinator, backend, expert_weights, expert_buffer |
| 328 | ): |
| 329 | try: |
| 330 | return create_eplb_communicator( |
| 331 | group_coordinator=group_coordinator, |
| 332 | backend=backend, |
| 333 | expert_weights=expert_weights, |
| 334 | expert_buffer=expert_buffer, |
| 335 | ) |
| 336 | except Exception as exc: |
| 337 | raise RuntimeError( |
| 338 | f"Failed to create EPLB communicator for backend={backend}: {exc}" |
| 339 | ) from exc |
| 340 | |
| 341 | |
| 342 | def _test_async_transfer_layer_without_mtp_worker( |
| 343 | env, |
| 344 | world_size: int, |
| 345 | num_layers: int, |
| 346 | num_local_experts: int, |
| 347 | num_logical_experts: int, |
| 348 | eplb_communicator: str, |
| 349 | ) -> None: |
| 350 | set_env_vars_and_device(env) |
| 351 | |
| 352 | vllm_config = VllmConfig() |
| 353 | vllm_config.parallel_config.tensor_parallel_size = world_size |
| 354 | |
| 355 | with set_current_vllm_config(vllm_config): |
| 356 | ensure_model_parallel_initialized( |
| 357 | tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 |
| 358 | ) |
| 359 | |
| 360 | ep_group_coordinator = get_tp_group() |
| 361 | ep_group = ep_group_coordinator.device_group |
| 362 | cpu_group = ep_group_coordinator.cpu_group |
| 363 | ep_rank = torch.distributed.get_rank() |
| 364 | accelerator_type = torch.accelerator.current_accelerator().type |
| 365 | device = torch.device(f"{accelerator_type}:{ep_rank}") |
| 366 | |
| 367 | total_physical_experts = world_size * num_local_experts |
| 368 | hidden_sizes = [16, 32] |
| 369 | |
| 370 | redundancy_config = create_redundancy_config( |
| 371 | num_logical_experts, |
| 372 | total_physical_experts, |
| 373 | ) |
| 374 | old_indices = create_expert_indices_with_redundancy( |
| 375 | num_layers, |
| 376 | num_logical_experts, |
| 377 | total_physical_experts, |
| 378 | redundancy_config, |
| 379 | ) |
| 380 | |
| 381 | new_redundancy_config = create_redundancy_config( |
| 382 | num_logical_experts, |
| 383 | total_physical_experts, |
| 384 | ) |
| 385 | new_indices = create_expert_indices_with_redundancy( |
| 386 | num_layers, |
| 387 | num_logical_experts, |
| 388 | total_physical_experts, |
| 389 | new_redundancy_config, |
| 390 | ) |
| 391 | |
| 392 | expert_weights = create_expert_weights( |
| 393 | num_layers, |
| 394 | num_local_experts, |
| 395 | hidden_sizes, |
| 396 | ep_rank, |
| 397 | device, |
| 398 | old_indices, |
| 399 | ) |
| 400 | old_indices_cpu = old_indices.cpu() |
| 401 | new_indices_cpu = new_indices.cpu() |
| 402 | |
| 403 | expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] |
| 404 | stream = torch.Stream(device=device) |
| 405 | |
| 406 | communicator = create_eplb_communicator_or_raise( |
| 407 | group_coordinator=ep_group_coordinator, |
| 408 | backend=eplb_communicator, |
| 409 | expert_weights=expert_weights, |
| 410 | expert_buffer=expert_buffer, |
| 411 | ) |
| 412 | communicator.set_stream(stream) |
| 413 | |
| 414 | for layer_idx in range(num_layers): |
| 415 | transfer_metadata = transfer_layer( |
| 416 | old_layer_indices=old_indices_cpu[layer_idx], |
| 417 | new_layer_indices=new_indices_cpu[layer_idx], |
| 418 | expert_weights=expert_weights[layer_idx], |
| 419 | expert_weights_buffer=expert_buffer, |
| 420 | ep_group=ep_group, |
| 421 | communicator=communicator, |
| 422 | stream=stream, |
| 423 | layer_idx=layer_idx, |
| 424 | ) |
| 425 | stream.synchronize() |
| 426 | move_from_buffer( |
| 427 | expert_weights=expert_weights[layer_idx], |
| 428 | expert_weights_buffers=expert_buffer, |
| 429 | transfer_metadata=transfer_metadata, |
| 430 | new_indices=new_indices_cpu[layer_idx].numpy(), |
| 431 | ep_rank=ep_rank, |
| 432 | ) |
| 433 | |
| 434 | local_ok = verify_expert_weights_after_shuffle( |
| 435 | expert_weights, |
| 436 | new_indices, |
| 437 | hidden_sizes, |
| 438 | ep_rank, |
| 439 | num_local_experts, |
| 440 | ) |
| 441 | local_ok = ( |
| 442 | verify_redundant_experts_have_same_weights( |
| 443 | expert_weights, |
| 444 | new_indices, |
| 445 | hidden_sizes, |
| 446 | ep_rank, |
| 447 | world_size, |
| 448 | num_local_experts, |
| 449 | cpu_group=cpu_group, |
| 450 | ) |
| 451 | and local_ok |
| 452 | ) |
| 453 | assert_verification_synced( |
| 454 | local_ok, |
| 455 | "Async transfer verification failed on at least one rank. " |
| 456 | "See logs for details.", |
| 457 | cpu_group=cpu_group, |
| 458 | ) |
| 459 | |
| 460 | |
| 461 | def _test_rearrange_expert_weights_with_redundancy( |
| 462 | env, |
| 463 | world_size, |
| 464 | num_layers, |
| 465 | num_local_experts, |
| 466 | num_logical_experts, |
| 467 | eplb_communicator: str, |
| 468 | ) -> None: |
| 469 | # Initialize model parallel (using tensor parallel as an entrypoint |
| 470 | # to expert parallel) |
| 471 | set_env_vars_and_device(env) |
| 472 | |
| 473 | vllm_config = VllmConfig() |
| 474 | vllm_config.parallel_config.tensor_parallel_size = world_size |
| 475 | |
| 476 | with set_current_vllm_config(vllm_config): |
| 477 | ensure_model_parallel_initialized( |
| 478 | tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 |
| 479 | ) |
| 480 | |
| 481 | ep_group_coordinator = get_tp_group() |
| 482 | ep_group = ep_group_coordinator.device_group |
| 483 | cpu_group = ep_group_coordinator.cpu_group |
| 484 | ep_rank = torch.distributed.get_rank() |
| 485 | accelerator_type = torch.accelerator.current_accelerator().type |
| 486 | device = torch.device(f"{accelerator_type}:{ep_rank}") |
| 487 | |
| 488 | # Test parameters |
| 489 | total_physical_experts = world_size * num_local_experts |
| 490 | hidden_sizes = [32, 64] # Two different weight matrices |
| 491 | |
| 492 | # Create old expert indices (with redundancy) |
| 493 | redundancy_config = create_redundancy_config( |
| 494 | num_logical_experts, total_physical_experts |
| 495 | ) |
| 496 | |
| 497 | old_indices = create_expert_indices_with_redundancy( |
| 498 | num_layers, |
| 499 | num_logical_experts, |
| 500 | total_physical_experts, |
| 501 | redundancy_config, |
| 502 | ) |
| 503 | |
| 504 | # Create new expert indices (with redundancy) |
| 505 | new_redundancy_config = create_redundancy_config( |
| 506 | num_logical_experts, total_physical_experts |
| 507 | ) |
| 508 | new_indices = create_expert_indices_with_redundancy( |
| 509 | num_layers, |
| 510 | num_logical_experts, |
| 511 | total_physical_experts, |
| 512 | new_redundancy_config, |
| 513 | ) |
| 514 | |
| 515 | # Create expert weights |
| 516 | expert_weights = create_expert_weights( |
| 517 | num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices |
| 518 | ) |
| 519 | |
| 520 | expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] |
| 521 | communicator = create_eplb_communicator_or_raise( |
| 522 | group_coordinator=ep_group_coordinator, |
| 523 | backend=eplb_communicator, |
| 524 | expert_weights=expert_weights, |
| 525 | expert_buffer=expert_buffer, |
| 526 | ) |
| 527 | |
| 528 | # Execute weight rearrangement |
| 529 | rearrange_expert_weights_inplace( |
| 530 | old_indices, |
| 531 | new_indices, |
| 532 | expert_weights, |
| 533 | expert_buffer, |
| 534 | ep_group, |
| 535 | communicator, |
| 536 | ) |
| 537 | |
| 538 | # Verify the rearrangement result |
| 539 | local_ok = verify_expert_weights_after_shuffle( |
| 540 | expert_weights, |
| 541 | new_indices, |
| 542 | hidden_sizes, |
| 543 | ep_rank, |
| 544 | num_local_experts, |
| 545 | ) |
| 546 | |
| 547 | local_ok = ( |
| 548 | verify_redundant_experts_have_same_weights( |
| 549 | expert_weights, |
| 550 | new_indices, |
| 551 | hidden_sizes, |
| 552 | ep_rank, |
| 553 | world_size, |
| 554 | num_local_experts, |
| 555 | cpu_group=cpu_group, |
| 556 | ) |
| 557 | and local_ok |
| 558 | ) |
| 559 | assert_verification_synced( |
| 560 | local_ok, |
| 561 | "Rearrange verification failed on at least one rank. See logs for details.", |
| 562 | cpu_group=cpu_group, |
| 563 | ) |
| 564 | |
| 565 | |
| 566 | @pytest.mark.parametrize( |
| 567 | "world_size,num_layers,num_local_experts,num_logical_experts", |
| 568 | [ |
| 569 | # 2 GPU, 2 experts per GPU |
| 570 | # 3 logical experts, 4 physical experts, 1 redundant experts |
| 571 | (2, 1, 2, 3), |
| 572 | # 2 GPU, 3 experts per GPU |
| 573 | # 4 logical experts, 6 physical experts, 2 redundant experts |
| 574 | (2, 2, 3, 4), |
| 575 | # 2 GPU, 8 experts per GPU |
| 576 | # 16 logical experts, 16 physical experts, 0 redundant experts |
| 577 | (2, 4, 8, 16), |
| 578 | # 4 GPU, 2 experts per GPU |
| 579 | # 6 logical experts, 8 physical experts, 2 redundant experts |
| 580 | (4, 1, 2, 6), |
| 581 | # 4 GPU, 2 experts per GPU |
| 582 | # 5 logical experts, 8 physical experts, 3 redundant experts |
| 583 | (4, 2, 2, 5), |
| 584 | # 4 GPU, 8 experts per GPU |
| 585 | # 16 logical experts, 32 physical experts, 16 redundant experts |
| 586 | (4, 8, 8, 16), |
| 587 | ], |
| 588 | ) |
| 589 | @pytest.mark.parametrize( |
| 590 | "eplb_communicator", ["torch_nccl", "torch_gloo", "torch_xccl", "pynccl", "nixl"] |
| 591 | ) |
| 592 | def test_rearrange_expert_weights_with_redundancy( |
| 593 | world_size, |
| 594 | num_layers, |
| 595 | num_local_experts, |
| 596 | num_logical_experts, |
| 597 | eplb_communicator, |
| 598 | ): |
| 599 | """Test the functionality of rearranging expert weights with redundancy.""" |
| 600 | if eplb_communicator == "nixl" and not has_nixl(): |
| 601 | pytest.skip("NIXL is not available") |
| 602 | if eplb_communicator == "nixl" and current_platform.is_xpu(): |
| 603 | pytest.skip("NIXL does not support XPU") |
| 604 | if eplb_communicator in ("torch_nccl", "pynccl") and not torch.cuda.is_available(): |
| 605 | pytest.skip(f"{eplb_communicator} requires CUDA") |
| 606 | if eplb_communicator == "torch_xccl" and not current_platform.is_xpu(): |
| 607 | pytest.skip("torch_xccl requires XPU") |
| 608 | if torch.accelerator.device_count() < world_size: |
| 609 | pytest.skip(f"Need at least {world_size} GPUs to run the test") |
| 610 | distributed_run( |
| 611 | _test_rearrange_expert_weights_with_redundancy, |
| 612 | world_size, |
| 613 | num_layers, |
| 614 | num_local_experts, |
| 615 | num_logical_experts, |
| 616 | eplb_communicator, |
| 617 | ) |
| 618 | |
| 619 | |
| 620 | def _test_rearrange_expert_weights_no_change(env, world_size) -> None: |
| 621 | set_env_vars_and_device(env) |
| 622 | |
| 623 | vllm_config = VllmConfig() |
| 624 | vllm_config.parallel_config.tensor_parallel_size = world_size |
| 625 | |
| 626 | with set_current_vllm_config(vllm_config): |
| 627 | ensure_model_parallel_initialized( |
| 628 | tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 |
| 629 | ) |
| 630 | |
| 631 | ep_group_coordinator = get_tp_group() |
| 632 | ep_group = ep_group_coordinator.device_group |
| 633 | cpu_group = ep_group_coordinator.cpu_group |
| 634 | ep_rank = torch.distributed.get_rank() |
| 635 | accelerator_type = torch.accelerator.current_accelerator().type |
| 636 | device = torch.device(f"{accelerator_type}:{ep_rank}") |
| 637 | |
| 638 | num_layers = 2 |
| 639 | num_local_experts = 2 |
| 640 | total_physical_experts = world_size * num_local_experts |
| 641 | num_logical_experts = total_physical_experts // 2 # Some redundancy |
| 642 | hidden_sizes = [32, 64] |
| 643 | |
| 644 | # Create redundancy configuration |
| 645 | redundancy_config = [2] * num_logical_experts |
| 646 | |
| 647 | # Same indices - no change |
| 648 | indices = create_expert_indices_with_redundancy( |
| 649 | num_layers, num_logical_experts, total_physical_experts, redundancy_config |
| 650 | ) |
| 651 | |
| 652 | expert_weights = create_expert_weights( |
| 653 | num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices |
| 654 | ) |
| 655 | |
| 656 | # Save original weights |
| 657 | original_weights = [] |
| 658 | for layer_weights in expert_weights: |
| 659 | layer_copy = [] |
| 660 | for weight in layer_weights: |
| 661 | layer_copy.append(weight.clone()) |
| 662 | original_weights.append(layer_copy) |
| 663 | |
| 664 | expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] |
| 665 | default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" |
| 666 | communicator = create_eplb_communicator_or_raise( |
| 667 | group_coordinator=ep_group_coordinator, |
| 668 | backend=default_backend, |
| 669 | expert_weights=expert_weights, |
| 670 | expert_buffer=expert_buffer, |
| 671 | ) |
| 672 | |
| 673 | # Execute rearrangement (should be no change) |
| 674 | rearrange_expert_weights_inplace( |
| 675 | indices, |
| 676 | indices, # Same indices |
| 677 | expert_weights, |
| 678 | expert_buffer, |
| 679 | ep_group, |
| 680 | communicator, |
| 681 | ) |
| 682 | |
| 683 | # Verify that the weights have not changed |
| 684 | local_ok = True |
| 685 | for layer in range(num_layers): |
| 686 | for weight_idx in range(len(hidden_sizes)): |
| 687 | if not torch.equal( |
| 688 | expert_weights[layer][weight_idx], |
| 689 | original_weights[layer][weight_idx], |
| 690 | ): |
| 691 | local_ok = False |
| 692 | print( |
| 693 | "test_rearrange_expert_weights_no_change failed: " |
| 694 | f"layer={layer}, weight_idx={weight_idx}", |
| 695 | flush=True, |
| 696 | ) |
| 697 | assert_verification_synced( |
| 698 | local_ok, |
| 699 | "No-change EPLB verification failed on at least one rank.", |
| 700 | cpu_group=cpu_group, |
| 701 | ) |
| 702 | |
| 703 | |
| 704 | @pytest.mark.parametrize( |
| 705 | "world_size,num_layers,num_local_experts,num_logical_experts", |
| 706 | [ |
| 707 | (2, 2, 2, 3), |
| 708 | ], |
| 709 | ) |
| 710 | @pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "torch_xccl", "nixl"]) |
| 711 | def test_async_transfer_layer_without_mtp( |
| 712 | world_size: int, |
| 713 | num_layers: int, |
| 714 | num_local_experts: int, |
| 715 | num_logical_experts: int, |
| 716 | eplb_communicator: str, |
| 717 | ): |
| 718 | """Exercise async EPLB transfer path without MTP/spec decode.""" |
| 719 | if eplb_communicator == "nixl" and not has_nixl(): |
| 720 | pytest.skip("NIXL is not available") |
| 721 | if eplb_communicator == "nixl" and current_platform.is_xpu(): |
| 722 | pytest.skip("NIXL does not support XPU") |
| 723 | if eplb_communicator == "torch_xccl" and not current_platform.is_xpu(): |
| 724 | pytest.skip("torch_xccl requires XPU") |
| 725 | if torch.accelerator.device_count() < world_size: |
| 726 | pytest.skip(f"Need at least {world_size} GPUs to run the test") |
| 727 | |
| 728 | distributed_run( |
| 729 | _test_async_transfer_layer_without_mtp_worker, |
| 730 | world_size, |
| 731 | num_layers, |
| 732 | num_local_experts, |
| 733 | num_logical_experts, |
| 734 | eplb_communicator, |
| 735 | ) |
| 736 | |
| 737 | |
| 738 | @pytest.mark.parametrize("world_size", [2, 4]) |
| 739 | def test_rearrange_expert_weights_no_change(world_size): |
| 740 | """Test that when the indices do not change, the weights should remain |
| 741 | unchanged. |
| 742 | """ |
| 743 | if torch.accelerator.device_count() < world_size: |
| 744 | pytest.skip(f"Need at least {world_size} GPUs to run the test") |
| 745 | distributed_run( |
| 746 | _test_rearrange_expert_weights_no_change, |
| 747 | world_size, |
| 748 | ) |
| 749 | |
| 750 | |
| 751 | def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: |
| 752 | set_env_vars_and_device(env) |
| 753 | |
| 754 | vllm_config = VllmConfig() |
| 755 | vllm_config.parallel_config.tensor_parallel_size = world_size |
| 756 | |
| 757 | with set_current_vllm_config(vllm_config): |
| 758 | ensure_model_parallel_initialized( |
| 759 | tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 |
| 760 | ) |
| 761 | |
| 762 | ep_group_coordinator = get_tp_group() |
| 763 | ep_group = ep_group_coordinator.device_group |
| 764 | cpu_group = ep_group_coordinator.cpu_group |
| 765 | ep_rank = torch.distributed.get_rank() |
| 766 | accelerator_type = torch.accelerator.current_accelerator().type |
| 767 | device = torch.device(f"{accelerator_type}:{ep_rank}") |
| 768 | |
| 769 | num_layers = 1 |
| 770 | num_local_experts = 2 |
| 771 | total_physical_experts = world_size * num_local_experts |
| 772 | num_logical_experts = total_physical_experts // 2 |
| 773 | hidden_sizes = [32] |
| 774 | |
| 775 | # Create different index distributions |
| 776 | old_redundancy = create_redundancy_config( |
| 777 | num_logical_experts, total_physical_experts |
| 778 | ) |
| 779 | new_redundancy = create_redundancy_config( |
| 780 | num_logical_experts, total_physical_experts |
| 781 | ) |
| 782 | |
| 783 | old_indices = create_expert_indices_with_redundancy( |
| 784 | num_layers, num_logical_experts, total_physical_experts, old_redundancy |
| 785 | ) |
| 786 | new_indices = create_expert_indices_with_redundancy( |
| 787 | num_layers, num_logical_experts, total_physical_experts, new_redundancy |
| 788 | ) |
| 789 | |
| 790 | expert_weights = create_expert_weights( |
| 791 | num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices |
| 792 | ) |
| 793 | |
| 794 | # Save original weights |
| 795 | original_weights = [] |
| 796 | for layer_weights in expert_weights: |
| 797 | layer_copy = [] |
| 798 | for weight in layer_weights: |
| 799 | layer_copy.append(weight.clone()) |
| 800 | original_weights.append(layer_copy) |
| 801 | |
| 802 | expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] |
| 803 | default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" |
| 804 | communicator = create_eplb_communicator_or_raise( |
| 805 | group_coordinator=ep_group_coordinator, |
| 806 | backend=default_backend, |
| 807 | expert_weights=expert_weights, |
| 808 | expert_buffer=expert_buffer, |
| 809 | ) |
| 810 | |
| 811 | # Execute profile mode rearrangement |
| 812 | rearrange_expert_weights_inplace( |
| 813 | old_indices, |
| 814 | new_indices, |
| 815 | expert_weights, |
| 816 | expert_buffer, |
| 817 | ep_group, |
| 818 | communicator, |
| 819 | is_profile=True, |
| 820 | ) |
| 821 | |
| 822 | # In profile mode, the weights should remain unchanged |
| 823 | local_ok = True |
| 824 | for layer in range(num_layers): |
| 825 | for weight_idx in range(len(hidden_sizes)): |
| 826 | if not torch.equal( |
| 827 | expert_weights[layer][weight_idx], |
| 828 | original_weights[layer][weight_idx], |
| 829 | ): |
| 830 | local_ok = False |
| 831 | print( |
| 832 | "test_rearrange_expert_weights_profile_mode failed: " |
| 833 | f"layer={layer}, weight_idx={weight_idx}", |
| 834 | flush=True, |
| 835 | ) |
| 836 | assert_verification_synced( |
| 837 | local_ok, |
| 838 | "Profile-mode EPLB verification failed on at least one rank.", |
| 839 | cpu_group=cpu_group, |
| 840 | ) |
| 841 | |
| 842 | |
| 843 | @pytest.mark.parametrize("world_size", [2, 4]) |
| 844 | def test_rearrange_expert_weights_profile_mode(world_size): |
| 845 | """Test profile mode (should not copy actual weights)""" |
| 846 | if torch.accelerator.device_count() < world_size: |
| 847 | pytest.skip(f"Need at least {world_size} GPUs to run the test") |
| 848 | distributed_run( |
| 849 | _test_rearrange_expert_weights_profile_mode, |
| 850 | world_size, |
| 851 | ) |