test_eplb_execute.py 25.8 KB
Newer Older
1
2
3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4
import asyncio
5
6
7
8
9
10
import random

import pytest
import torch
import torch.distributed

11
from vllm.config import VllmConfig, set_current_vllm_config
12
13
14
15
from vllm.distributed.eplb.eplb_communicator import (
    create_eplb_communicator,
    has_nixl,
)
16
17
18
19
20
from vllm.distributed.eplb.rebalance_execute import (
    move_from_buffer,
    rearrange_expert_weights_inplace,
    transfer_layer,
)
21
22
23
24
from vllm.distributed.parallel_state import (
    ensure_model_parallel_initialized,
    get_tp_group,
)
25

26
from .eplb_utils import distributed_run, set_env_vars_and_device
27
28
29


def create_expert_indices_with_redundancy(
30
31
32
33
    num_layers: int,
    num_logical_experts: int,
    total_physical_experts: int,
    redundancy_config: list[int],  # redundancy for each logical expert
34
35
36
) -> torch.Tensor:
    """
    Create expert indices with redundancy.
37

38
39
40
41
42
    Args:
        num_layers: number of layers
        num_logical_experts: number of logical experts
        total_physical_experts: total number of physical experts
        redundancy_config: redundancy for each logical expert
43

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    Returns:
        indices: Shape (num_layers, total_physical_experts)
    """
    assert sum(redundancy_config) == total_physical_experts
    assert len(redundancy_config) == num_logical_experts

    indices = torch.zeros(num_layers, total_physical_experts, dtype=torch.long)

    for layer in range(num_layers):
        physical_pos = 0
        for logical_expert_id, redundancy in enumerate(redundancy_config):
            for _ in range(redundancy):
                indices[layer, physical_pos] = logical_expert_id
                physical_pos += 1

    # Shuffle the indices at dim 1
    for layer in range(num_layers):
        indices[layer] = indices[layer][torch.randperm(indices.shape[1])]

    return indices


def create_expert_weights(
    num_layers: int,
    num_local_experts: int,
    hidden_sizes: list[int],
    rank: int,
    device: torch.device,
    physical_to_logical_mapping: torch.Tensor,
) -> list[list[torch.Tensor]]:
    """
    Create fake expert weights tensor for testing.
76

77
78
79
    Use `arange` to generate predictable weights values, based on logical
    expert ID.
    All replicas of the same logical expert should have the same weights.
80

81
82
83
84
85
86
87
88
89
    Args:
        physical_to_logical_mapping: Shape (num_layers, num_local_experts)
            mapping[layer, physical_pos] = logical_expert_id
    """
    expert_weights = []

    for layer in range(num_layers):
        layer_weights = []
        for weight_idx, hidden_size in enumerate(hidden_sizes):
90
91
92
            weight_tensor = torch.zeros(
                num_local_experts, hidden_size, device=device, dtype=torch.float32
            )
93
94
95
96
97

            for local_expert in range(num_local_experts):
                # Get the logical expert ID for this physical expert
                global_pos = rank * num_local_experts + local_expert
                logical_expert_id = physical_to_logical_mapping[
98
99
                    layer, global_pos
                ].item()
100
101
102
103

                # Generate weights based on logical expert ID
                # (so that all replicas of the same logical expert have the
                # same weights)
104
105
106
107
108
109
110
                base_value = logical_expert_id * 1000 + layer * 100 + weight_idx * 10
                weight_tensor[local_expert] = torch.arange(
                    base_value,
                    base_value + hidden_size,
                    device=device,
                    dtype=torch.float32,
                )
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

            layer_weights.append(weight_tensor)
        expert_weights.append(layer_weights)

    return expert_weights


def create_redundancy_config(
    num_logical_experts: int,
    num_physical_experts: int,
) -> list[int]:
    """Create a redundancy configuration."""
    redundancy_config = [1] * num_logical_experts
    remaining = num_physical_experts - num_logical_experts
    # Randomly assign the remaining physical experts to the logical experts
    for _ in range(remaining):
        redundancy_config[random.choice(range(num_logical_experts))] += 1
    return redundancy_config


