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

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 10, 14:00 UTC

[EPLB] Optimize EPLB with numpy (#29499)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

Ilya Markov committed Jan 7, 2026 at 21:21 UTC 6170d47d2220bc464fb561b93caf799a7560e304

8 files changed +738 -272

tests/distributed/test_eplb_algo.py

+140

@@ -310,3 +310,143 @@ if __name__ == "__main__":

310 print(phy2log)

311

312 test_basic_rebalance()

313 +

314 +

315 +def _make_phy_replicas_idx_from_phy2log(phy2log: torch.Tensor) -> torch.Tensor:

316 + """Create replicas indices mapping from phy2log"""

317 + pr = torch.zeros_like(phy2log)

318 + for layer in range(phy2log.shape[0]):

319 + seen: dict[int, int] = {}

320 + row = phy2log[layer].tolist()

321 + for i, expert in enumerate(row):

322 + r = seen.get(expert, 0)

323 + pr[layer, i] = r

324 + seen[expert] = r + 1

325 + return pr

326 +

327 +

328 +def _validate_intragpu_rearrangement(

329 + old_global_expert_indices: torch.Tensor,

330 + new_phy2log: torch.Tensor,

331 + new_phy_replicas_idx: torch.Tensor,

332 + post_phy2log: torch.Tensor,

333 + post_phy_replicas_idx: torch.Tensor,

334 + num_ranks: int,

335 + slots_per_gpu: int,

336 +):

337 + # Per-GPU checks

338 + for gpu_idx in range(num_ranks):

339 + start = gpu_idx * slots_per_gpu

340 + end = start + slots_per_gpu

341 + old_seg = old_global_expert_indices[0, start:end]

342 + new_seg = new_phy2log[0, start:end]

343 + new_rnk = new_phy_replicas_idx[0, start:end]

344 + post_seg = post_phy2log[0, start:end]

345 + post_rnk = post_phy_replicas_idx[0, start:end]

346 +

347 + # Pairwise equality for (expert, rank) pairs to ensure nothing is lost

348 + def sorted_pairs(seg: torch.Tensor, rnk: torch.Tensor):

349 + pairs = list(zip(seg.tolist(), rnk.tolist()))

350 + pairs.sort()

351 + return pairs

352 +

353 + assert sorted_pairs(post_seg, post_rnk) == sorted_pairs(new_seg, new_rnk), (

354 + f"Per-GPU pairs of (expert,rank) must match new mapping for GPU {gpu_idx}"

355 + )

356 +

357 + # For experts that remain on the same GPU, the old slot is preserved

358 + # for at least one occurrence; rank at that slot must be valid for that expert

359 + old_list = old_seg.tolist()

360 + new_list = new_seg.tolist()

361 + post_list = post_seg.tolist()

362 + remained = set(old_list) & set(new_list)

363 + new_ranks_for_expert: dict[int, list[int]] = {}

364 + for v, r in zip(new_list, new_rnk.tolist()):

365 + new_ranks_for_expert.setdefault(v, []).append(r)

366 + for expert in remained:

367 + old_pos = old_list.index(expert)

368 + assert post_list[old_pos] == expert, (

369 + f"Expert {expert} on GPU {gpu_idx} should stay at old slot {old_pos}"

370 + )

371 + # Rank at preserved slot must be one of the ranks

372 + # the expert has in new mapping

373 + assert post_rnk.tolist()[old_pos] in new_ranks_for_expert[expert], (

374 + f"Rank for expert {expert} at preserved slot on GPU {gpu_idx} "

375 + "must come from new mapping"

376 + )

377 +

378 +

379 +@pytest.mark.parametrize(

380 + "num_ranks, slots_per_gpu, old_phy2log, new_phy2log",

381 + [

382 + pytest.param(

383 + # Setup: 2 GPUs, 4 slots each, 1 layer

384 + # Old mapping: GPU0 -> [0,1,2,3], GPU1 -> [4,5,6,7]

385 + # New mapping shuffles within GPU0 and brings 4,5 into GPU0.

386 + # GPU0 new -> [1,5,0,4]; GPU1 new -> [6,2,7,3]

387 + 2,

388 + 4,

389 + torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]]),

390 + torch.tensor([[1, 5, 0, 4, 6, 2, 7, 3]]),

391 + id="simple",

392 + ),

393 + pytest.param(

394 + # Setup: 2 GPUs, 5 slots each (total 10 physical experts), 1 layer

395 + # Old mapping:

396 + # GPU0 -> [0, 1, 0, 2, 3] (expert 0 duplicated)

397 + # GPU1 -> [4, 5, 6, 1, 2]

398 + # New mapping reorders within GPUs and moves some experts across GPUs,

399 + # while still including duplicates:

400 + # GPU0 new -> [0, 5, 4, 0, 1] (expert 0 duplicated, 4/5 incoming)

401 + # GPU1 new -> [6, 2, 3, 2, 1] (expert 2 duplicated)

402 + 2,

403 + 5,

404 + torch.tensor([[0, 1, 0, 2, 3, 4, 5, 6, 1, 2]]),

405 + torch.tensor([[0, 5, 4, 0, 1, 6, 2, 3, 2, 1]]),

406 + id="duplicates",

407 + ),

408 + pytest.param(

409 + # Setup: 3 GPUs, 4 slots each (total 12 physical experts), 1 layer

410 + # Old mapping:

411 + # GPU0 -> [0, 1, 2, 3]

412 + # GPU1 -> [0, 1, 2, 3]

413 + # GPU2 -> [0, 1, 2, 3]

414 + # New mapping decides to use one expert on 2 GPUs and shuffles

415 + # experts on the third GPU,

416 + # GPU0 new -> [0, 0, 0, 0]

417 + # GPU1 new -> [0, 0, 0, 0]

418 + # GPU2 new -> [1, 2, 3, 0]

419 + 3,

420 + 4,

421 + torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]),

422 + torch.tensor([[0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 0]]),

423 + id="skewed_expert",

424 + ),

425 + ],

426 +)

427 +def test_preserve_intragpu_slots(

428 + num_ranks: int,

429 + slots_per_gpu: int,

430 + old_phy2log: torch.Tensor,

431 + new_phy2log: torch.Tensor,

432 +):

433 + """Experts that stay on a GPU keep their old slots; incoming not lost."""

434 + phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(new_phy2log)

435 +

436 + post_phy2log, post_phy_replicas_idx = DefaultEplbPolicy.preserve_intragpu_slots(

437 + new_phy2log, phy_replicas_idx, num_ranks, old_phy2log

438 + )

439 +

440 + # Shapes preserved

441 + assert post_phy2log.shape == new_phy2log.shape

442 + assert post_phy_replicas_idx.shape == phy_replicas_idx.shape

443 +

444 + _validate_intragpu_rearrangement(

445 + old_phy2log,

446 + new_phy2log,

447 + phy_replicas_idx,

448 + post_phy2log,

449 + post_phy_replicas_idx,

450 + num_ranks,

451 + slots_per_gpu,

452 + )

tests/distributed/test_eplb_execute.py

+9 -8

@@ -286,15 +286,17 @@ def _test_async_transfer_layer_without_mtp_worker(

286 device,

287 old_indices,

288 )

289 + old_indices_cpu = old_indices.cpu()

290 + new_indices_cpu = new_indices.cpu()

291

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

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

294

295 for layer_idx in range(num_layers):

