@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / abdbb68386

Fork 0

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 5, 18:30 UTC

[EPLB] Add alternative communication for EPLB weight exchange (#33176)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com>

Ilya Markov committed Mar 31, 2026 at 14:17 UTC abdbb683861a2586817ae6836bad308c09919d5c

12 files changed +635 -183

.buildkite/test_areas/expert_parallelism.yaml

+2 -2

@@ -13,8 +13,8 @@ steps:

13 - pytest -v -s distributed/test_eplb_algo.py

14 - pytest -v -s distributed/test_eplb_utils.py

16 -- label: EPLB Execution

17 - timeout_in_minutes: 20

16 +- label: EPLB Execution # 17min

17 + timeout_in_minutes: 27

18 working_dir: "/vllm-workspace/tests"

19 num_devices: 4

20 source_file_dependencies:

tests/distributed/eplb_utils.py

+39 -2

@@ -1,9 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import atexit

5 import os

6 import random

8 +import pytest

9 import torch

10 import torch.multiprocessing as mp

@@ -16,9 +18,20 @@ from vllm.utils.system_utils import update_environment_variables

18 mp.set_start_method("spawn", force=True)

21 +def _distributed_worker_wrapper(fn, env, world_size, args, rank, skip_queue):

22 + try:

23 + fn(env, world_size, *args)

24 + except BaseException as exc:

25 + if isinstance(exc, pytest.skip.Exception):

26 + skip_queue.put((rank, str(exc)))

27 + return

28 + raise

29 +

30 +

31 def distributed_run(fn, world_size, *args):

32 number_of_processes = world_size

33 processes: list[mp.Process] = []

34 + skip_queue: mp.SimpleQueue = mp.SimpleQueue()

35 for i in range(number_of_processes):

36 env: dict[str, str] = {}

37 env["RANK"] = str(i)

@@ -27,13 +40,32 @@ def distributed_run(fn, world_size, *args):

40 env["LOCAL_WORLD_SIZE"] = str(number_of_processes)

41 env["MASTER_ADDR"] = "localhost"

42 env["MASTER_PORT"] = "12345"

30 - p = mp.Process(target=fn, args=(env, world_size, *args))

43 + p = mp.Process(

44 + target=_distributed_worker_wrapper,

45 + args=(fn, env, world_size, args, i, skip_queue),

46 + )

47 processes.append(p)

48 p.start()

50 for p in processes:

51 p.join()

53 + skipped: list[tuple[int, str]] = []

54 + while not skip_queue.empty():

55 + rank, reason = skip_queue.get()

56 + skipped.append((rank, reason))

57 +

58 + if len(skipped) == number_of_processes:

59 + reason = skipped[0][1]

60 + pytest.skip(reason)

61 + if 0 < len(skipped) < number_of_processes:

62 + skipped_ranks = sorted(rank for rank, _ in skipped)

63 + raise AssertionError(

64 + "Distributed test had partial skips; expected either all ranks "

65 + f"to skip or none. Skipped ranks: {skipped_ranks}, "

66 + f"total ranks: {number_of_processes}"

67 + )

68 +

69 for p in processes:

70 assert p.exitcode == 0

@@ -48,7 +80,12 @@ def set_env_vars_and_device(env: dict[str, str]) -> None:

80 vllm_config = VllmConfig()

81 with set_current_vllm_config(vllm_config):

82 init_distributed_environment()

51 -

83 + atexit.register(_destroy_process_group_if_initialized)

84 # Ensure each worker process has the same random seed

85 random.seed(42)

86 torch.manual_seed(42)

87 +

88 +

89 +def _destroy_process_group_if_initialized() -> None:

90 + if torch.distributed.is_available() and torch.distributed.is_initialized():

91 + torch.distributed.destroy_process_group()

tests/distributed/test_eplb_execute.py

+195 -55

@@ -9,6 +9,7 @@ import torch

9 import torch.distributed

11 from vllm.config import VllmConfig, set_current_vllm_config

12 +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

13 from vllm.distributed.eplb.rebalance_execute import (

14 move_from_buffer,

15 rearrange_expert_weights_inplace,

@@ -130,9 +131,10 @@ def verify_expert_weights_after_shuffle(

131 hidden_sizes: list[int],

132 ep_rank: int,

133 num_local_experts: int,

133 -):

134 +) -> bool:

135 """Verify the weights after shuffling are correct."""

136 num_layers = len(expert_weights)

137 + ok = True

138

139 for layer in range(num_layers):

140 for weight_idx, hidden_size in enumerate(hidden_sizes):

@@ -155,29 +157,38 @@ def verify_expert_weights_after_shuffle(

157 dtype=actual_weights.dtype,

158 )

159

158 - torch.testing.assert_close(

159 - actual_weights,

160 - expected_weights,

161 - msg=f"Layer {layer}, weight {weight_idx},"

162 - f"local expert {local_expert}: "

163 - f"weights do not match. "

164 - f"Expected logical expert {expected_logical_expert}",

165 - )

160 + if not torch.equal(actual_weights, expected_weights):

161 + ok = False

162 + actual_head = actual_weights[:8].detach().cpu().tolist()

163 + expected_head = expected_weights[:8].detach().cpu().tolist()

164 + print(

165 + "verify_expert_weights_after_shuffle failed: "

166 + f"rank={ep_rank}, "

167 + f"layer={layer}, weight_idx={weight_idx}, "

168 + f"local_expert={local_expert}, "

169 + f"expected_logical_expert={expected_logical_expert}, "

170 + f"actual_head={actual_head}, expected_head={expected_head}",

171 + flush=True,

172 + )

173 +

174 + return ok

175

176

177 def verify_redundant_experts_have_same_weights(

178 expert_weights: list[list[torch.Tensor]],

179 indices: torch.Tensor,

180 hidden_sizes: list[int],

181 + ep_rank: int,

182 world_size: int,

183 num_local_experts: int,

174 -):

184 +) -> bool:

185 """

186 Verify that all replicas of the same logical expert have the same weights.

187 """

188 num_layers = len(expert_weights)

189 total_physical_experts = world_size * num_local_experts

190

191 + ok = True

192 for layer in range(num_layers):

193 # Collect weights for all physical experts for each weight matrix