def verify_expert_weights_after_shuffle(
    expert_weights: list[list[torch.Tensor]],
    new_indices: torch.Tensor,
    hidden_sizes: list[int],
    ep_rank: int,
    num_local_experts: int,
137
) -> bool:
138
139
    """Verify the weights after shuffling are correct."""
    num_layers = len(expert_weights)
140
    ok = True
141
142
143
144
145
146
147
148
149
150
151
152

    for layer in range(num_layers):
        for weight_idx, hidden_size in enumerate(hidden_sizes):
            weight_tensor = expert_weights[layer][weight_idx]

            for local_expert in range(num_local_experts):
                # Calculate the global expert ID for this local expert
                global_pos = ep_rank * num_local_experts + local_expert
                expected_logical_expert = new_indices[layer, global_pos].item()

                # Check if the weights are correct
                actual_weights = weight_tensor[local_expert]
153
154
155
156
157
158
159
160
161
                expected_base = (
                    expected_logical_expert * 1000 + layer * 100 + weight_idx * 10
                )
                expected_weights = torch.arange(
                    expected_base,
                    expected_base + hidden_size,
                    device=actual_weights.device,
                    dtype=actual_weights.dtype,
                )
162

163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
                if not torch.equal(actual_weights, expected_weights):
                    ok = False
                    actual_head = actual_weights[:8].detach().cpu().tolist()
                    expected_head = expected_weights[:8].detach().cpu().tolist()
                    print(
                        "verify_expert_weights_after_shuffle failed: "
                        f"rank={ep_rank}, "
                        f"layer={layer}, weight_idx={weight_idx}, "
                        f"local_expert={local_expert}, "
                        f"expected_logical_expert={expected_logical_expert}, "
                        f"actual_head={actual_head}, expected_head={expected_head}",
                        flush=True,
                    )

    return ok
178
179
180
181
182
183


def verify_redundant_experts_have_same_weights(
    expert_weights: list[list[torch.Tensor]],
    indices: torch.Tensor,
    hidden_sizes: list[int],
184
    ep_rank: int,
185
186
    world_size: int,
    num_local_experts: int,
187
) -> bool:
188
189
190
191
192
193
    """
    Verify that all replicas of the same logical expert have the same weights.
    """
    num_layers = len(expert_weights)
    total_physical_experts = world_size * num_local_experts

194
    ok = True
195
196
197
198
199
200
201
202
203
204
205
    for layer in range(num_layers):
        # Collect weights for all physical experts for each weight matrix
        all_weights: list[torch.Tensor] = []

        for weight_idx, hidden_size in enumerate(hidden_sizes):
            # Create tensor to store all expert weights
            # Shape: [total_physical_experts, hidden_size]
            gathered_weights = torch.zeros(
                total_physical_experts,
                hidden_size,
                device=expert_weights[layer][weight_idx].device,
206
207
                dtype=expert_weights[layer][weight_idx].dtype,
            )
208
209
210
211
212

            # Use all_gather to collect expert weights from current node
            # expert_weights[layer][weight_idx] shape:
            # [num_local_experts, hidden_size]
            local_weights = expert_weights[layer][
213
214
                weight_idx
            ]  # [num_local_experts, hidden_size]
215
216

            # Split tensor along dim 0 into a list for all_gather
217
            gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0)
218
219
220
221

            torch.distributed.all_gather(
                # Output list: each element corresponds to one rank's weights
                list(gathered_weights_list),
222
                local_weights,  # Input: current rank's local weights
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
            )

            all_weights.append(gathered_weights)

        # Verify that all replicas of the same logical expert have the same
        # weights
        logical_expert_weights: dict[int, dict[int, torch.Tensor]] = {}

        for physical_pos in range(total_physical_experts):
            logical_expert_id = int(indices[layer, physical_pos].item())

            if logical_expert_id not in logical_expert_weights:
                # First time encountering this logical expert, save its weights
                logical_expert_weights[logical_expert_id] = {
                    weight_idx: all_weights[weight_idx][physical_pos]
                    for weight_idx in range(len(hidden_sizes))
                }
            else:
                # Verify that current physical expert's weights match the
                # previously saved logical expert weights
                for weight_idx in range(len(hidden_sizes)):