294 - is_unchanged, is_received_locally, experts_recv_loc = asyncio.run(

296 + is_unchanged, is_received_locally, recv_metadata = asyncio.run(

297 transfer_layer(

296 - old_global_expert_indices=old_indices,

297 - new_global_expert_indices=new_indices,

298 + old_global_expert_indices=old_indices_cpu,

299 + new_global_expert_indices=new_indices_cpu,

300 expert_weights=expert_weights,

301 expert_weights_buffer=expert_buffer,

302 ep_group=ep_group,

@@ -302,16 +304,15 @@ def _test_async_transfer_layer_without_mtp_worker(

304 cuda_stream=cuda_stream,

305 )

306 )

305 -

307 cuda_stream.synchronize()

308 move_from_buffer(

309 expert_weights=expert_weights[layer_idx],

309 - expert_weights_buffer=expert_buffer,

310 + expert_weights_buffers=expert_buffer,

311 is_unchanged=is_unchanged,

312 is_received_locally=is_received_locally,

312 - experts_recv_loc=experts_recv_loc,

313 - new_indices=new_indices[layer_idx].tolist(),

314 - ep_group=ep_group,

313 + recv_metadata=recv_metadata,

314 + new_indices=new_indices_cpu[layer_idx],

315 + ep_rank=ep_rank,

316 )

317

318 verify_expert_weights_after_shuffle(

vllm/config/parallel.py

+12

@@ -69,6 +69,10 @@ class EPLBConfig:

69 Log the balancedness each step of expert parallelism.

70 This is turned off by default since it will cause communication overhead.

71 """

72 + log_balancedness_interval: int = 1

73 + """

74 + Interval for logging the balancedness.

75 + """

76 use_async: bool = False

77 """

78 Whether to use non-blocking EPLB.

@@ -77,6 +81,14 @@ class EPLBConfig:

81 policy: EPLBPolicyOption = "default"

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

84 + @model_validator(mode="after")

85 + def _validate_eplb_config(self) -> Self:

86 + if self.use_async and self.policy != "default":

87 + raise ValueError("Async EPLB is only supported with the default policy.")

88 + if self.log_balancedness and self.log_balancedness_interval <= 0:

89 + raise ValueError("log_balancedness_interval must be greater than 0.")

90 + return self

91 +

93 @config

94 @dataclass

vllm/distributed/eplb/async_worker.py

+1 -1

@@ -89,7 +89,7 @@ async def transfer_run_periodically(

89 (

90 model_state.is_unchanged,

91 model_state.is_received_locally,

92 - model_state.experts_recv_loc,

92 + model_state.recv_metadata,

93 ) = await transfer_layer(

94 old_global_expert_indices=model_state.physical_to_logical_map,

95 new_global_expert_indices=model_state.new_physical_to_logical_map,

vllm/distributed/eplb/eplb_state.py

+60 -32

@@ -27,10 +27,10 @@ physical experts.

27 """

29 import threading

30 -import time

30 from collections.abc import Sequence

31 from dataclasses import dataclass

33 +import numpy as np

34 import torch

35 from torch.distributed import ProcessGroup, all_reduce

@@ -46,7 +46,11 @@ from vllm.model_executor.models.interfaces import MixtureOfExperts

47 from .async_worker import start_async_worker

48 from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy

49 -from .rebalance_execute import move_from_buffer, rearrange_expert_weights_inplace

49 +from .rebalance_execute import (

50 + RecvMetadata,

51 + move_from_buffer,

52 + rearrange_expert_weights_inplace,

53 +)

55 logger = init_logger(__name__)

@@ -164,20 +168,19 @@ class EplbModelState:

168 """

169 Whether the async EPLB needs to poll peers for buffer readiness.

170 """

167 - is_unchanged: list[bool]

171 + is_unchanged: np.ndarray

172 """

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

174 The size is same as the num of physical experts in the current layer.

175 """

172 - is_received_locally: list[bool]

176 + is_received_locally: np.ndarray

177 """

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

179 The size is same as the num of physical experts in the current layer.

180 """

177 - experts_recv_loc: dict[int, int]

181 + recv_metadata: RecvMetadata

182 """

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

180 - The size is same as the num of physical experts in the current layer.

184 """

185 is_async_enabled: bool

186 """

@@ -507,9 +510,14 @@ class EplbState:

510 layer_to_transfer=0,

511 rebalanced=False,

512 pending_global_ready_check=False,

510 - is_unchanged=[],

511 - is_received_locally=[],

512 - experts_recv_loc={},

513 + is_unchanged=np.array([]),

514 + is_received_locally=np.array([]),

515 + recv_metadata=RecvMetadata(

516 + recv_primary_mask=np.array([]),

517 + recv_count=0,

518 + recv_expert_ids=np.array([]),

519 + recv_dst_rows=np.array([]),

520 + ),

521 is_async_enabled=self.is_async,

522 cuda_device_index=self.cuda_device_index,

523 new_physical_to_logical_map=new_physical_to_logical_map,

@@ -553,7 +561,12 @@ class EplbState:

561 for eplb_model_state in self.model_states.values():

562 eplb_model_state.expert_load_pass.zero_()

563

556 - if log_stats:

564 + if (

565 + log_stats

566 + and self.expert_rearrangement_step

567 + % self.parallel_config.eplb_config.log_balancedness_interval

568 + == 0

569 + ):

570 # Sync the expert load pass for each model (main and drafter).

571 # expert_load_pass: (num_moe_layers, num_physical_experts)

572 expert_load_pass_list = self._sync_load_pass()

@@ -586,12 +599,15 @@ class EplbState:

599 if ep_group.rank() == 0:

600 logger.info(

601 "EPLB step: %d for model %s: avg_tokens=%.2f, "

589 - "max_tokens=%d, balancedness=%.4f",

602 + "max_tokens=%d, balancedness=%.4f, "

603 + "steps until the next rearrangement: %d",

604 self.expert_rearrangement_step,

605 eplb_model_state.model_name,

606 avg_tokens,

607 max_tokens,

608 balancedness,

609 + self.expert_rearrangement_step_interval

610 + - self.expert_rearrangement_step,

611 )

612

613 # Update the expert load sliding window

@@ -684,11 +700,14 @@ class EplbState:

700 ep_group = get_ep_group().device_group

701 ep_rank = ep_group.rank()

702

687 - time_start = None

703 + start_event = None

704 + end_event = None

705 is_main_rank = ep_rank == 0

706 if is_main_rank:

690 - torch.cuda.synchronize()

691 - time_start = time.perf_counter()

707 + if not self.is_async or is_profile:

708 + start_event = torch.cuda.Event(enable_timing=True)

709 + end_event = torch.cuda.Event(enable_timing=True)

710 + start_event.record()

711 logger.info(

712 "Rearranging experts %s %s...",

713 "(async mode)" if self.is_async else "sync mode",

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

819 num_groups,

820 num_nodes,

821 num_gpus,

822 + eplb_model_state.physical_to_logical_map,

823 )

824

825 if not eplb_model_state.is_async_enabled or is_profile:

@@ -848,17 +868,17 @@ class EplbState:

868 new_logical_replica_count

869 )

870 if is_main_rank:

851 - assert time_start is not None

852 - torch.cuda.synchronize()

853 - time_end = time.perf_counter()

871 + assert start_event is not None

872 + assert end_event is not None

873 + end_event.record()

874 + end_event.synchronize()

875 + gpu_elapsed = start_event.elapsed_time(end_event) / 1000.0

876 logger.info(

855 - "Rearranged experts%sin %.2f seconds.",

877 + "Rearranged experts %s in %.2f s.",

878 " (profile) " if is_profile else " ",

857 - time_end - time_start,

879 + gpu_elapsed,

880 )

881 else:

860 - device = eplb_model_state.physical_to_logical_map.device

861 - new_physical = new_physical_to_logical_map.to(device)

882 max_slots = eplb_model_state.logical_to_physical_map.shape[-1]

883 padded_logical = torch.nn.functional.pad(

884 new_logical_to_physical_map,

@@ -869,7 +889,10 @@ class EplbState:

889 eplb_model_state.logical_replica_count.device

890 )

891

872 - eplb_model_state.new_physical_to_logical_map = new_physical

892 + # Move map to cpu in advance

893 + eplb_model_state.new_physical_to_logical_map = (

894 + new_physical_to_logical_map.cpu()

895 + )

896 eplb_model_state.new_logical_to_physical_map = padded_logical

897 eplb_model_state.new_logical_replica_count = new_replica

898

@@ -968,25 +991,30 @@ class EplbState:

991 stream = torch.cuda.current_stream(device=device_index)

992 stream.wait_event(model_state.buffer_ready_event)

993 model_state.buffer_ready_event = None

994 + expert_weights = model_state.model.expert_weights[

995 + model_state.layer_to_transfer

996 + ]

997 + expert_weights_buffer = model_state.expert_buffer

998 + new_indices = (

999 + model_state.new_physical_to_logical_map[model_state.layer_to_transfer]

1000 + .cpu()

1001 + .numpy()

1002 + )

1003 move_from_buffer(

972 - expert_weights=model_state.model.expert_weights[

973 - model_state.layer_to_transfer

974 - ],

975 - expert_weights_buffer=model_state.expert_buffer,

1004 + expert_weights=expert_weights,

1005 + expert_weights_buffers=expert_weights_buffer,

1006 is_unchanged=model_state.is_unchanged,

1007 is_received_locally=model_state.is_received_locally,

978 - experts_recv_loc=model_state.experts_recv_loc,

979 - new_indices=model_state.new_physical_to_logical_map[

980 - model_state.layer_to_transfer

981 - ].tolist(),

982 - ep_group=ep_group,

1008 + recv_metadata=model_state.recv_metadata,

1009 + new_indices=new_indices,

1010 + ep_rank=ep_group.rank(),

1011 )

1012 transferred_layer = model_state.layer_to_transfer

1013 self._update_layer_mapping_from_new(model_state, transferred_layer)

1014 # After the main thread consumes, advance layer_to_transfer

1015 model_state.layer_to_transfer += 1

1016 model_state.ep_buffer_ready = 0

989 - logger.info(

1017 + logger.debug(

1018 "model %s successfully move_to_workspace layer %d",

1019 model_state.model_name,

1020 transferred_layer,

vllm/distributed/eplb/policy/abstract.py

+4 -1

@@ -16,6 +16,7 @@ class AbstractEplbPolicy(ABC):

16 num_groups: int,

17 num_nodes: int,

18 num_ranks: int,

19 + old_global_expert_indices: torch.Tensor | None = None,

20 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

21 """

22 Entry point for expert-parallelism load balancer.

@@ -28,7 +29,9 @@ class AbstractEplbPolicy(ABC):

29 num_groups: number of expert groups

30 num_nodes: number of server nodes

31 num_ranks: number of ranks, must be a multiple of `num_nodes`

31 -

32 + old_global_expert_indices: [layers, num_logical_experts], the old global

33 + expert indices. Used to avoid unnecessary weight copying

34 + for experts moving within one rank.

35 Returns:

36 physical_to_logical_map: [layers, num_replicas], the expert

37 index of each replica

vllm/distributed/eplb/policy/default.py

+124 -12

@@ -93,7 +93,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

94 Returns:

95 phy2log: [X, num_phy], logical expert id of each physical expert

96 - rank: [X, num_phy], the replica rank

96 + replica_idx: [X, num_phy], the index of the replica for each logical expert

97 logcnt: [X, num_log], number of replicas for each logical expert

98 """

99 n, num_log = weight.shape

@@ -101,15 +101,15 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

101 assert num_redundant >= 0

102 device = weight.device

103 phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)

104 - rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)

104 + replica_idx = torch.zeros(n, num_phy, dtype=torch.int64, device=device)

105 logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)

106 arangen = torch.arange(n, dtype=torch.int64, device=device)

107 for i in range(num_log, num_phy):

108 redundant_indices = (weight / logcnt).max(dim=-1).indices

109 phy2log[:, i] = redundant_indices

110 - rank[:, i] = logcnt[arangen, redundant_indices]

110 + replica_idx[:, i] = logcnt[arangen, redundant_indices]

111 logcnt[arangen, redundant_indices] += 1

112 - return phy2log, rank, logcnt

112 + return phy2log, replica_idx, logcnt

113

114 @classmethod

115 def rebalance_experts_hierarchical(

@@ -132,7 +132,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

132 Returns:

133 phy2log: [layers, num_replicas], the expert

134 index of each replica

135 - log2phy: [layers, num_logical_experts, X],

135 + pphy_replicas_idx: [layers, num_logical_experts, X],

136 the replica indices for each expert

137 logcnt: [layers, num_logical_experts], number of

138 physical replicas for each logical expert

@@ -177,7 +177,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

177 tokens_per_mlog = weight.gather(-1, mlog2log).view(

178 -1, num_logical_experts // num_nodes

179 )

180 - phy2mlog, phyrank, mlogcnt = cls.replicate_experts(

180 + phy2mlog, replicas_idx, mlogcnt = cls.replicate_experts(

181 tokens_per_mlog, num_physical_experts // num_nodes

182 )

183

@@ -203,9 +203,109 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

203 ).view(1, -1, 1)

204 ).flatten(-2)

205 pphy2log = mlog2log.gather(-1, pphy2mlog)

206 - pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)

206 + pphy_replicas_idx = replicas_idx.gather(-1, pphy2phy).view(num_layers, -1)

207 logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)

208 - return pphy2log, pphyrank, logcnt

208 + return pphy2log, pphy_replicas_idx, logcnt

209 +

210 + @classmethod

211 + def preserve_intragpu_slots(

212 + cls,

213 + phy2log: torch.Tensor,

214 + phy_replicas_idx: torch.Tensor,

215 + num_ranks: int,

216 + old_global_expert_indices: torch.Tensor,

217 + ) -> tuple[torch.Tensor, torch.Tensor]:

218 + """

219 + Reorder the new mapping per GPU so that experts that remain on the same GPU

220 + keep their previous slot positions when possible. Incoming experts to that GPU

221 + fill any remaining available slots. This is applied only when the number of GPUs

222 + is unchanged and the slots per GPU remain the same between

223 + the old and new mappings.

224 + """

225 + device = phy2log.device

226 + num_phy_experts = phy2log.shape[1]

227 + if num_ranks <= 0 or num_phy_experts % num_ranks != 0:

228 + return phy2log, phy_replicas_idx

229 +

230 + # Move to CPU and convert to NumPy for processing

231 + new_phy2log_np = phy2log.cpu().numpy()

232 + replicas_idx_np = phy_replicas_idx.cpu().numpy()

233 + old_phy2log_np = old_global_expert_indices.cpu().numpy()

234 +

235 + slots_per_gpu = num_phy_experts // num_ranks

236 + num_layers = new_phy2log_np.shape[0]

237 +

238 + post_phy2log_np = new_phy2log_np.copy()

239 + post_phy_replicas_idx_np = replicas_idx_np.copy()

240 +

241 + for gpu_idx in range(num_ranks):

242 + start = gpu_idx * slots_per_gpu

243 + end = start + slots_per_gpu

244 + # Experts across all layers for this GPU

245 + old_local = old_phy2log_np[:, start:end] # [layers, slots]

246 + new_local = new_phy2log_np[:, start:end] # [layers, slots]

247 + new_ridx = replicas_idx_np[:, start:end] # [layers, slots]

248 +

249 + used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool)