194 all_weights: list[torch.Tensor] = []

@@ -227,14 +238,54 @@ def verify_redundant_experts_have_same_weights(

238 # Verify that current physical expert's weights match the

239 # previously saved logical expert weights

240 for weight_idx in range(len(hidden_sizes)):

230 - torch.testing.assert_close(

241 + if not torch.equal(

242 all_weights[weight_idx][physical_pos],

243 logical_expert_weights[logical_expert_id][weight_idx],

233 - msg=f"Layer {layer}, weight {weight_idx},"

234 - f"logical expert {logical_expert_id}: "

235 - f"Physical expert {physical_pos} has different weights"

236 - f"than expected",

237 - )

244 + ):

245 + ok = False

246 + actual_head = (

247 + all_weights[weight_idx][physical_pos][:8]

248 + .detach()

249 + .cpu()

250 + .tolist()

251 + )

252 + reference_head = (

253 + logical_expert_weights[logical_expert_id][weight_idx][:8]

254 + .detach()

255 + .cpu()

256 + .tolist()

257 + )

258 + print(

259 + "verify_redundant_experts_have_same_weights failed: "

260 + f"rank={ep_rank}, "

261 + f"layer={layer}, weight_idx={weight_idx}, "

262 + f"logical_expert={logical_expert_id}, "

263 + f"physical_pos={physical_pos}, "

264 + f"actual_head={actual_head}, "

265 + f"reference_head={reference_head}",

266 + flush=True,

267 + )

268 +

269 + return ok

270 +

271 +

272 +def assert_verification_synced(local_ok: bool, msg: str) -> None:

273 + ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32)

274 + torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)

275 + assert bool(ok_tensor.item()), msg

276 +

277 +

278 +def create_eplb_communicator_or_raise(*, group_coordinator, backend, expert_weights):

279 + try:

280 + return create_eplb_communicator(

281 + group_coordinator=group_coordinator,

282 + backend=backend,

283 + expert_weights=expert_weights,

284 + )

285 + except Exception as exc:

286 + raise RuntimeError(

287 + f"Failed to create EPLB communicator for backend={backend}: {exc}"

288 + ) from exc

289

290

291 def _test_async_transfer_layer_without_mtp_worker(

@@ -243,6 +294,7 @@ def _test_async_transfer_layer_without_mtp_worker(

294 num_layers: int,

295 num_local_experts: int,

296 num_logical_experts: int,

297 + eplb_communicator: str,

298 ) -> None:

299 set_env_vars_and_device(env)

300

@@ -254,8 +306,8 @@ def _test_async_transfer_layer_without_mtp_worker(

306 tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

307 )

308

257 - tp_group = get_tp_group()

258 - ep_group = tp_group.device_group

309 + ep_group_coordinator = get_tp_group()

310 + ep_group = ep_group_coordinator.device_group

311 ep_rank = torch.distributed.get_rank()

312 device = torch.device(f"cuda:{ep_rank}")

313

@@ -298,6 +350,13 @@ def _test_async_transfer_layer_without_mtp_worker(

350 expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

351 cuda_stream = torch.cuda.Stream(device=device)

352

353 + communicator = create_eplb_communicator_or_raise(

354 + group_coordinator=ep_group_coordinator,

355 + backend=eplb_communicator,

356 + expert_weights=expert_weights[0],

357 + )

358 + communicator.set_stream(cuda_stream)

359 +

360 for layer_idx in range(num_layers):

361 is_unchanged, is_received_locally, recv_metadata = asyncio.run(

362 transfer_layer(

@@ -306,6 +365,7 @@ def _test_async_transfer_layer_without_mtp_worker(

365 expert_weights=expert_weights[layer_idx],

366 expert_weights_buffer=expert_buffer,

367 ep_group=ep_group,

368 + communicator=communicator,

369 cuda_stream=cuda_stream,

370 )

371 )

@@ -320,24 +380,38 @@ def _test_async_transfer_layer_without_mtp_worker(

380 ep_rank=ep_rank,

381 )

382

323 - verify_expert_weights_after_shuffle(

324 - expert_weights,

325 - new_indices,

326 - hidden_sizes,

327 - ep_rank,

328 - num_local_experts,

329 - )

383 + local_ok = verify_expert_weights_after_shuffle(

384 + expert_weights,

385 + new_indices,

386 + hidden_sizes,

387 + ep_rank,

388 + num_local_experts,

389 + )

390 + local_ok = (

391 verify_redundant_experts_have_same_weights(

392 expert_weights,

393 new_indices,

394 hidden_sizes,

395 + ep_rank,

396 world_size,

397 num_local_experts,

398 )

399 + and local_ok

400 + )

401 + assert_verification_synced(

402 + local_ok,

403 + "Async transfer verification failed on at least one rank. "

404 + "See logs for details.",

405 + )

406

407

408 def _test_rearrange_expert_weights_with_redundancy(

340 - env, world_size, num_layers, num_local_experts, num_logical_experts

409 + env,

410 + world_size,

411 + num_layers,

412 + num_local_experts,

413 + num_logical_experts,

414 + eplb_communicator: str,

415 ) -> None:

416 # Initialize model parallel (using tensor parallel as an entrypoint

417 # to expert parallel)

@@ -351,7 +425,8 @@ def _test_rearrange_expert_weights_with_redundancy(

425 tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

426 )

427

354 - ep_group = get_tp_group().cpu_group

428 + ep_group_coordinator = get_tp_group()

429 + ep_group = ep_group_coordinator.cpu_group

430 ep_rank = torch.distributed.get_rank()

431 device = torch.device(f"cuda:{ep_rank}")

432

@@ -387,6 +462,12 @@ def _test_rearrange_expert_weights_with_redundancy(

462 num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices

463 )

464

465 + communicator = create_eplb_communicator_or_raise(

466 + group_coordinator=ep_group_coordinator,

467 + backend=eplb_communicator,

468 + expert_weights=expert_weights[0],

469 + )

470 +

471 # Execute weight rearrangement

472 rearrange_expert_weights_inplace(

473 old_indices,

@@ -394,24 +475,33 @@ def _test_rearrange_expert_weights_with_redundancy(

475 expert_weights,

476 ep_group,

477 is_profile=False,

478 + communicator=communicator,

479 )

480

399 - # Verify the rearrangement result

400 - verify_expert_weights_after_shuffle(

401 - expert_weights,

402 - new_indices,

403 - hidden_sizes,

404 - ep_rank,

405 - num_local_experts,

406 - )

481 + # Verify the rearrangement result

482 + local_ok = verify_expert_weights_after_shuffle(

483 + expert_weights,

484 + new_indices,

485 + hidden_sizes,

486 + ep_rank,

487 + num_local_experts,

488 + )

489

490 + local_ok = (

491 verify_redundant_experts_have_same_weights(

492 expert_weights,

493 new_indices,

494 hidden_sizes,

495 + ep_rank,

496 world_size,

497 num_local_experts,

498 )

499 + and local_ok

500 + )

501 + assert_verification_synced(

502 + local_ok,

503 + "Rearrange verification failed on at least one rank. See logs for details.",

504 + )

505

506

507 @pytest.mark.parametrize(

@@ -437,8 +527,13 @@ def _test_rearrange_expert_weights_with_redundancy(

527 (4, 8, 8, 16),

528 ],

529 )

530 +@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])

531 def test_rearrange_expert_weights_with_redundancy(

441 - world_size, num_layers, num_local_experts, num_logical_experts

532 + world_size,

533 + num_layers,

534 + num_local_experts,

535 + num_logical_experts,

536 + eplb_communicator,

537 ):

538 """Test the functionality of rearranging expert weights with redundancy."""

539

@@ -450,6 +545,7 @@ def test_rearrange_expert_weights_with_redundancy(

545 num_layers,

546 num_local_experts,

547 num_logical_experts,

548 + eplb_communicator,

549 )

550

551

@@ -464,7 +560,8 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

560 tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

561 )

562

467 - ep_group = get_tp_group().cpu_group

563 + ep_group_coordinator = get_tp_group()

564 + ep_group = ep_group_coordinator.cpu_group

565 ep_rank = torch.distributed.get_rank()

566 device = torch.device(f"cuda:{ep_rank}")

567

@@ -494,24 +591,40 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

591 layer_copy.append(weight.clone())

592 original_weights.append(layer_copy)

593

594 + communicator = create_eplb_communicator_or_raise(

595 + group_coordinator=ep_group_coordinator,

596 + backend="torch_nccl",

597 + expert_weights=expert_weights[0],

598 + )

599 +

600 # Execute rearrangement (should be no change)

601 rearrange_expert_weights_inplace(

602 indices,

603 indices, # Same indices

604 expert_weights,

605 ep_group,

606 + communicator,

607 is_profile=False,

608 )

609

506 - # Verify that the weights have not changed

507 - for layer in range(num_layers):

508 - for weight_idx in range(len(hidden_sizes)):

509 - torch.testing.assert_close(

510 - expert_weights[layer][weight_idx],

511 - original_weights[layer][weight_idx],

512 - msg=f"""Layer {layer}, weight {weight_idx}

513 - should remain unchanged""",

610 + # Verify that the weights have not changed

611 + local_ok = True

612 + for layer in range(num_layers):

613 + for weight_idx in range(len(hidden_sizes)):

614 + if not torch.equal(

615 + expert_weights[layer][weight_idx],

616 + original_weights[layer][weight_idx],

617 + ):

618 + local_ok = False

619 + print(

620 + "test_rearrange_expert_weights_no_change failed: "

621 + f"layer={layer}, weight_idx={weight_idx}",

622 + flush=True,

623 )

624 + assert_verification_synced(

625 + local_ok,

626 + "No-change EPLB verification failed on at least one rank.",

627 + )

628

629

630 @pytest.mark.parametrize(

@@ -520,11 +633,13 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

633 (2, 2, 2, 3),

634 ],

635 )

636 +@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])