244
                    if not torch.equal(
245
246
                        all_weights[weight_idx][physical_pos],
                        logical_expert_weights[logical_expert_id][weight_idx],
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
                    ):
                        ok = False
                        actual_head = (
                            all_weights[weight_idx][physical_pos][:8]
                            .detach()
                            .cpu()
                            .tolist()
                        )
                        reference_head = (
                            logical_expert_weights[logical_expert_id][weight_idx][:8]
                            .detach()
                            .cpu()
                            .tolist()
                        )
                        print(
                            "verify_redundant_experts_have_same_weights failed: "
                            f"rank={ep_rank}, "
                            f"layer={layer}, weight_idx={weight_idx}, "
                            f"logical_expert={logical_expert_id}, "
                            f"physical_pos={physical_pos}, "
                            f"actual_head={actual_head}, "
                            f"reference_head={reference_head}",
                            flush=True,
                        )

    return ok


def assert_verification_synced(local_ok: bool, msg: str) -> None:
    ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32)
    torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)
    assert bool(ok_tensor.item()), msg


def create_eplb_communicator_or_raise(*, group_coordinator, backend, expert_weights):
    try:
        return create_eplb_communicator(
            group_coordinator=group_coordinator,
            backend=backend,
            expert_weights=expert_weights,
        )
    except Exception as exc:
        raise RuntimeError(
            f"Failed to create EPLB communicator for backend={backend}: {exc}"
        ) from exc
292
293


294
295
296
297
298
299
def _test_async_transfer_layer_without_mtp_worker(
    env,
    world_size: int,
    num_layers: int,
    num_local_experts: int,
    num_logical_experts: int,
300
    eplb_communicator: str,
301
302
303
) -> None:
    set_env_vars_and_device(env)

304
305
    vllm_config = VllmConfig()
    vllm_config.parallel_config.tensor_parallel_size = world_size
306

307
308
309
310
    with set_current_vllm_config(vllm_config):
        ensure_model_parallel_initialized(
            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
        )
311

312
313
        ep_group_coordinator = get_tp_group()
        ep_group = ep_group_coordinator.device_group
314
315
        ep_rank = torch.distributed.get_rank()
        device = torch.device(f"cuda:{ep_rank}")
316

317
318
        total_physical_experts = world_size * num_local_experts
        hidden_sizes = [16, 32]
319

320
321
322
323
324
325
326
327
328
329
        redundancy_config = create_redundancy_config(
            num_logical_experts,
            total_physical_experts,
        )
        old_indices = create_expert_indices_with_redundancy(
            num_layers,
            num_logical_experts,
            total_physical_experts,
            redundancy_config,
        )
330

331
332
333
334
335
336
337
338
339
340
        new_redundancy_config = create_redundancy_config(
            num_logical_experts,
            total_physical_experts,
        )
        new_indices = create_expert_indices_with_redundancy(
            num_layers,
            num_logical_experts,
            total_physical_experts,
            new_redundancy_config,
        )
341

342
343
344
345
346
347
348
349
350
351
352
353
354
355
        expert_weights = create_expert_weights(
            num_layers,
            num_local_experts,
            hidden_sizes,
            ep_rank,
            device,
            old_indices,
        )
        old_indices_cpu = old_indices.cpu()
        new_indices_cpu = new_indices.cpu()

        expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
        cuda_stream = torch.cuda.Stream(device=device)

356
357
358
359
360
361
362
        communicator = create_eplb_communicator_or_raise(
            group_coordinator=ep_group_coordinator,
            backend=eplb_communicator,
            expert_weights=expert_weights[0],
        )
        communicator.set_stream(cuda_stream)

363
        for layer_idx in range(num_layers):
364
            transfer_metadata = asyncio.run(
365
366
367
368
369
370
                transfer_layer(
                    old_layer_indices=old_indices_cpu[layer_idx],
                    new_layer_indices=new_indices_cpu[layer_idx],
                    expert_weights=expert_weights[layer_idx],
                    expert_weights_buffer=expert_buffer,
                    ep_group=ep_group,
371
                    communicator=communicator,
372
373
374
375
376
                    cuda_stream=cuda_stream,
                )
            )
            cuda_stream.synchronize()
            move_from_buffer(
377
                expert_weights=expert_weights[layer_idx],
378
                expert_weights_buffers=expert_buffer,
379
                transfer_metadata=transfer_metadata,
380
381
                new_indices=new_indices_cpu[layer_idx].numpy(),
                ep_rank=ep_rank,
382
            )