250 + preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool)

251 +

252 + # First pass: preserve same-logical experts in their previous slots

253 + for slot_idx in range(slots_per_gpu):

254 + # matches: [layers, slots], True where new local experts have

255 + # the same logical value as the old from 'slot_idx' and not checked yet

256 + matches = (new_local == old_local[:, slot_idx][:, None]) & (

257 + ~used_new_indices

258 + )

259 + has_any = matches.any(axis=1)

260 + if np.any(has_any):

261 + first_idx = np.argmax(matches, axis=1)

262 + layer_indices = np.nonzero(has_any)[0]

263 + matched_new_positions = first_idx[layer_indices]

264 + post_phy2log_np[layer_indices, start + slot_idx] = new_local[

265 + layer_indices, matched_new_positions

266 + ]

267 + post_phy_replicas_idx_np[layer_indices, start + slot_idx] = (

268 + new_ridx[layer_indices, matched_new_positions]

269 + )

270 + used_new_indices[layer_indices, matched_new_positions] = True

271 + preserved_positions[layer_indices, slot_idx] = True

272 +

273 + # Second pass: fill remaining slots with remaining new experts

274 + remaining_mask = ~used_new_indices # [layers, slots]

275 + fill_mask = ~preserved_positions # [layers, slots]

276 + if remaining_mask.any() and fill_mask.any():