637 def test_async_transfer_layer_without_mtp(

638 world_size: int,

639 num_layers: int,

640 num_local_experts: int,

641 num_logical_experts: int,

642 + eplb_communicator: str,

643 ):

644 """Exercise async EPLB transfer path without MTP/spec decode."""

645

@@ -537,6 +652,7 @@ def test_async_transfer_layer_without_mtp(

652 num_layers,

653 num_local_experts,

654 num_logical_experts,

655 + eplb_communicator,

656 )

657

658

@@ -549,7 +665,10 @@ def test_rearrange_expert_weights_no_change(world_size):

665

666 if torch.accelerator.device_count() < world_size:

667 pytest.skip(f"Need at least {world_size} GPUs to run the test")

552 - distributed_run(_test_rearrange_expert_weights_no_change, world_size)

668 + distributed_run(

669 + _test_rearrange_expert_weights_no_change,

670 + world_size,

671 + )

672

673

674 def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

@@ -563,7 +682,8 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

682 tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

683 )

684

566 - ep_group = get_tp_group().cpu_group

685 + ep_group_coordinator = get_tp_group()

686 + ep_group = ep_group_coordinator.cpu_group

687 ep_rank = torch.distributed.get_rank()

688 device = torch.device(f"cuda:{ep_rank}")

689

@@ -600,23 +720,40 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

720 layer_copy.append(weight.clone())

721 original_weights.append(layer_copy)

722

723 + communicator = create_eplb_communicator_or_raise(

724 + group_coordinator=ep_group_coordinator,

725 + backend="torch_nccl",

726 + expert_weights=expert_weights[0],

727 + )

728 +

729 # Execute profile mode rearrangement

730 rearrange_expert_weights_inplace(

731 old_indices,

732 new_indices,

733 expert_weights,

734 ep_group,

735 + communicator,

736 is_profile=True, # Profile mode

737 )

738

612 - # In profile mode, the weights should remain unchanged

613 - for layer in range(num_layers):

614 - for weight_idx in range(len(hidden_sizes)):

