A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 22:15 UTC
[ci][amd] fix EPLB execution test (#28742)
Signed-off-by: Bradley Davis <bradleyhd@meta.com>
Bradley D committed
Nov 19, 2025 at 23:53 UTC
1e1c06789e63a760d91aaf6e4ddfeabfe382c301
1 file changed
+213
-210
tests/distributed/test_eplb_execute.py
+213
-210
@@ -1,13 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import multiprocessing
4
import os
5
import random
6
7
import pytest
8
import torch
9
import torch.distributed
10
+import torch.multiprocessing as mp
11
12
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
13
from vllm.distributed.parallel_state import (
@@ -17,10 +17,12 @@ from vllm.distributed.parallel_state import (
17
)
18
from vllm.utils.system_utils import update_environment_variables
19
20
+mp.set_start_method("spawn", force=True)
21
21
-def distributed_run(fn, world_size):
22
+
23
+def distributed_run(fn, world_size, *args):
24
number_of_processes = world_size
23
- processes: list[multiprocessing.Process] = []
25
+ processes: list[mp.Process] = []
26
for i in range(number_of_processes):
27
env: dict[str, str] = {}
28
env["RANK"] = str(i)
@@ -29,7 +31,7 @@ def distributed_run(fn, world_size):
31
env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
32
env["MASTER_ADDR"] = "localhost"
33
env["MASTER_PORT"] = "12345"
32
- p = multiprocessing.Process(target=fn, args=(env,))
34
+ p = mp.Process(target=fn, args=(env, world_size, *args))
35
processes.append(p)
36
p.start()
37
@@ -40,24 +42,16 @@ def distributed_run(fn, world_size):
42
assert p.exitcode == 0
43
44
43
-def worker_fn_wrapper(fn):
44
- # `multiprocessing.Process` cannot accept environment variables directly
45
- # so we need to pass the environment variables as arguments
46
- # and update the environment variables in the function
47
- def wrapped_fn(env):
48
- update_environment_variables(env)
49
- local_rank = os.environ["LOCAL_RANK"]
50
- device = torch.device(f"cuda:{local_rank}")
51
- torch.cuda.set_device(device)
52
- init_distributed_environment()
53
-
54
- # Ensure each worker process has the same random seed
55
- random.seed(42)
56
- torch.manual_seed(42)
57
-
58
- fn()
45
+def set_env_vars_and_device(env: dict[str, str]) -> None:
46
+ update_environment_variables(env)
47
+ local_rank = os.environ["LOCAL_RANK"]
48
+ device = torch.device(f"cuda:{local_rank}")
49
+ torch.cuda.set_device(device)
50
+ init_distributed_environment()
51
60
- return wrapped_fn
52
+ # Ensure each worker process has the same random seed
53
+ random.seed(42)
54
+ torch.manual_seed(42)
55
56
57
def create_expert_indices_with_redundancy(
@@ -275,6 +269,79 @@ def verify_redundant_experts_have_same_weights(
269
)
270
271
272
+def _test_rearrange_expert_weights_with_redundancy(
273
+ env, world_size, num_layers, num_local_experts, num_logical_experts
274
+) -> None:
275
+ # Initialize model parallel (using tensor parallel as an entrypoint
276
+ # to expert parallel)
277
+ set_env_vars_and_device(env)
278
+ ensure_model_parallel_initialized(
279
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
280
+ )
281
+
282
+ ep_group = get_tp_group().cpu_group
283
+ ep_rank = torch.distributed.get_rank()
284
+ device = torch.device(f"cuda:{ep_rank}")
285
+
286
+ # Test parameters
287
+ total_physical_experts = world_size * num_local_experts
288
+ hidden_sizes = [32, 64] # Two different weight matrices
289
+
290
+ # Create old expert indices (with redundancy)
291
+ redundancy_config = create_redundancy_config(
292
+ num_logical_experts, total_physical_experts
293
+ )
294
+
295
+ old_indices = create_expert_indices_with_redundancy(
296
+ num_layers,
297
+ num_logical_experts,
298
+ total_physical_experts,
299
+ redundancy_config,
300
+ )
301
+
302
+ # Create new expert indices (with redundancy)
303
+ new_redundancy_config = create_redundancy_config(
304
+ num_logical_experts, total_physical_experts
305
+ )
306
+ new_indices = create_expert_indices_with_redundancy(
307
+ num_layers,
308
+ num_logical_experts,
309
+ total_physical_experts,
310
+ new_redundancy_config,
311
+ )
312
+
313
+ # Create expert weights
314
+ expert_weights = create_expert_weights(
315
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
316
+ )
317
+
318
+ # Execute weight rearrangement
319
+ rearrange_expert_weights_inplace(
320
+ old_indices,
321
+ new_indices,
322
+ expert_weights,
323
+ ep_group,
324
+ is_profile=False,
325
+ )
326
+
327
+ # Verify the rearrangement result
328
+ verify_expert_weights_after_shuffle(
329
+ expert_weights,
330
+ new_indices,
331
+ hidden_sizes,
332
+ ep_rank,
333
+ num_local_experts,
334
+ )
335
+
336
+ verify_redundant_experts_have_same_weights(
337
+ expert_weights,
338
+ new_indices,
339
+ hidden_sizes,
340
+ world_size,
341
+ num_local_experts,
342
+ )
343
+
344
+
345
@pytest.mark.parametrize(
346
"world_size,num_layers,num_local_experts,num_logical_experts",
347
[
@@ -305,78 +372,69 @@ def test_rearrange_expert_weights_with_redundancy(
372
373
if torch.cuda.device_count() < world_size:
374
pytest.skip(f"Need at least {world_size} GPUs to run the test")
308
-
309
- @worker_fn_wrapper
310
- def worker_fn():
311
- # Initialize model parallel (using tensor parallel as an entrypoint
312
- # to expert parallel)
313
- ensure_model_parallel_initialized(
314
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
315
- )
316
-
317
- ep_group = get_tp_group().cpu_group
318
- ep_rank = torch.distributed.get_rank()
319
- device = torch.device(f"cuda:{ep_rank}")
320
-
321
- # Test parameters
322
- total_physical_experts = world_size * num_local_experts
323
- hidden_sizes = [32, 64] # Two different weight matrices
324
-
325
- # Create old expert indices (with redundancy)
326
- redundancy_config = create_redundancy_config(
327
- num_logical_experts, total_physical_experts
328
- )
329
-
330
- old_indices = create_expert_indices_with_redundancy(
331
- num_layers,
332
- num_logical_experts,
333
- total_physical_experts,
334
- redundancy_config,
335
- )
336
-
337
- # Create new expert indices (with redundancy)
338
- new_redundancy_config = create_redundancy_config(
339
- num_logical_experts, total_physical_experts
340
- )
341
- new_indices = create_expert_indices_with_redundancy(
342
- num_layers,
343
- num_logical_experts,
344
- total_physical_experts,
345
- new_redundancy_config,
346
- )
347
-
348
- # Create expert weights
349
- expert_weights = create_expert_weights(
350
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
351
- )
352
-
353
- # Execute weight rearrangement
354
- rearrange_expert_weights_inplace(
355
- old_indices,
356
- new_indices,
357
- expert_weights,
358
- ep_group,
359
- is_profile=False,
360
- )
361
-
362
- # Verify the rearrangement result
363
- verify_expert_weights_after_shuffle(
364
- expert_weights,
365
- new_indices,
366
- hidden_sizes,
367
- ep_rank,
368
- num_local_experts,
369
- )
370
-
371
- verify_redundant_experts_have_same_weights(
372
- expert_weights,
373
- new_indices,
374
- hidden_sizes,
375
- world_size,
376
- num_local_experts,
377
- )
378
-
379
- distributed_run(worker_fn, world_size)
375
+ distributed_run(
376
+ _test_rearrange_expert_weights_with_redundancy,
377
+ world_size,
378
+ num_layers,
379
+ num_local_experts,
380
+ num_logical_experts,
381
+ )
382
+
383
+
384
+def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
385
+ set_env_vars_and_device(env)
386
+ ensure_model_parallel_initialized(
387
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
388
+ )
389
+
390
+ ep_group = get_tp_group().cpu_group
391
+ ep_rank = torch.distributed.get_rank()
392
+ device = torch.device(f"cuda:{ep_rank}")
393
+
394
+ num_layers = 2
395
+ num_local_experts = 2
396
+ total_physical_experts = world_size * num_local_experts
397
+ num_logical_experts = total_physical_experts // 2 # Some redundancy
398
+ hidden_sizes = [32, 64]
399
+
400
+ # Create redundancy configuration
401
+ redundancy_config = [2] * num_logical_experts
402
+
403
+ # Same indices - no change
404
+ indices = create_expert_indices_with_redundancy(
405
+ num_layers, num_logical_experts, total_physical_experts, redundancy_config
406
+ )
407
+
408
+ expert_weights = create_expert_weights(
409
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
410
+ )
411
+
412
+ # Save original weights
413
+ original_weights = []
414
+ for layer_weights in expert_weights:
415
+ layer_copy = []
416
+ for weight in layer_weights:
417
+ layer_copy.append(weight.clone())
418
+ original_weights.append(layer_copy)
419
+
420
+ # Execute rearrangement (should be no change)
421
+ rearrange_expert_weights_inplace(
422
+ indices,
423
+ indices, # Same indices
424
+ expert_weights,
425
+ ep_group,
426
+ is_profile=False,
427
+ )
428
+
429
+ # Verify that the weights have not changed
430
+ for layer in range(num_layers):
431
+ for weight_idx in range(len(hidden_sizes)):
432
+ torch.testing.assert_close(
433
+ expert_weights[layer][weight_idx],
434
+ original_weights[layer][weight_idx],
435
+ msg=f"""Layer {layer}, weight {weight_idx}
436
+ should remain unchanged""",
437
+ )
438
439
440
@pytest.mark.parametrize("world_size", [2, 4])
@@ -388,62 +446,69 @@ def test_rearrange_expert_weights_no_change(world_size):
446
447
if torch.cuda.device_count() < world_size:
448
pytest.skip(f"Need at least {world_size} GPUs to run the test")
449
+ distributed_run(_test_rearrange_expert_weights_no_change, world_size)
450
392
- @worker_fn_wrapper
393
- def worker_fn():
394
- ensure_model_parallel_initialized(
395
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
396
- )
397
-
398
- ep_group = get_tp_group().cpu_group
399
- ep_rank = torch.distributed.get_rank()
400
- device = torch.device(f"cuda:{ep_rank}")
401
-
402
- num_layers = 2
403
- num_local_experts = 2
404
- total_physical_experts = world_size * num_local_experts
405
- num_logical_experts = total_physical_experts // 2 # Some redundancy
406
- hidden_sizes = [32, 64]
407
-
408
- # Create redundancy configuration
409
- redundancy_config = [2] * num_logical_experts
410
-
411
- # Same indices - no change
412
- indices = create_expert_indices_with_redundancy(
413
- num_layers, num_logical_experts, total_physical_experts, redundancy_config
414
- )
415
-
416
- expert_weights = create_expert_weights(
417
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
418
- )
419
-
420
- # Save original weights
421
- original_weights = []
422
- for layer_weights in expert_weights:
423
- layer_copy = []
424
- for weight in layer_weights:
425
- layer_copy.append(weight.clone())
426
- original_weights.append(layer_copy)
427
-
428
- # Execute rearrangement (should be no change)
429
- rearrange_expert_weights_inplace(
430
- indices,
431
- indices, # Same indices
432
- expert_weights,
433
- ep_group,
434
- is_profile=False,
435
- )
436
-
437
- # Verify that the weights have not changed
438
- for layer in range(num_layers):
439
- for weight_idx in range(len(hidden_sizes)):
440
- torch.testing.assert_close(
441
- expert_weights[layer][weight_idx],
442
- original_weights[layer][weight_idx],
443
- msg=f"Layer {layer}, weight {weight_idx} should remain unchanged",
444
- )
451
446
- distributed_run(worker_fn, world_size)
452
+def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
453
+ set_env_vars_and_device(env)
454
+ ensure_model_parallel_initialized(
455
+ tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
456
+ )
457
+
458
+ ep_group = get_tp_group().cpu_group
459
+ ep_rank = torch.distributed.get_rank()
460
+ device = torch.device(f"cuda:{ep_rank}")
461
+
462
+ num_layers = 1
463
+ num_local_experts = 2
464
+ total_physical_experts = world_size * num_local_experts
465
+ num_logical_experts = total_physical_experts // 2
466
+ hidden_sizes = [32]
467
+
468
+ # Create different index distributions
469
+ old_redundancy = create_redundancy_config(
470
+ num_logical_experts, total_physical_experts
471
+ )
472
+ new_redundancy = create_redundancy_config(
473
+ num_logical_experts, total_physical_experts
474
+ )
475
+
476
+ old_indices = create_expert_indices_with_redundancy(
477
+ num_layers, num_logical_experts, total_physical_experts, old_redundancy
478
+ )
479
+ new_indices = create_expert_indices_with_redundancy(
480
+ num_layers, num_logical_experts, total_physical_experts, new_redundancy
481
+ )
482
+
483
+ expert_weights = create_expert_weights(
484
+ num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
485
+ )
486
+
487
+ # Save original weights
488
+ original_weights = []
489
+ for layer_weights in expert_weights:
490
+ layer_copy = []
491
+ for weight in layer_weights:
492
+ layer_copy.append(weight.clone())
493
+ original_weights.append(layer_copy)
494
+
495
+ # Execute profile mode rearrangement
496
+ rearrange_expert_weights_inplace(
497
+ old_indices,
498
+ new_indices,
499
+ expert_weights,
500
+ ep_group,
501
+ is_profile=True, # Profile mode
502
+ )
503
+
504
+ # In profile mode, the weights should remain unchanged
505
+ for layer in range(num_layers):
506
+ for weight_idx in range(len(hidden_sizes)):
507
+ torch.testing.assert_close(
508
+ expert_weights[layer][weight_idx],
509
+ original_weights[layer][weight_idx],
510
+ msg="In profile mode, the weights should remain unchanged",
511
+ )
512
513
514
@pytest.mark.parametrize("world_size", [2, 4])
@@ -452,66 +517,4 @@ def test_rearrange_expert_weights_profile_mode(world_size):
517
518
if torch.cuda.device_count() < world_size:
519
pytest.skip(f"Need at least {world_size} GPUs to run the test")
455
-
456
- @worker_fn_wrapper
457
- def worker_fn():
458
- ensure_model_parallel_initialized(
459
- tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
460
- )
461
-
462
- ep_group = get_tp_group().cpu_group
463
- ep_rank = torch.distributed.get_rank()
464
- device = torch.device(f"cuda:{ep_rank}")
465
-
466
- num_layers = 1
467
- num_local_experts = 2
468
- total_physical_experts = world_size * num_local_experts
469
- num_logical_experts = total_physical_experts // 2
470
- hidden_sizes = [32]
471
-
472
- # Create different index distributions
473
- old_redundancy = create_redundancy_config(
474
- num_logical_experts, total_physical_experts
475
- )
476
- new_redundancy = create_redundancy_config(
477
- num_logical_experts, total_physical_experts
478
- )
479
-
480
- old_indices = create_expert_indices_with_redundancy(
481
- num_layers, num_logical_experts, total_physical_experts, old_redundancy
482
- )
483
- new_indices = create_expert_indices_with_redundancy(
484
- num_layers, num_logical_experts, total_physical_experts, new_redundancy
485
- )
486
-
487
- expert_weights = create_expert_weights(
488
- num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
489
- )
490
-
491
- # Save original weights
492
- original_weights = []
493
- for layer_weights in expert_weights:
494
- layer_copy = []
495
- for weight in layer_weights:
496
- layer_copy.append(weight.clone())
497
- original_weights.append(layer_copy)
498
-
499
- # Execute profile mode rearrangement
500
- rearrange_expert_weights_inplace(
501
- old_indices,
502
- new_indices,
503
- expert_weights,
504
- ep_group,
505
- is_profile=True, # Profile mode
506
- )
507
-
508
- # In profile mode, the weights should remain unchanged
509
- for layer in range(num_layers):
510
- for weight_idx in range(len(hidden_sizes)):
511
- torch.testing.assert_close(
512
- expert_weights[layer][weight_idx],
513
- original_weights[layer][weight_idx],
514
- msg="In profile mode, the weights should remain unchanged",
515
- )
516
-
517
- distributed_run(worker_fn, world_size)
520
+ distributed_run(_test_rearrange_expert_weights_profile_mode, world_size)