277 + idx_base = np.tile(np.arange(slots_per_gpu), (num_layers, 1))

278 + # Sentinel value for unavailable positions.

279 + large = slots_per_gpu + 1

280 + # Priorities: keep original index for available spots, set sentinel

281 + # for unavailable; lower is earlier.

282 + remaining_priority = np.where(remaining_mask, idx_base, large)

283 + fill_priority = np.where(fill_mask, idx_base, large)

284 + # Sort to get ordered indices of available src/dst positions per layer.

285 + remaining_indices = np.argsort(remaining_priority, axis=1)

286 + fill_indices = np.argsort(fill_priority, axis=1)

287 + # Fill count per layer (cannot exceed either side).

288 + remaining_counts = remaining_mask.sum(axis=1)

289 + fill_counts = fill_mask.sum(axis=1)

290 + take_counts = np.minimum(remaining_counts, fill_counts)

291 + # Assign remaining new experts to remaining slots per layer.

292 + for layer_idx in range(num_layers):

293 + k = int(take_counts[layer_idx])

294 + if k <= 0:

295 + continue

296 + src_pos = remaining_indices[layer_idx, :k]

297 + dst_pos = fill_indices[layer_idx, :k]

298 + post_phy2log_np[layer_idx, start + dst_pos] = new_local[

299 + layer_idx, src_pos

300 + ]

301 + post_phy_replicas_idx_np[layer_idx, start + dst_pos] = new_ridx[

302 + layer_idx, src_pos

303 + ]

304 +

305 + # Convert back to torch and move to original device

306 + post_phy2log = torch.from_numpy(post_phy2log_np).to(device)

307 + post_phy_replicas_idx = torch.from_numpy(post_phy_replicas_idx_np).to(device)

308 + return post_phy2log, post_phy_replicas_idx

309

310 @classmethod