615 - torch.testing.assert_close(

616 - expert_weights[layer][weight_idx],

617 - original_weights[layer][weight_idx],

618 - msg="In profile mode, the weights should remain unchanged",

739 + # In profile mode, the weights should remain unchanged

740 + local_ok = True

741 + for layer in range(num_layers):

742 + for weight_idx in range(len(hidden_sizes)):

743 + if not torch.equal(

744 + expert_weights[layer][weight_idx],

745 + original_weights[layer][weight_idx],

746 + ):

747 + local_ok = False

748 + print(

749 + "test_rearrange_expert_weights_profile_mode failed: "

750 + f"layer={layer}, weight_idx={weight_idx}",

751 + flush=True,

752 )

753 + assert_verification_synced(

754 + local_ok,

755 + "Profile-mode EPLB verification failed on at least one rank.",

756 + )

757

758

759 @pytest.mark.parametrize("world_size", [2, 4])

@@ -625,4 +762,7 @@ def test_rearrange_expert_weights_profile_mode(world_size):

762

763 if torch.accelerator.device_count() < world_size:

764 pytest.skip(f"Need at least {world_size} GPUs to run the test")

628 - distributed_run(_test_rearrange_expert_weights_profile_mode, world_size)

765 + distributed_run(

766 + _test_rearrange_expert_weights_profile_mode,

767 + world_size,

768 + )

vllm/config/parallel.py

+22 -10

@@ -35,6 +35,7 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]

35 DataParallelBackend = Literal["ray", "mp"]

36 EPLBPolicyOption = Literal["default"]

37 DCPCommBackend = Literal["ag_rs", "a2a"]

38 +EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "pynccl"]

39 All2AllBackend = Literal[

40 "naive",

41 "pplx",

@@ -83,6 +84,15 @@ class EPLBConfig:

84 policy: EPLBPolicyOption = "default"

85 """The policy type for expert parallel load balancing (EPLB)."""

87 + communicator: EPLBCommunicatorBackend | None = None

88 + """

89 + Backend for EPLB expert weight communication:

90 + - "torch_nccl": Use torch.distributed on the device process group

91 + - "torch_gloo": Use torch.distributed gloo with CPU staging

92 + - "pynccl": Use PyNccl send/recv

93 + - None: Auto-select backend ("torch_gloo" for async, "torch_nccl" for sync)

94 + """

95 +

96 @model_validator(mode="after")

97 def _validate_eplb_config(self) -> Self:

98 if self.use_async and self.policy != "default":

@@ -764,16 +774,18 @@ class ParallelConfig:

774 "backend is mp, uni or external_launcher."

775 )

776

767 - if (

768 - self.all2all_backend in ("allgather_reducescatter")

769 - and self.eplb_config.use_async

770 - ):

771 - logger.warning(

772 - "Async EPLB causes hangs with the '%s' all2all backend. "

773 - "Forcing synchronous EPLB.",

774 - self.all2all_backend,

775 - )

776 - self.eplb_config.use_async = False

777 + if self.enable_eplb and self.eplb_config.communicator is None:

778 + if self.enable_elastic_ep:

779 + # Elastic EP requires stateless mode

780 + # (torch.distributed.batch_isend_irecv doesn't

781 + # support stateless mode), so we use PyNCCL backend

782 + self.eplb_config.communicator = "pynccl"

783 + elif self.eplb_config.use_async:

784 + # Torch Gloo is a backend that allows avoiding hangs

785 + # due to NCCL multi-thread conflicts in async EPLB

786 + self.eplb_config.communicator = "torch_gloo"

787 + else:

788 + self.eplb_config.communicator = "torch_nccl"

789

790 @property

791 def use_ray(self) -> bool:

vllm/distributed/device_communicators/pynccl_wrapper.py

+27 -18

@@ -23,6 +23,7 @@

23 # variable in the code.

25 import ctypes

26 +import functools

27 import platform

28 from dataclasses import dataclass

29 from typing import Any

@@ -75,26 +76,34 @@ class ncclDataTypeEnum:

76 ncclFloat8e4m3 = 10

77 ncclNumTypes = 11

79 + @classmethod

80 + @functools.lru_cache(maxsize=1)

81 + def _torch_to_nccl_map(cls) -> dict[torch.dtype, int]:

82 + return {

83 + torch.int8: cls.ncclInt8,

84 + torch.uint8: cls.ncclUint8,

85 + torch.int32: cls.ncclInt32,

86 + torch.int64: cls.ncclInt64,

87 + torch.float16: cls.ncclFloat16,

88 + torch.float32: cls.ncclFloat32,

89 + torch.float64: cls.ncclFloat64,

90 + torch.bfloat16: cls.ncclBfloat16,

91 + current_platform.fp8_dtype(): cls.ncclFloat8e4m3,

92 + }

93 +

94 + @classmethod

95 + def supports_torch_dtype(cls, dtype: torch.dtype) -> bool:

96 + return dtype in cls._torch_to_nccl_map()

97 +

98 + @classmethod

99 + def try_from_torch(cls, dtype: torch.dtype) -> int | None:

100 + return cls._torch_to_nccl_map().get(dtype)

101 +

102 @classmethod

103 def from_torch(cls, dtype: torch.dtype) -> int:

80 - if dtype == torch.int8:

81 - return cls.ncclInt8

82 - if dtype == torch.uint8:

83 - return cls.ncclUint8

84 - if dtype == torch.int32:

85 - return cls.ncclInt32

86 - if dtype == torch.int64:

87 - return cls.ncclInt64

88 - if dtype == torch.float16:

89 - return cls.ncclFloat16

90 - if dtype == torch.float32:

91 - return cls.ncclFloat32

92 - if dtype == torch.float64:

93 - return cls.ncclFloat64

94 - if dtype == torch.bfloat16:

95 - return cls.ncclBfloat16

96 - if dtype == current_platform.fp8_dtype():

97 - return cls.ncclFloat8e4m3

104 + nccl_dtype = cls.try_from_torch(dtype)

105 + if nccl_dtype is not None:

106 + return nccl_dtype