383

384
385
386
387
388
389
390
391
    local_ok = verify_expert_weights_after_shuffle(
        expert_weights,
        new_indices,
        hidden_sizes,
        ep_rank,
        num_local_experts,
    )
    local_ok = (
392
393
394
395
        verify_redundant_experts_have_same_weights(
            expert_weights,
            new_indices,
            hidden_sizes,
396
            ep_rank,
397
398
            world_size,
            num_local_experts,
399
        )
400
401
402
403
404
405
406
        and local_ok
    )
    assert_verification_synced(
        local_ok,
        "Async transfer verification failed on at least one rank. "
        "See logs for details.",
    )
407
408


409
def _test_rearrange_expert_weights_with_redundancy(
410
411
412
413
414
415
    env,
    world_size,
    num_layers,
    num_local_experts,
    num_logical_experts,
    eplb_communicator: str,
416
417
418
419
420
) -> None:
    # Initialize model parallel (using tensor parallel as an entrypoint
    # to expert parallel)
    set_env_vars_and_device(env)

421
422
    vllm_config = VllmConfig()
    vllm_config.parallel_config.tensor_parallel_size = world_size
423

424
425
426
427
    with set_current_vllm_config(vllm_config):
        ensure_model_parallel_initialized(
            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
        )
428

429
430
        ep_group_coordinator = get_tp_group()
        ep_group = ep_group_coordinator.cpu_group
431
432
        ep_rank = torch.distributed.get_rank()
        device = torch.device(f"cuda:{ep_rank}")
433

434
435
436
        # Test parameters
        total_physical_experts = world_size * num_local_experts
        hidden_sizes = [32, 64]  # Two different weight matrices
437

438
439
440
441
        # Create old expert indices (with redundancy)
        redundancy_config = create_redundancy_config(
            num_logical_experts, total_physical_experts
        )
442

443
444
445
446
447
448
        old_indices = create_expert_indices_with_redundancy(
            num_layers,
            num_logical_experts,
            total_physical_experts,
            redundancy_config,
        )
449

450
451
452
453
454
455
456
457
458
459
        # Create new expert indices (with redundancy)
        new_redundancy_config = create_redundancy_config(
            num_logical_experts, total_physical_experts
        )
        new_indices = create_expert_indices_with_redundancy(
            num_layers,
            num_logical_experts,
            total_physical_experts,
            new_redundancy_config,
        )
460

461
462
463
464
        # Create expert weights
        expert_weights = create_expert_weights(
            num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
        )
465

466
467
468
469
470
471
        communicator = create_eplb_communicator_or_raise(
            group_coordinator=ep_group_coordinator,
            backend=eplb_communicator,
            expert_weights=expert_weights[0],
        )

472
473
474
475
476
477
478
        # Execute weight rearrangement
        rearrange_expert_weights_inplace(
            old_indices,
            new_indices,
            expert_weights,
            ep_group,
            is_profile=False,
479
            communicator=communicator,
480
481
        )

482
483
484
485
486
487
488
489
    # Verify the rearrangement result
    local_ok = verify_expert_weights_after_shuffle(
        expert_weights,
        new_indices,
        hidden_sizes,
        ep_rank,
        num_local_experts,
    )
490

491
    local_ok = (
492
493
494
495
        verify_redundant_experts_have_same_weights(
            expert_weights,
            new_indices,
            hidden_sizes,
496
            ep_rank,
497
498
499
            world_size,
            num_local_experts,
        )
500
501
502
503
504
505
        and local_ok
    )
    assert_verification_synced(
        local_ok,
        "Rearrange verification failed on at least one rank. See logs for details.",
    )
506
507