311 def rebalance_experts(

@@ -215,6 +315,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

315 num_groups: int,

316 num_nodes: int,

317 num_ranks: int,

318 + old_global_expert_indices: torch.Tensor | None = None,

319 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

320 """

321 Entry point for expert-parallelism load balancer.

@@ -228,7 +329,9 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

329 num_nodes: number of server nodes, where the intra-node network

330 (e.g, NVLink) is faster

331 num_ranks: number of ranks, must be a multiple of `num_nodes`

231 -

332 + old_global_expert_indices: [layers, num_logical_experts], the old global

333 + expert indices. Used to avoid unnecessary weight copying

334 + for experts moving within one rank.

335 Returns:

336 phy2log: [layers, num_replicas], the expert

337 index of each replica

@@ -241,14 +344,23 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

344 weight = weight.float()

345 if num_groups % num_nodes == 0:

346 # use hierarchical load-balance policy

244 - phy2log, phyrank, logcnt = cls.rebalance_experts_hierarchical(

347 + phy2log, phy_replicas_idx, logcnt = cls.rebalance_experts_hierarchical(

348 weight, num_replicas, num_groups, num_nodes, num_ranks

349 )

350 else:

351 # use global load-balance policy

249 - phy2log, phyrank, logcnt = cls.rebalance_experts_hierarchical(

352 + phy2log, phy_replicas_idx, logcnt = cls.rebalance_experts_hierarchical(

353 weight, num_replicas, 1, 1, num_ranks

354 )

355 + # Optional postprocessing to preserve slots for experts moving

356 + # within the same GPU

357 + # Only apply when the number of GPUs and slots per GPU remain unchanged.

358 + # Helps to avoid unnecessary weight copying when experts move

359 + # within the same GPU.

360 + if old_global_expert_indices is not None:

361 + phy2log, phy_replicas_idx = cls.preserve_intragpu_slots(

362 + phy2log, phy_replicas_idx, num_ranks, old_global_expert_indices

363 + )

364 num_redundant_experts = num_replicas - num_logical_experts

365 maxlogcnt = num_redundant_experts + 1

366 log2phy: torch.Tensor = torch.full(

@@ -259,7 +371,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

371 )

372 log2phy.view(num_layers, -1).scatter_(

373 -1,

262 - phy2log * maxlogcnt + phyrank,

374 + phy2log * maxlogcnt + phy_replicas_idx,

375 torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(

376 num_layers, -1

377 ),

vllm/distributed/eplb/rebalance_execute.py

+388 -218

@@ -6,9 +6,10 @@ The actual execution of the rearrangement.

6 This involves the exchange of expert weights between GPUs.

7 """

9 -from collections.abc import Iterable, MutableSequence, Sequence

10 -from functools import partial

9 +from collections.abc import Iterable, Sequence

10 +from dataclasses import dataclass

12 +import numpy as np

13 import torch

14 from torch.distributed import (

15 P2POp,

@@ -18,213 +19,317 @@ from torch.distributed import (

19 get_global_rank,

20 )

22 +from vllm.logger import init_logger

22 -def idx_local_to_global(

23 - local_idx: int,

24 - local_cnt: int,

25 - ep_rank: int,

26 -) -> int:

27 - """

28 - Convert a local expert index to a global expert index.

29 - """

30 - return ep_rank * local_cnt + local_idx

24 +logger = init_logger(__name__)

33 -def idx_global_to_local(

34 - global_idx: int,

35 - local_cnt: int,

36 - ep_rank: int,

37 -) -> int:

38 - """

39 - Convert a global expert index to a local expert index.

40 - """

41 - return global_idx - ep_rank * local_cnt

27 +@dataclass

28 +class RecvMetadata:

29 + """Metadata describing remote receives during EPLB rebalancing."""

31 + recv_primary_mask: np.ndarray

32 + """Mask of (num_local_experts,) indicating primary experts received."""

33 + recv_count: int

34 + """Number of received experts for the layer."""

35 + recv_expert_ids: np.ndarray

36 + """Expert ids (num_local_experts,) of remote primary experts."""

37 + recv_dst_rows: np.ndarray

38 + """Target expert indices (num_local_experts,) in local tensors to send."""

44 -def global_idx_to_rank(

45 - global_idx: int,

46 - local_cnt: int,

47 -) -> int:

48 - """

49 - Convert a global expert index to a rank index.

50 - """

51 - return global_idx // local_cnt

41 +# Type alias for the result of move_to_buffer or transfer_layer

42 +MoveToBufferResult = tuple[np.ndarray, np.ndarray, RecvMetadata]

54 -def get_ep_ranks_with_expert(

55 - idx: int,

44 +

45 +def get_ep_ranks_with_experts_batch(

46 + expert_ids: np.ndarray,

47 num_local_experts: int,

57 - old_indices: Sequence[int],

58 - new_indices: Sequence[int],

59 -) -> tuple[MutableSequence[int], MutableSequence[int]]:

48 + old_indices: np.ndarray,

49 + new_indices: np.ndarray,

50 +) -> tuple[dict[int, list[int]], dict[int, list[int]]]:

51 """

52 Get the ranks of the experts that need to be exchanged.

54 Args:

64 - idx: The index of the expert.

55 + expert_ids: 1D array of expert indices to query.

56 num_local_experts: The number of local experts.

57 old_indices: The old indices of the experts.

58 new_indices: The new indices of the experts.

60 Returns:

70 - A tuple of two lists:

71 - - The ranks of the experts that need to be sent.

72 - - The ranks of the experts that need to be received.

61 + A tuple of two dictionaries mapping expert_id to:

62 + - ranks_to_send: The ranks that have this expert and need to send.

63 + - ranks_to_recv: The ranks that need to receive this expert.

64 """

74 - global2rank = partial(

75 - global_idx_to_rank,

76 - local_cnt=num_local_experts,

77 - )

65 + ranks_to_send_map: dict[int, list[int]] = {}

66 + ranks_to_recv_map: dict[int, list[int]] = {}

67 +

68 + # Fast path: if no experts, return empty dicts

69 + if expert_ids.size == 0:

70 + return ranks_to_send_map, ranks_to_recv_map

71 +

72 + unique_experts = np.unique(expert_ids)

73 + num_positions = len(old_indices)

74 + position_indices = np.arange(num_positions, dtype=np.int32)

75 +

76 + # Vectorized approach: find all positions matching any query expert in one pass

77 + # Use np.isin to get boolean masks for all relevant positions at once

78 + old_relevant_mask = np.isin(old_indices, unique_experts)

79 + new_relevant_mask = np.isin(new_indices, unique_experts)

80 +

81 + # Process old_indices (send ranks)

82 + if np.any(old_relevant_mask):

83 + old_relevant_positions = position_indices[old_relevant_mask]

84 + old_relevant_experts = old_indices[old_relevant_mask]

85 + old_relevant_ranks = old_relevant_positions // num_local_experts

86 +

87 + # Sort by expert first, then by position (to maintain first-appearance order)

88 + sort_order = np.lexsort((old_relevant_positions, old_relevant_experts))

89 + sorted_experts = old_relevant_experts[sort_order]

90 + sorted_ranks = old_relevant_ranks[sort_order]

91 +

92 + # Find boundaries where expert changes

93 + expert_boundaries = np.concatenate(

94 + [[0], np.where(np.diff(sorted_experts) != 0)[0] + 1, [len(sorted_experts)]]

95 + )

79 - ranks_to_send: list[int] = []

80 - ranks_to_recv: list[int] = []

81 -

82 - for i, e in enumerate(old_indices):

83 - if e == idx:

84 - rank = global2rank(i)

85 - if not ranks_to_send or ranks_to_send[-1] != rank:

86 - ranks_to_send.append(rank)

87 -

88 - for i, e in enumerate(new_indices):

89 - if e == idx:

90 - rank = global2rank(i)

91 - if not ranks_to_recv or ranks_to_recv[-1] != rank:

92 - ranks_to_recv.append(rank)

93 -

94 - # Remove those ranks that can get this expert locally.

95 - ranks_to_send_set = set(ranks_to_send)

96 - ranks_to_recv_actual = [

97 - rank for rank in ranks_to_recv if rank not in ranks_to_send_set

98 - ]

97 + # For each expert, extract unique ranks in order of first appearance

98 + for i in range(len(expert_boundaries) - 1):

99 + start, end = expert_boundaries[i], expert_boundaries[i + 1]

100 + expert = int(sorted_experts[start])

101 + expert_ranks = sorted_ranks[start:end]

102 +

103 + # Get unique ranks preserving order

104 + _, unique_idx = np.unique(expert_ranks, return_index=True)

105 + unique_ranks = expert_ranks[np.sort(unique_idx)]

106 + ranks_to_send_map[expert] = unique_ranks.tolist()

107 +

108 + # Process new_indices (recv ranks)

109 + if np.any(new_relevant_mask):

110 + new_relevant_positions = position_indices[new_relevant_mask]

111 + new_relevant_experts = new_indices[new_relevant_mask]

112 + new_relevant_ranks = new_relevant_positions // num_local_experts

113 +

114 + # Sort by expert first, then by position

115 + sort_order = np.lexsort((new_relevant_positions, new_relevant_experts))

116 + sorted_experts = new_relevant_experts[sort_order]

117 + sorted_ranks = new_relevant_ranks[sort_order]

118 +

119 + # Find boundaries where expert changes

120 + expert_boundaries = np.concatenate(

121 + [[0], np.where(np.diff(sorted_experts) != 0)[0] + 1, [len(sorted_experts)]]

122 + )

123

100 - return ranks_to_send, ranks_to_recv_actual

124 + # For each expert, extract unique ranks and exclude local copies

125 + for i in range(len(expert_boundaries) - 1):

126 + start, end = expert_boundaries[i], expert_boundaries[i + 1]

127 + expert = int(sorted_experts[start])

128 + expert_ranks = sorted_ranks[start:end]

129 +

130 + # Get unique ranks preserving order

131 + _, unique_idx = np.unique(expert_ranks, return_index=True)

132 + unique_ranks = expert_ranks[np.sort(unique_idx)]

133 +

134 + # Remove ranks that have local copies (in send map)

135 + send_ranks_set = set(ranks_to_send_map.get(expert, []))

136 + recv_ranks_actual = [

137 + int(r) for r in unique_ranks if r not in send_ranks_set

138 + ]

139 + ranks_to_recv_map[expert] = recv_ranks_actual

140 +

141 + # Handle experts that only appear in old (send only) or new (recv only)

142 + for expert in unique_experts:

143 + expert = int(expert)

144 + if expert not in ranks_to_send_map:

145 + ranks_to_send_map[expert] = []

146 + if expert not in ranks_to_recv_map:

147 + ranks_to_recv_map[expert] = []

148 +

149 + return ranks_to_send_map, ranks_to_recv_map

150

151

152 def move_to_buffer(

153 num_local_experts: int,

105 - old_indices: Sequence[int],

106 - new_indices: Sequence[int],

154 + old_indices: np.ndarray,

155 + new_indices: np.ndarray,

156 expert_weights: Iterable[torch.Tensor],

108 - expert_weights_buffer: Sequence[torch.Tensor],

157 + expert_weights_buffers: Sequence[torch.Tensor],

158 cuda_stream: torch.cuda.Stream | None,

159 ep_group: ProcessGroup,

111 -) -> tuple[list[bool], list[bool], dict[int, int]]:

160 +) -> MoveToBufferResult:

161 """

113 - Perform expert weights rearrangement of one layer.

162 + Rearranges expert weights during EPLB rebalancing.

163 +

164 + Args:

165 + num_local_experts: Number of local experts.

166 + old_indices: (num_experts_total,) ndarray of current (old)

167 + global-to-local expert assignments.

168 + new_indices: (num_experts_total,) ndarray of desired (new)

169 + global-to-local assignments after rebalance.

170 + expert_weights: Original expert weights for the layer.

171 + expert_weights_buffers: Intermediate buffers (one per tensor).

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

173 + ep_group: Distributed process group for expert parallel comms.

174 +

175 + Returns:

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

177 + is unchanged after rebalance.

178 + is_received_locally (np.ndarray): (num_local_experts,), True where a row

179 + can be updated from local data.

180 + RecvMetadata: Metadata needed for completing remote weight transfers.

181 """

182 + assert old_indices.shape == new_indices.shape

183 ep_rank = ep_group.rank()

116 - local2global = partial(

117 - idx_local_to_global,

118 - local_cnt=num_local_experts,

119 - ep_rank=ep_rank,

120 - )

184

122 - # 0. Do nothing for experts that did not change.

123 - is_unchanged = [

124 - old_indices[local2global(i)] == new_indices[local2global(i)]

125 - for i in range(num_local_experts)

126 - ]

185 + recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_)

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

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