107 raise ValueError(

108 f"Unsupported dtype {dtype}: should be one of "

109 f"int8, uint8, int32, int64, float16, float32, float64, bfloat16,"

vllm/distributed/elastic_ep/elastic_execute.py

@@ -29,8 +29,10 @@ from vllm.distributed.elastic_ep.standby_state import (

29 get_standby_ep_group,

30 pop_standby_groups,

31 )

32 +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

33 from vllm.distributed.parallel_state import (

34 _replace_active_groups,

35 + get_eplb_group,

36 prepare_communication_buffer_for_model,

37 )

38 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

@@ -411,6 +413,13 @@ class ElasticEPScalingExecutor:

413 module.quant_method = module.quant_method.old_quant_method

414 module.runner = module._init_runner()

415 prepare_communication_buffer_for_model(self.worker.model_runner.model)

416 +

417 + eplb_model_state.communicator = create_eplb_communicator(

418 + group_coordinator=get_eplb_group(),

419 + backend=parallel_config.eplb_config.communicator,

420 + expert_weights=model.expert_weights[0],

421 + )

422 +

423 if (

424 self.worker.vllm_config.compilation_config.mode

425 == CompilationMode.STOCK_TORCH_COMPILE

vllm/distributed/eplb/async_worker.py

@@ -98,6 +98,8 @@ async def transfer_run_periodically(

99 assert state.is_async

100 for model_state in state.model_states.values():

101 + # Set the async worker's CUDA stream on the communicator

102 + model_state.communicator.set_stream(cuda_stream)

103 rebalancing_algorithm_executed = False

104 physical_to_logical_map_cpu = None

105 current_num_layers = model_state.model.num_moe_layers

@@ -157,6 +159,7 @@ async def transfer_run_periodically(

159 expert_weights=model_state.model.expert_weights[layer_idx],

160 expert_weights_buffer=model_state.expert_buffer,

161 ep_group=eplb_group,

162 + communicator=model_state.communicator,

163 is_profile=is_profile,

164 cuda_stream=cuda_stream,

165 )

vllm/distributed/eplb/eplb_communicator.py new

+277

@@ -0,0 +1,277 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +EPLB communicator implementations and factory.

5 +"""

6 +

7 +from abc import ABC, abstractmethod

8 +from collections.abc import Sequence

9 +

10 +import torch

11 +from torch.distributed import (

12 + P2POp,

13 + ProcessGroup,

14 + batch_isend_irecv,

15 +)

16 +

17 +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

18 +from vllm.distributed.device_communicators.pynccl_wrapper import (

19 + ncclDataTypeEnum,

20 +)

21 +from vllm.distributed.parallel_state import GroupCoordinator, is_local_first_rank

22 +from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

23 +from vllm.logger import init_logger

24 +

25 +logger = init_logger(__name__)

26 +

27 +

28 +class EplbCommunicator(ABC):

29 + """Abstract EPLB communicator for expert weight transfers."""

30 +

31 + @abstractmethod

32 + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:

33 + pass

34 +

35 + @abstractmethod

36 + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:

37 + pass

38 +

39 + @abstractmethod

40 + def execute(self) -> None:

41 + pass

42 +

43 + def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:

44 + self._cuda_stream = cuda_stream

45 +

46 + def _log_initialized(self) -> None:

47 + if is_local_first_rank():

48 + logger.info("Initialized EPLB communicator: %s.", self.__class__.__name__)

49 +

50 +

51 +class TorchDistNcclEplbCommunicator(EplbCommunicator):

52 + """EPLB communicator backed by torch.distributed isend/irecv."""

53 +

54 + def __init__(

55 + self,

56 + ep_group: ProcessGroup,

57 + cuda_stream: torch.cuda.Stream | None = None,

58 + ) -> None:

59 + self._ep_group = ep_group

60 + self._cuda_stream = cuda_stream

61 + self._p2p_ops: list[P2POp] = []

62 + self._log_initialized()

63 +

64 + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:

65 + self._p2p_ops.append(

66 + P2POp(

67 + torch.distributed.isend,

68 + tensor,

69 + dst_rank,

70 + self._ep_group,

71 + )

72 + )

73 +

74 + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:

75 + self._p2p_ops.append(

76 + P2POp(

77 + torch.distributed.irecv,

78 + tensor,

79 + src_rank,

80 + self._ep_group,

81 + )

82 + )

83 +

84 + def execute(self) -> None:

85 + if not self._p2p_ops:

86 + return

87 + try:

88 + with torch.cuda.stream(self._cuda_stream):

89 + reqs = batch_isend_irecv(self._p2p_ops)

90 + for req in reqs:

91 + req.wait()

92 + finally:

93 + self._p2p_ops.clear()

94 +

95 +

96 +class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

97 + """EPLB communicator using gloo P2P with CPU staging."""

98 +

99 + def __init__(

100 + self,

101 + cpu_group: ProcessGroup,

102 + cuda_stream: torch.cuda.Stream | None = None,

103 + ) -> None:

104 + self._cpu_group = cpu_group

105 + self._cuda_stream = cuda_stream

106 + self._ops: list[tuple[str, torch.Tensor, int]] = []

107 + self._log_initialized()

108 +

109 + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:

110 + self._ops.append(("send", tensor, dst_rank))

111 +

112 + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:

113 + self._ops.append(("recv", tensor, src_rank))

114 +

115 + def execute(self) -> None:

116 + if not self._ops:

117 + return

118 +

119 + p2p_ops: list[P2POp] = []

120 + recv_staging: list[tuple[torch.Tensor, torch.Tensor]] = []

121 +

122 + def build_ops() -> None:

123 + for op, tensor, peer_rank in self._ops:

124 + if op == "send":

125 + cpu_tensor = tensor.to(device="cpu", non_blocking=True)

126 + p2p_ops.append(

127 + P2POp(

128 + torch.distributed.isend,

129 + cpu_tensor,

130 + peer_rank,

131 + self._cpu_group,

132 + )

133 + )

134 + continue

135 + cpu_tensor = torch.empty_like(tensor, device="cpu")

136 + p2p_ops.append(

137 + P2POp(

138 + torch.distributed.irecv,

139 + cpu_tensor,

140 + peer_rank,

141 + self._cpu_group,

142 + )

143 + )

144 + recv_staging.append((tensor, cpu_tensor))

145 +

146 + try:

147 + with torch.cuda.stream(self._cuda_stream):

148 + build_ops()

149 + finally:

150 + self._ops.clear()

151 +

152 + # Wait for all D2H copies to finish

153 + # before issuing gloo batch_isend_irecv operations.

154 + if self._cuda_stream is not None:

155 + self._cuda_stream.synchronize()

156 + else:

157 + torch.cuda.current_stream().synchronize()

158 +

159 + reqs = batch_isend_irecv(p2p_ops)

160 + for req in reqs:

161 + req.wait()

162 +

163 + if not recv_staging:

164 + return

165 + with torch.cuda.stream(self._cuda_stream):

166 + for dst_tensor, cpu_tensor in recv_staging:

167 + dst_tensor.copy_(cpu_tensor, non_blocking=True)

168 +

169 +

170 +class PyNcclEplbCommunicator(EplbCommunicator):

171 + """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv."""

172 +

173 + def __init__(

174 + self,

175 + pynccl_comm: PyNcclCommunicator,

176 + cuda_stream: torch.cuda.Stream | None = None,

177 + ) -> None:

178 + self._pynccl_comm = pynccl_comm

179 + self._cuda_stream = cuda_stream

180 + self._group_started = False

181 + self._log_initialized()

182 +

183 + def _ensure_group_started(self) -> None:

184 + if not self._group_started:

185 + self._pynccl_comm.group_start()

186 + self._group_started = True

187 +

188 + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:

189 + self._ensure_group_started()

190 + self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)

191 +

192 + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:

193 + self._ensure_group_started()

194 + self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)

195 +

196 + def execute(self) -> None:

197 + if self._group_started:

198 + self._pynccl_comm.group_end()

199 + self._group_started = False

200 +

201 +

202 +def create_eplb_communicator(

203 + group_coordinator: GroupCoordinator,

204 + backend: str | None,

205 + expert_weights: Sequence[torch.Tensor],

206 +) -> EplbCommunicator:

207 + # Keep a safe default for callers that have not resolved communicator yet.

208 + if backend is None:

209 + backend = "torch_nccl"

210 +

211 + tensor_device_type = expert_weights[0].device.type if expert_weights else "cpu"

212 + torch_group = (

213 + group_coordinator.cpu_group

214 + if tensor_device_type == "cpu"

215 + else group_coordinator.device_group

216 + )

217 +

218 + def _create_pynccl() -> EplbCommunicator:

219 + if tensor_device_type == "cpu":

220 + raise RuntimeError(

221 + "EPLB communicator 'pynccl' supports only cuda-like devices "

222 + f"(got {tensor_device_type})."

223 + )

224 + unsupported_dtypes = sorted(

225 + {

226 + tensor.dtype

227 + for tensor in expert_weights

228 + if not ncclDataTypeEnum.supports_torch_dtype(tensor.dtype)

229 + },

230 + key=str,

231 + )

232 + if unsupported_dtypes:

233 + raise RuntimeError(

234 + "EPLB communicator 'pynccl' requested but expert weights contain "

235 + "unsupported dtypes: "

236 + f"({', '.join(str(dtype) for dtype in unsupported_dtypes)})."

237 + )

238 +

239 + device_comm = group_coordinator.device_communicator

240 + pynccl_comm = (

241 + getattr(device_comm, "pynccl_comm", None)

242 + if device_comm is not None

243 + else None

244 + )

245 + if pynccl_comm is None or pynccl_comm.disabled or not pynccl_comm.available:

246 + raise RuntimeError("EPLB communicator 'pynccl' requested but unavailable.")

247 + try:

248 + return PyNcclEplbCommunicator(pynccl_comm=pynccl_comm)

249 + except Exception as exc:

250 + raise RuntimeError(

251 + f"Failed to initialize PyNcclEplbCommunicator ({exc})."

252 + ) from exc

253 +

254 + is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)

255 + if is_stateless:

256 + if backend not in ("torch_nccl", "pynccl"):

257 + raise ValueError(

258 + f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator "

259 + f"(got '{backend}'). torch_gloo is not supported with stateless groups."

260 + )

261 + if backend == "torch_nccl":

262 + logger.warning(

263 + "Stateless elastic EP requires PyNCCL backend. "

264 + "Forcing EPLB communicator to 'pynccl'."

265 + )

266 + backend = "pynccl"

267 + return _create_pynccl()

268 +

269 + if backend == "torch_gloo":

270 + return TorchDistGlooStagedEplbCommunicator(

271 + cpu_group=group_coordinator.cpu_group,

272 + )

273 + elif backend == "torch_nccl":

274 + return TorchDistNcclEplbCommunicator(ep_group=torch_group)

275 + elif backend == "pynccl":

276 + return _create_pynccl()

277 + raise ValueError(f"Unknown EPLB communicator backend: {backend}")

vllm/distributed/eplb/eplb_state.py

+23 -5

@@ -37,6 +37,7 @@ from torch.distributed import ProcessGroup, all_reduce

37 from vllm.config import ModelConfig, ParallelConfig

38 from vllm.distributed.parallel_state import (

39 get_ep_group,

40 + get_eplb_group,

41 get_node_count,

42 in_the_same_node_as,

43 )

@@ -46,6 +47,7 @@ from vllm.logger import init_logger

47 from vllm.model_executor.models.interfaces import MixtureOfExperts

49 from .async_worker import start_async_worker

50 +from .eplb_communicator import EplbCommunicator, create_eplb_communicator

51 from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy

52 from .rebalance_execute import (

53 RecvMetadata,

@@ -225,6 +227,10 @@ class EplbModelState:

227 """

228 CUDA device index for the async EPLB worker thread.

229 """

230 + communicator: EplbCommunicator

231 + """

232 + The communicator for expert weight transfers.

233 + """

234 new_physical_to_logical_map: torch.Tensor | None = None

235 """

236 intermediate variable between `move_to_buffer` and `move_to_workspace`.

@@ -472,6 +478,12 @@ class EplbState:

478 self._init_should_record_tensor(model)

479 expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]

480

481 + communicator = create_eplb_communicator(

482 + group_coordinator=get_eplb_group(),

483 + backend=self.parallel_config.eplb_config.communicator,

484 + expert_weights=model.expert_weights[0],

485 + )

486 +

487 model_state = EplbModelState(

488 physical_to_logical_map=physical_to_logical_map,

489 logical_to_physical_map=logical_to_physical_map,

@@ -498,6 +510,7 @@ class EplbState:

510 recv_dst_rows=np.array([]),

511 ),

512 cuda_device_index=self.cuda_device_index,

513 + communicator=communicator,

514 new_physical_to_logical_map=None,

515 )

516 self.model_states[model_config.compute_hash()] = model_state

@@ -800,6 +813,7 @@ class EplbState:

813 new_physical_to_logical_map,

814 eplb_model_state.model.expert_weights,

815 ep_group,

816 + eplb_model_state.communicator,

817 is_profile,

818 rank_mapping,

819 )