508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
@pytest.mark.parametrize(
    "world_size,num_layers,num_local_experts,num_logical_experts",
    [
        # 2 GPU, 2 experts per GPU
        # 3 logical experts, 4 physical experts, 1 redundant experts
        (2, 1, 2, 3),
        # 2 GPU, 3 experts per GPU
        # 4 logical experts, 6 physical experts, 2 redundant experts
        (2, 2, 3, 4),
        # 2 GPU, 8 experts per GPU
        # 16 logical experts, 16 physical experts, 0 redundant experts
        (2, 4, 8, 16),
        # 4 GPU, 2 experts per GPU
        # 6 logical experts, 8 physical experts, 2 redundant experts
        (4, 1, 2, 6),
        # 4 GPU, 2 experts per GPU
        # 5 logical experts, 8 physical experts, 3 redundant experts
        (4, 2, 2, 5),
        # 4 GPU, 8 experts per GPU
        # 16 logical experts, 32 physical experts, 16 redundant experts
        (4, 8, 8, 16),
529
530
    ],
)
531
532
533
@pytest.mark.parametrize(
    "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
)
534
def test_rearrange_expert_weights_with_redundancy(
535
536
537
538
539
    world_size,
    num_layers,
    num_local_experts,
    num_logical_experts,
    eplb_communicator,
540
):
541
542
    """Test the functionality of rearranging expert weights with redundancy."""

543
544
    if eplb_communicator == "nixl" and not has_nixl():
        pytest.skip("NIXL is not available")
545
    if torch.accelerator.device_count() < world_size:
546
        pytest.skip(f"Need at least {world_size} GPUs to run the test")
547
548
549
550
551
552
    distributed_run(
        _test_rearrange_expert_weights_with_redundancy,
        world_size,
        num_layers,
        num_local_experts,
        num_logical_experts,
553
        eplb_communicator,
554
555
556
557
558
559
    )


def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
    set_env_vars_and_device(env)

560
561
    vllm_config = VllmConfig()
    vllm_config.parallel_config.tensor_parallel_size = world_size
562

563
564
565
566
    with set_current_vllm_config(vllm_config):
        ensure_model_parallel_initialized(
            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
        )
567

568
569
        ep_group_coordinator = get_tp_group()
        ep_group = ep_group_coordinator.cpu_group
570
571
        ep_rank = torch.distributed.get_rank()
        device = torch.device(f"cuda:{ep_rank}")
572

573
574
575
576
577
        num_layers = 2
        num_local_experts = 2
        total_physical_experts = world_size * num_local_experts
        num_logical_experts = total_physical_experts // 2  # Some redundancy
        hidden_sizes = [32, 64]
578

579
580
        # Create redundancy configuration
        redundancy_config = [2] * num_logical_experts
581

582
583
584
585
        # Same indices - no change
        indices = create_expert_indices_with_redundancy(
            num_layers, num_logical_experts, total_physical_experts, redundancy_config
        )
586

587
588
589
590
591
592
593
594
595
596
597
598
        expert_weights = create_expert_weights(
            num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
        )

        # Save original weights
        original_weights = []
        for layer_weights in expert_weights:
            layer_copy = []
            for weight in layer_weights:
                layer_copy.append(weight.clone())
            original_weights.append(layer_copy)

599
600
601
602
603
604
        communicator = create_eplb_communicator_or_raise(
            group_coordinator=ep_group_coordinator,
            backend="torch_nccl",
            expert_weights=expert_weights[0],
        )

605
606
607
608
609
610
        # Execute rearrangement (should be no change)
        rearrange_expert_weights_inplace(
            indices,
            indices,  # Same indices
            expert_weights,
            ep_group,
611
            communicator,
612
613
614
            is_profile=False,
        )

615
616
617
618
619
620
621
622
623
624
625
626
627
    # Verify that the weights have not changed
    local_ok = True
    for layer in range(num_layers):
        for weight_idx in range(len(hidden_sizes)):
            if not torch.equal(
                expert_weights[layer][weight_idx],
                original_weights[layer][weight_idx],
            ):
                local_ok = False
                print(
                    "test_rearrange_expert_weights_no_change failed: "
                    f"layer={layer}, weight_idx={weight_idx}",
                    flush=True,
628
                )
629
630
631
632
    assert_verification_synced(
        local_ok,
        "No-change EPLB verification failed on at least one rank.",
    )
633
634


635
636
637
638
639
640
@pytest.mark.parametrize(
    "world_size,num_layers,num_local_experts,num_logical_experts",
    [
        (2, 2, 2, 3),
    ],
)
641
642
643
@pytest.mark.parametrize(
    "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
)
644
645
646
647
648
def test_async_transfer_layer_without_mtp(
    world_size: int,
    num_layers: int,
    num_local_experts: int,
    num_logical_experts: int,
649
    eplb_communicator: str,
650
651
652
):
    """Exercise async EPLB transfer path without MTP/spec decode."""