188 + recv_expert_ids = np.full((num_local_experts,), -1, dtype=np.int64)

189 + recv_dst_rows = np.full((num_local_experts,), -1, dtype=np.int32)

190

128 - # 1. Perform weight copy inside the local rank.

129 - is_received_locally = is_unchanged[:]

130 - for src in range(num_local_experts):

131 - src_global = local2global(src)

132 - for dst in range(num_local_experts):

133 - dst_global = local2global(dst)

134 - if is_received_locally[dst]:

135 - continue

136 - if old_indices[src_global] == -1 or new_indices[dst_global] == -1:

137 - continue

138 - if old_indices[src_global] == new_indices[dst_global]:

139 - is_received_locally[dst] = True

140 - for weight, buffer in zip(expert_weights, expert_weights_buffer):

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

142 - buffer[dst].copy_(weight[src], non_blocking=True)

191 + base = ep_rank * num_local_experts

192 + local_rows = np.arange(num_local_experts, dtype=np.int32)

193 + local_global = base + local_rows

194 +

195 + old_local_expert_ids = old_indices[local_global]

196 + new_local_expert_ids = new_indices[local_global]

197 +

198 + # Unchanged mask

199 + is_unchanged = old_local_expert_ids == new_local_expert_ids

200 +

201 + # Local receive eligibility

202 + new_valid = new_local_expert_ids != -1

203 + can_recv_local = np.isin(

204 + new_local_expert_ids, old_local_expert_ids, assume_unique=False

205 + )

206 + is_received_locally = np.logical_or(

207 + is_unchanged, np.logical_and(new_valid, can_recv_local)

208 + )

209 +

210 + # Send map: first src row per unique expert present locally in old mapping

211 + send_count = 0

212 + valid_old = old_local_expert_ids != -1

213 + if np.any(valid_old):

214 + uniq_experts, first_idx = np.unique(

215 + old_local_expert_ids[valid_old], return_index=True

216 + )

217 + filtered_rows = local_rows[valid_old]

218 + src_rows = filtered_rows[first_idx]

219 + send_count = int(uniq_experts.shape[0])

220 + send_expert_ids[:send_count] = uniq_experts

221 + send_src_rows[:send_count] = src_rows

222 +

223 + # Recv map: primary dst per unique expert needed remotely

224 + recv_count = 0

225 + need_recv_mask = np.logical_and(~is_received_locally, new_valid)

226 + if np.any(need_recv_mask):

227 + desired_experts = new_local_expert_ids[need_recv_mask]

228 + desired_dsts = local_rows[need_recv_mask]

229 + uniq_recv_experts, uniq_indices = np.unique(desired_experts, return_index=True)

230 + dst_rows = desired_dsts[uniq_indices]

231 + recv_count = int(uniq_recv_experts.shape[0])

232 + recv_expert_ids[:recv_count] = uniq_recv_experts

233 + recv_dst_rows[:recv_count] = dst_rows

234 + recv_primary_mask[dst_rows] = True

235 +

236 + eligible_local_buffer_mask = np.logical_and(~is_unchanged, is_received_locally)

237 +

238 + # 1. Local moves into tmp buffers

239 + if bool(eligible_local_buffer_mask.any()) and send_count > 0:

240 + dest_indices = np.nonzero(eligible_local_buffer_mask)[0].tolist()

241 + expert_to_src_map = dict(

242 + zip(send_expert_ids[:send_count], send_src_rows[:send_count])

243 + )

244 + for dst in dest_indices:

245 + expert = new_local_expert_ids[dst]

246 + src_local = expert_to_src_map.get(expert, -1)

247 + if src_local != -1:

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

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

250

251 p2p_ops: list[P2POp] = []

252

146 - # 2. Initiate sending of weights.

147 - experts_send_loc: dict[int, int] = {}

148 - for src in range(num_local_experts):

149 - expert = old_indices[local2global(src)]

150 - if expert == -1:

151 - continue

152 - if expert in experts_send_loc:

153 - continue

154 - experts_send_loc[expert] = src

155 -

156 - # We need to sort here to match send/recv

157 - for expert, src in sorted(experts_send_loc.items()):

158 - ranks_to_send, ranks_to_recv = get_ep_ranks_with_expert(

159 - expert,

253 + # Pre-compute global ranks mapping

254 + ep_size = ep_group.size()

255 + rank_to_global = {rank: get_global_rank(ep_group, rank) for rank in range(ep_size)}

256 +

257 + # 2. Post sends

258 + if send_count > 0:

259 + experts = send_expert_ids[:send_count]

260 + srcs = send_src_rows[:send_count]

261 + order = np.argsort(experts, kind="stable")

262 + experts = experts[order]

263 + srcs = srcs[order]

264 +

265 + send_map, recv_map = get_ep_ranks_with_experts_batch(

266 + experts,

267 num_local_experts,

268 old_indices,

269 new_indices,

270 )

271

165 - # Calculate the ranks to send by this rank

166 - num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)

167 - sender_pos = ranks_to_send.index(ep_rank)

168 - recv_begin = sender_pos * num_dst_per_sender

169 - recv_end = recv_begin + num_dst_per_sender

170 - recv_ranks = ranks_to_recv[recv_begin:recv_end]

171 -

172 - # Tackle remainders

173 - remainder_start = len(ranks_to_send) * num_dst_per_sender

174 - recver_pos = remainder_start + sender_pos

175 - if recver_pos < len(ranks_to_recv):

176 - recv_ranks.append(ranks_to_recv[recver_pos])

177 -

178 - for dst in recv_ranks:

179 - dst_global = get_global_rank(ep_group, dst)

180 - p2p_ops += [

181 - P2POp(

182 - torch.distributed.isend,

183 - weight[src],

184 - dst_global,

185 - )

186 - for weight in expert_weights

187 - ]