@@ -923,11 +937,8 @@ class EplbState:

937 new_indices=new_indices,

938 ep_rank=ep_group.rank(),

939 )

926 - # Record event after consuming buffer to signal async thread

927 - # that it's safe to overwrite the intermediate buffer

928 - consumed_event = torch.cuda.Event()

929 - consumed_event.record()

930 - model_state.buffer_consumed_event = consumed_event

940 +

941 + transferred_layer = model_state.layer_to_transfer

942

943 transferred_layer = model_state.layer_to_transfer

944 assert model_state.new_physical_to_logical_map is not None

@@ -936,6 +947,13 @@ class EplbState:

947 new_physical_to_logical_map=model_state.new_physical_to_logical_map,

948 layer=transferred_layer,

949 )

950 +

951 + # Record event after consuming buffer to signal async thread

952 + # that it's safe to overwrite the intermediate buffer

953 + consumed_event = torch.cuda.Event()

954 + consumed_event.record()

955 + model_state.buffer_consumed_event = consumed_event

956 +

957 # After the main thread consumes, advance layer_to_transfer

958 model_state.layer_to_transfer += 1

959 model_state.ep_buffer_ready = 0

vllm/distributed/eplb/eplb_utils.py

+13 -2

@@ -21,6 +21,10 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:

21 is_eplb_enabled = parallel_config.enable_eplb

22 async_eplb = parallel_config.eplb_config.use_async

23 is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency"

24 + is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in (

25 + "torch_nccl",

26 + "pynccl",

27 + )

29 # Override NCCL_MAX_CTAS to avoid hangs when using async EPLB with the

30 # DeepEP low-latency backend.

@@ -39,7 +43,13 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:

43 # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP

44 # cooperative kernel to launch and complete, breaking the deadlock.

45 # See: https://github.com/deepseek-ai/DeepEP/issues/496

42 - if is_data_parallel and is_eplb_enabled and is_deepep_ll and async_eplb:

46 + if (

47 + is_data_parallel

48 + and is_eplb_enabled

49 + and is_deepep_ll

50 + and async_eplb

51 + and is_nccl_based_eplb_communicator

52 + ):

53 current_value_str = os.getenv("NCCL_MAX_CTAS")

55 if current_value_str and current_value_str.isdigit():

@@ -49,6 +59,7 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:

59 os.environ["NCCL_MAX_CTAS"] = str(override_value)

60 logger.info_once(

61 f"EPLB: Setting NCCL_MAX_CTAS={override_value} "

52 - "for expert parallel with EPLB and deepep_low_latency backend",

62 + "for expert parallel with NCCL-based EPLB communicator and "

63 + "deepep_low_latency backend",

64 scope="global",

65 )

vllm/distributed/eplb/rebalance_execute.py

+24 -84

@@ -11,19 +11,9 @@ from dataclasses import dataclass

12 import numpy as np

13 import torch

14 -from torch.distributed import (

15 - P2POp,

16 - ProcessGroup,

17 - all_gather,

18 - batch_isend_irecv,

19 - get_global_rank,

20 -)

14 +from torch.distributed import ProcessGroup, all_gather

22 -from vllm.distributed.parallel_state import get_ep_group

23 -from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

24 -from vllm.logger import init_logger

25 -

26 -logger = init_logger(__name__)

16 +from .eplb_communicator import EplbCommunicator

19 @dataclass

@@ -158,7 +148,8 @@ def move_to_buffer(

148 expert_weights: Sequence[torch.Tensor],

149 expert_weights_buffers: Sequence[torch.Tensor],

150 cuda_stream: torch.cuda.Stream | None,

161 - ep_group: ProcessGroup,

151 + ep_rank: int,

152 + communicator: EplbCommunicator,

153 ) -> MoveToBufferResult:

154 """

155 Rearranges expert weights during EPLB rebalancing.

@@ -172,7 +163,8 @@ def move_to_buffer(

163 expert_weights: Original expert weights for the layer.

164 expert_weights_buffers: Intermediate buffers (one per tensor).

165 cuda_stream: CUDA stream for async copies (can be None for sync mode).

175 - ep_group: Distributed process group for expert parallel comms.

166 + ep_rank: Rank of this process in expert parallel group.

167 + communicator: EplbCommunicator instance for P2P communication.

168

169 Returns:

170 is_unchanged (np.ndarray): (num_local_experts,), True where an expert row

@@ -182,8 +174,6 @@ def move_to_buffer(

174 RecvMetadata: Metadata needed for completing remote weight transfers.

175 """

176 assert old_indices.shape == new_indices.shape

185 - ep_rank = ep_group.rank()

186 -

177 recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_)

178 send_expert_ids = np.full((num_local_experts,), -1, dtype=np.int64)

179 send_src_rows = np.full((num_local_experts,), -1, dtype=np.int32)