653
654
    if eplb_communicator == "nixl" and not has_nixl():
        pytest.skip("NIXL is not available")
655
    if torch.accelerator.device_count() < world_size:
656
657
658
659
660
661
662
663
        pytest.skip(f"Need at least {world_size} GPUs to run the test")

    distributed_run(
        _test_async_transfer_layer_without_mtp_worker,
        world_size,
        num_layers,
        num_local_experts,
        num_logical_experts,
664
        eplb_communicator,
665
666
667
    )


668
669
670
671
672
673
674
@pytest.mark.parametrize("world_size", [2, 4])
def test_rearrange_expert_weights_no_change(world_size):
    """
    Test that when the indices do not change, the weights should remain
    unchanged.
    """

675
    if torch.accelerator.device_count() < world_size:
676
        pytest.skip(f"Need at least {world_size} GPUs to run the test")
677
678
679
680
    distributed_run(
        _test_rearrange_expert_weights_no_change,
        world_size,
    )
681
682


683
684
685
def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
    set_env_vars_and_device(env)

686
687
    vllm_config = VllmConfig()
    vllm_config.parallel_config.tensor_parallel_size = world_size
688

689
690
691
692
    with set_current_vllm_config(vllm_config):
        ensure_model_parallel_initialized(
            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
        )
693

694
695
        ep_group_coordinator = get_tp_group()
        ep_group = ep_group_coordinator.cpu_group
696
697
        ep_rank = torch.distributed.get_rank()
        device = torch.device(f"cuda:{ep_rank}")
698

699
700
701
702
703
        num_layers = 1
        num_local_experts = 2
        total_physical_experts = world_size * num_local_experts
        num_logical_experts = total_physical_experts // 2
        hidden_sizes = [32]
704

705
706
707
708
709
710
711
        # Create different index distributions
        old_redundancy = create_redundancy_config(
            num_logical_experts, total_physical_experts
        )
        new_redundancy = create_redundancy_config(
            num_logical_experts, total_physical_experts
        )
712

713
714
715
716
717
718
        old_indices = create_expert_indices_with_redundancy(
            num_layers, num_logical_experts, total_physical_experts, old_redundancy
        )
        new_indices = create_expert_indices_with_redundancy(
            num_layers, num_logical_experts, total_physical_experts, new_redundancy
        )
719

720
721
722
723
724
725
726
727
728
729
730
731
        expert_weights = create_expert_weights(
            num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
        )

        # Save original weights
        original_weights = []
        for layer_weights in expert_weights:
            layer_copy = []
            for weight in layer_weights:
                layer_copy.append(weight.clone())
            original_weights.append(layer_copy)

732
733
734
735
736
737
        communicator = create_eplb_communicator_or_raise(
            group_coordinator=ep_group_coordinator,
            backend="torch_nccl",
            expert_weights=expert_weights[0],
        )

738
739
740
741
742
743
        # Execute profile mode rearrangement
        rearrange_expert_weights_inplace(
            old_indices,
            new_indices,
            expert_weights,
            ep_group,
744
            communicator,
745
746
747
            is_profile=True,  # Profile mode
        )

748
749
750
751
752
753
754
755
756
757
758
759
760
    # In profile mode, the weights should remain unchanged
    local_ok = True
    for layer in range(num_layers):
        for weight_idx in range(len(hidden_sizes)):
            if not torch.equal(
                expert_weights[layer][weight_idx],
                original_weights[layer][weight_idx],
            ):
                local_ok = False
                print(
                    "test_rearrange_expert_weights_profile_mode failed: "
                    f"layer={layer}, weight_idx={weight_idx}",
                    flush=True,
761
                )
762
763
764
765
    assert_verification_synced(
        local_ok,
        "Profile-mode EPLB verification failed on at least one rank.",
    )
766
767
768
769
770
771


@pytest.mark.parametrize("world_size", [2, 4])
def test_rearrange_expert_weights_profile_mode(world_size):
    """Test profile mode (should not copy actual weights)"""

772
    if torch.accelerator.device_count() < world_size:
773
        pytest.skip(f"Need at least {world_size} GPUs to run the test")
774
775
776
777
    distributed_run(
        _test_rearrange_expert_weights_profile_mode,
        world_size,
    )