188 -

189 - # 3. Initiate receiving of weights.

190 - experts_recv_loc: dict[int, int] = {}

191 - for dst in range(num_local_experts):

192 - if is_received_locally[dst]:

193 - continue

194 - expert = new_indices[local2global(dst)]

195 - if expert == -1:

196 - continue

197 - if expert in experts_recv_loc:

198 - continue

199 - experts_recv_loc[expert] = dst

200 -

201 - # We need to sort here to match send/recv

202 - for expert, dst in sorted(experts_recv_loc.items()):

203 - ranks_to_send, ranks_to_recv = get_ep_ranks_with_expert(

204 - expert,

272 + for expert, src in zip(experts.tolist(), srcs.tolist()):

273 + ranks_to_send = send_map[expert]

274 + ranks_to_recv = recv_map[expert]

275 + if not ranks_to_send or not ranks_to_recv:

276 + continue

277 + num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)

278 + sender_pos = ranks_to_send.index(ep_rank)

279 + recv_begin = sender_pos * num_dst_per_sender

280 + recv_end = recv_begin + num_dst_per_sender

281 + recv_ranks = ranks_to_recv[recv_begin:recv_end]

282 + remainder_start = len(ranks_to_send) * num_dst_per_sender

283 + recver_pos = remainder_start + sender_pos

284 + if recver_pos < len(ranks_to_recv):

285 + recv_ranks.append(ranks_to_recv[recver_pos])

286 + for dst in recv_ranks:

287 + dst_global = rank_to_global[dst]

288 + p2p_ops += [

289 + P2POp(

290 + torch.distributed.isend,

291 + w[src],

292 + dst_global,

293 + )

294 + for w in expert_weights

295 + ]

296 +

297 + # 3. Post recvs

298 + if recv_count > 0:

299 + experts = recv_expert_ids[:recv_count]

300 + dsts = recv_dst_rows[:recv_count]

301 + order = np.argsort(experts, kind="stable")

302 + experts = experts[order]

303 + dsts = dsts[order]

304 +

305 + send_map, recv_map = get_ep_ranks_with_experts_batch(

306 + experts,

307 num_local_experts,

308 old_indices,

309 new_indices,

310 )

311

210 - # Calculate the rank to recv by this rank

211 - num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)

212 - recver_pos = ranks_to_recv.index(ep_rank)

213 - remainder_start = len(ranks_to_send) * num_dst_per_sender

214 - if recver_pos < remainder_start:

215 - src = ranks_to_send[recver_pos // num_dst_per_sender]

216 - else:

217 - src = ranks_to_send[recver_pos - remainder_start]

218 -

219 - src_global = get_global_rank(ep_group, src)

220 - p2p_ops += [

221 - P2POp(

222 - torch.distributed.irecv,

223 - weight[dst],

224 - src_global,

225 - )

226 - for weight in expert_weights_buffer

227 - ]

312 + for expert, dst in zip(experts.tolist(), dsts.tolist()):

313 + ranks_to_send = send_map[expert]

314 + ranks_to_recv = recv_map[expert]

315 + if not ranks_to_send or not ranks_to_recv:

316 + continue

317 + num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)

318 + recver_pos = ranks_to_recv.index(ep_rank)

319 + remainder_start = len(ranks_to_send) * num_dst_per_sender

320 + if recver_pos < remainder_start:

321 + src = ranks_to_send[recver_pos // num_dst_per_sender]

322 + else:

323 + src = ranks_to_send[recver_pos - remainder_start]

324 + src_global = rank_to_global[src]

325 + p2p_ops += [

326 + P2POp(

327 + torch.distributed.irecv,

328 + b[dst],

329 + src_global,

330 + )

331 + for b in expert_weights_buffers

332 + ]

333

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

335 if p2p_ops and cuda_stream is not None:

@@ -237,38 +342,95 @@ def move_to_buffer(

342 for req in reqs:

343 req.wait()

344 # wait for the communication to finish

240 - return is_unchanged, is_received_locally, experts_recv_loc

345 + return (

346 + is_unchanged,

347 + is_received_locally,

348 + RecvMetadata(

349 + recv_primary_mask=recv_primary_mask,

350 + recv_count=recv_count,

351 + recv_expert_ids=recv_expert_ids,

352 + recv_dst_rows=recv_dst_rows,

353 + ),

354 + )

355

356

357 def move_from_buffer(

358 expert_weights: Iterable[torch.Tensor],

245 - expert_weights_buffer: list[torch.Tensor],

246 - is_unchanged: list[bool],

247 - is_received_locally: list[bool],

248 - experts_recv_loc: dict[int, int],

249 - new_indices: Sequence[int],

250 - ep_group: ProcessGroup,

359 + expert_weights_buffers: list[torch.Tensor],

360 + is_unchanged: np.ndarray,

361 + is_received_locally: np.ndarray,

362 + recv_metadata: RecvMetadata,

363 + new_indices: np.ndarray,

364 + ep_rank: int,

365 ) -> None:

252 - ep_rank = ep_group.rank()

253 - num_local_experts = len(is_unchanged)

366 + """

367 + Copies expert weights from communication buffers back to the target weight tensors

368 + after EPLB rebalancing.

369 +

370 + Args:

371 + expert_weights: List of the actual MoE layer weights used in the execution.

372 + expert_weights_buffers: Intermediate buffers containing the experts weights

373 + after the transfer is completed.

374 + is_unchanged: (num_local_experts,), True where an expert row is unchanged.

375 + is_received_locally: (num_local_experts,), True where a row is updated locally.

376 + recv_metadata: RecvMetadata containing remote receive metadata.

377 + new_indices: (num_experts_total,) mapping from local rows to desired

378 + (possibly global) expert id, after rebalance.

379 + ep_rank: Rank of the process in the expert parallel group.

380 + """

381 + recv_primary_mask = recv_metadata.recv_primary_mask

382 + recv_count = recv_metadata.recv_count

383 + recv_expert_ids = recv_metadata.recv_expert_ids

384 + recv_dst_rows = recv_metadata.recv_dst_rows

385 + num_local_experts = is_unchanged.shape[0]

386 +

387 + # Mask for rows to copy back from buffers:

388 + # copy if locally received OR remote primary recv

389 + copy_mask = np.logical_or(is_received_locally, recv_primary_mask)

390 + dest_mask_np = np.logical_and(~is_unchanged, copy_mask)

391 + if bool(dest_mask_np.any()):

392 + dest_indices = np.nonzero(dest_mask_np)[0].tolist()

393 + for dst in dest_indices:

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

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

396 +

397 + if recv_count == 0:

398 + return

399

255 - local2global = partial(

256 - idx_local_to_global, local_cnt=num_local_experts, ep_rank=ep_rank

400 + # Duplicate remote received rows to non-primary duplicate dsts

401 + base = ep_rank * num_local_experts

402 + local_experts = new_indices[base + np.arange(num_local_experts, dtype=np.int32)]

403 + duplicate_mask = np.logical_and(

404 + np.logical_and(~is_unchanged, ~is_received_locally),

405 + np.logical_and(~recv_primary_mask, local_experts != -1),

406 )

407 + # All received experts are unique in the destination, so no need to copy duplicates

408 + if not bool(duplicate_mask.any()):

409 + return

410

259 - for dst in range(num_local_experts):

260 - if is_unchanged[dst]:

261 - continue

262 - if is_received_locally[dst]:

263 - for weight, buffer in zip(expert_weights, expert_weights_buffer):

264 - weight[dst].copy_(buffer[dst], non_blocking=True)

265 - else:

266 - expert = new_indices[local2global(dst)]

267 - if expert == -1:

268 - continue

269 - src = experts_recv_loc[expert]

270 - for weight, buffer in zip(expert_weights, expert_weights_buffer):

271 - weight[dst].copy_(buffer[src], non_blocking=True)

411 + dup_dst_rows = np.nonzero(duplicate_mask)[0]

412 + dup_experts = local_experts[dup_dst_rows]

413 +

414 + prim_experts = recv_expert_ids[:recv_count]

415 + prim_dsts = recv_dst_rows[:recv_count]

416 + order = np.argsort(prim_experts, kind="stable")

417 + prim_experts_sorted = prim_experts[order]

418 + prim_dsts_sorted = prim_dsts[order]

419 + pos = np.searchsorted(prim_experts_sorted, dup_experts)

420 + valid = np.logical_and(

421 + pos < prim_experts_sorted.shape[0],

422 + prim_experts_sorted[np.minimum(pos, prim_experts_sorted.shape[0] - 1)]

423 + == dup_experts,

424 + )

425 + if not bool(valid.any()):

426 + return

427 +

428 + matched_dst_rows = dup_dst_rows[valid]

429 + matched_src_rows = prim_dsts_sorted[pos[valid]]

430 +

431 + for dst, src in zip(matched_dst_rows.tolist(), matched_src_rows.tolist()):

432 + for w in expert_weights:

433 + w[dst].copy_(w[src], non_blocking=True)

434

435

436 async def transfer_layer(

@@ -281,7 +443,7 @@ async def transfer_layer(

443 layer: int = 0,

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

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

284 -) -> tuple[list[bool], list[bool], dict[int, int]]:

446 +) -> MoveToBufferResult:

447 """

448 Rearranges the expert weights in place according to the new expert indices.

449

@@ -299,6 +461,13 @@ async def transfer_layer(

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

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

463 communications to reserve enough memory for the buffers.

464 +

465 + Returns:

466 + is_unchanged (np.ndarray): (1, num_local_experts), True where expert

467 + is left unchanged.

468 + is_received_locally (np.ndarray): (1, num_local_experts), True where expert

469 + can be received locally.

470 + RecvMetadata: Metadata needed for completing remote weight transfers.

471 """

472 ep_size = ep_group.size()

473 if rank_mapping is not None:

@@ -323,16 +492,19 @@ async def transfer_layer(

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

493 assert num_physical_experts == ep_size * num_local_physical_experts

494

326 - is_unchanged, is_received_locally, experts_recv_loc = move_to_buffer(

495 + old_global_expert_indices_np = old_global_expert_indices.cpu().numpy()

496 + new_global_expert_indices_np = new_global_expert_indices.cpu().numpy()

497 +

498 + is_unchanged, is_received_locally, recv_metadata = move_to_buffer(

499 num_local_experts=num_local_physical_experts,

328 - old_indices=old_global_expert_indices[layer].tolist(),

329 - new_indices=new_global_expert_indices[layer].tolist(),

500 + old_indices=old_global_expert_indices_np[layer],

501 + new_indices=new_global_expert_indices_np[layer],

502 expert_weights=expert_weights[layer],

331 - expert_weights_buffer=expert_weights_buffer,

503 + expert_weights_buffers=expert_weights_buffer,

504 cuda_stream=cuda_stream,

505 ep_group=ep_group,

506 )

335 - return is_unchanged, is_received_locally, experts_recv_loc

507 + return is_unchanged, is_received_locally, recv_metadata

508

509

510 def rearrange_expert_weights_inplace(

@@ -388,19 +560,17 @@ def rearrange_expert_weights_inplace(

560 ep_size = ep_group.size()

561 assert num_physical_experts == ep_size * num_local_physical_experts

562

391 - # A buffer to hold the expert weights in one layer during the exchange.

563 + first_layer_weights = list(expert_weights[0])

564 + # Buffers to hold the expert weights during the exchange.

565 # NOTE: Currently we assume the same weights across different layers

566 # have the same shape.

394 - expert_weights_buffer = [torch.empty_like(w) for w in expert_weights[0]]

395 -

567 + weights_buffer: list[torch.Tensor] = [

568 + torch.empty_like(w) for w in first_layer_weights

569 + ]

570 if is_profile:

397 - # Maximum send size is to send all local experts to all ranks,

398 - # So we use a dummy `all_gather` to reserve enough communication buffer

399 - for weight, buffer in zip(expert_weights[0], expert_weights_buffer):

400 - # A `/dev/null`-like buffer to avoid real memory allocation

571 + # Reserve communication buffers via a minimal dummy all_gather on first layer

572 + for weight, buffer in zip(expert_weights[0], weights_buffer):

573 dummy_recv_buffer = [buffer for _ in range(ep_size)]

402 - # NOTE(bowen): Needed this barrier to avoid OOM during actual

403 - # execution. I'm not very sure why this is needed

574 torch.distributed.barrier()

575 all_gather(

576 dummy_recv_buffer,

@@ -409,32 +579,32 @@ def rearrange_expert_weights_inplace(

579 )

580 return

581

412 - old_global_expert_indices_cpu = old_global_expert_indices.cpu()

413 - new_global_expert_indices_cpu = new_global_expert_indices.cpu()

414 -

582 # NOTE(bowen): We need this synchronize to run, but I don't know why.

583 # If you figure out the reason, please let me know -- thank you!

584 torch.cuda.synchronize()

585

419 - for layer in range(num_moe_layers):

420 - is_unchanged, is_received_locally, experts_recv_loc = move_to_buffer(

586 + old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()

587 + new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()

588 +

589 + for layer_idx in range(num_moe_layers):

590 + is_unchanged, is_received_locally, recv_metadata = move_to_buffer(

591 num_local_experts=num_local_physical_experts,

422 - old_indices=old_global_expert_indices_cpu[layer].tolist(),

423 - new_indices=new_global_expert_indices_cpu[layer].tolist(),

424 - expert_weights=expert_weights[layer],

425 - expert_weights_buffer=expert_weights_buffer,

592 + old_indices=old_global_expert_indices_cpu[layer_idx],

593 + new_indices=new_global_expert_indices_cpu[layer_idx],

594 + expert_weights=expert_weights[layer_idx],

595 + expert_weights_buffers=weights_buffer,

596 cuda_stream=None,

597 ep_group=ep_group,

598 )

599

600 move_from_buffer(

431 - expert_weights=expert_weights[layer],

432 - expert_weights_buffer=expert_weights_buffer,

601 + expert_weights=expert_weights[layer_idx],

602 + expert_weights_buffers=weights_buffer,

603 is_unchanged=is_unchanged,

604 is_received_locally=is_received_locally,

435 - experts_recv_loc=experts_recv_loc,

436 - new_indices=new_global_expert_indices[layer].tolist(),

437 - ep_group=ep_group,

605 + recv_metadata=recv_metadata,

606 + new_indices=new_global_expert_indices_cpu[layer_idx],

607 + ep_rank=ep_group.rank(),

608 )

609

610

@@ -526,4 +696,4 @@ def _map_new_expert_indices_with_rank_mapping(

696 return mapped_expert_indices

697

698

529 -__all__ = ["transfer_layer", "move_from_buffer"]

699 +__all__ = ["transfer_layer", "move_from_buffer", "RecvMetadata"]