@@ -247,22 +237,9 @@ def move_to_buffer(

237 expert = new_local_expert_ids[dst]

238 src_local = expert_to_src_map.get(expert, -1)

239 if src_local != -1:

250 - for w, b in zip(expert_weights, expert_weights_buffers):

251 - b[dst].copy_(w[src_local], non_blocking=True)

252 -

253 - p2p_ops: list[P2POp] = []

254 - if isinstance(get_ep_group(), StatelessGroupCoordinator):

255 - ep_group = get_ep_group()

256 - is_stateless = True

257 - else:

258 - is_stateless = False

259 -

260 - # Pre-compute global ranks mapping (only needed for non-stateless groups)

261 - ep_size = ep_group.size()

262 - if not is_stateless:

263 - rank_to_global = {

264 - rank: get_global_rank(ep_group, rank) for rank in range(ep_size)

265 - }

240 + with torch.cuda.stream(cuda_stream):

241 + for w, b in zip(expert_weights, expert_weights_buffers):

242 + b[dst].copy_(w[src_local], non_blocking=True)

243

244 # 2. Post sends

245 if send_count > 0:

@@ -294,23 +271,8 @@ def move_to_buffer(

271 if recver_pos < len(ranks_to_recv):

272 recv_ranks.append(ranks_to_recv[recver_pos])

273 for dst in recv_ranks:

297 - if is_stateless:

298 - for w in expert_weights:

299 - op = object.__new__(P2POp)

300 - op.op = torch.distributed.isend

301 - op.tensor = w[src]

302 - op.group_peer = dst

303 - p2p_ops.append(op)

304 - else:

305 - dst_global = rank_to_global[dst]

306 - p2p_ops += [

307 - P2POp(

308 - torch.distributed.isend,

309 - w[src],

310 - dst_global,

311 - )

312 - for w in expert_weights

313 - ]

274 + for w in expert_weights:

275 + communicator.add_send(w[src], dst)

276

277 # 3. Post recvs

278 if recv_count > 0:

@@ -339,40 +301,11 @@ def move_to_buffer(

301 src = ranks_to_send[recver_pos // num_dst_per_sender]

302 else:

303 src = ranks_to_send[recver_pos - remainder_start]

342 - if is_stateless:

343 - for b in expert_weights_buffers:

344 - op = object.__new__(P2POp)

345 - op.op = torch.distributed.irecv

346 - op.tensor = b[dst]

347 - op.group_peer = src

348 - p2p_ops.append(op)

349 - else:

350 - src_global = rank_to_global[src]

351 - p2p_ops += [

352 - P2POp(

353 - torch.distributed.irecv,

354 - b[dst],

355 - src_global,

356 - )

357 - for b in expert_weights_buffers

358 - ]

304 + for b in expert_weights_buffers:

305 + communicator.add_recv(b[dst], src)

306

307 # 4. Execute the P2P operations. The real communication happens here.

361 - if p2p_ops and cuda_stream is not None:

362 - with torch.cuda.stream(cuda_stream):

363 - if is_stateless:

364 - ep_group.device_communicator.batch_isend_irecv(p2p_ops)

365 - else:

366 - reqs = batch_isend_irecv(p2p_ops)

367 - for req in reqs:

368 - req.wait()

369 - elif p2p_ops:

370 - if is_stateless:

371 - ep_group.device_communicator.batch_isend_irecv(p2p_ops)

372 - else:

373 - reqs = batch_isend_irecv(p2p_ops)

374 - for req in reqs:

375 - req.wait()

308 + communicator.execute()

309 # wait for the communication to finish

310 return (

311 is_unchanged,

@@ -471,6 +404,7 @@ async def transfer_layer(

404 expert_weights: Sequence[torch.Tensor],

405 expert_weights_buffer: Sequence[torch.Tensor],

406 ep_group: ProcessGroup,

407 + communicator: EplbCommunicator,

408 is_profile: bool = False,

409 cuda_stream: torch.cuda.Stream | None = None,

410 rank_mapping: dict[int, int] | None = None,

@@ -489,6 +423,7 @@ async def transfer_layer(

423 For example, a linear layer may have up and down projection.

424 expert_weights_buffer: Intermediate buffers (one per weight tensor).

425 ep_group: The device process group for expert parallelism.

426 + communicator: EplbCommunicator instance for P2P communication.

427 is_profile (bool): If `True`, do not perform any actual weight copy.

428 This is used during profile run, where we only perform dummy

429 communications to reserve enough memory for the buffers.

@@ -542,7 +477,8 @@ async def transfer_layer(

477 expert_weights=expert_weights,

478 expert_weights_buffers=expert_weights_buffer,

479 cuda_stream=cuda_stream,

545 - ep_group=ep_group,

480 + ep_rank=ep_group.rank(),

481 + communicator=communicator,

482 )

483 return is_unchanged, is_received_locally, recv_metadata

484

@@ -552,6 +488,7 @@ def rearrange_expert_weights_inplace(

488 new_global_expert_indices: torch.Tensor,

489 expert_weights: Sequence[Sequence[torch.Tensor]],

490 ep_group: ProcessGroup,

491 + communicator: EplbCommunicator,

492 is_profile: bool = False,

493 rank_mapping: dict[int, int] | None = None,

494 ) -> None:

@@ -569,6 +506,7 @@ def rearrange_expert_weights_inplace(

506 For example, a linear layer may have up and down projection,

507 so weight_count = 2. Each weight's hidden size can be different.

508 ep_group: The device process group for expert parallelism.

509 + communicator: EplbCommunicator instance for P2P communication.

510 is_profile (bool): If `True`, do not perform any actual weight copy.

511 This is used during profile run, where we only perform dummy

512 communications to reserve enough memory for the buffers.

@@ -599,6 +537,7 @@ def rearrange_expert_weights_inplace(

537 assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)

538

539 ep_size = ep_group.size()

540 + ep_rank = ep_group.rank()

541 assert num_physical_experts == ep_size * num_local_physical_experts

542

543 first_layer_weights = list(expert_weights[0])

@@ -635,7 +574,8 @@ def rearrange_expert_weights_inplace(

574 expert_weights=expert_weights[layer_idx],

575 expert_weights_buffers=weights_buffer,

576 cuda_stream=None,

638 - ep_group=ep_group,

577 + ep_rank=ep_rank,

578 + communicator=communicator,

579 )

580

581 move_from_buffer(

@@ -645,7 +585,7 @@ def rearrange_expert_weights_inplace(

585 is_received_locally=is_received_locally,

586 recv_metadata=recv_metadata,

587 new_indices=new_global_expert_indices_cpu[layer_idx],

648 - ep_rank=ep_group.rank(),

588 + ep_rank=ep_rank,

589 )

590

591

vllm/distributed/parallel_state.py

+1 -5

@@ -1690,11 +1690,7 @@ def initialize_model_parallel(

1690 # using torch.distributed in execution with torch.distributed in EPLB.

1691 global _EPLB

1692 assert _EPLB is None, "EPLB group is already initialized"

1693 - if (

1694 - config is not None

1695 - and config.parallel_config is not None

1696 - and config.parallel_config.enable_eplb

1697 - ):

1693 + if config.parallel_config.enable_eplb:

1694 if enable_elastic_ep:

1695 _EPLB = _init_stateless_group(

1696 group_ranks,