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 8, 12:30 UTC
[Perf] Move eplb rebalance algo to async thread (#30888)
Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Ilya Markov committed
Feb 10, 2026 at 23:19 UTC
67132945bbad23233fd583e6106ebebe859c8366
5 files changed
+251
-103
tests/distributed/test_eplb_execute.py
+3
-4
@@ -295,12 +295,11 @@ def _test_async_transfer_layer_without_mtp_worker(
295
for layer_idx in range(num_layers):
296
is_unchanged, is_received_locally, recv_metadata = asyncio.run(
297
transfer_layer(
298
- old_global_expert_indices=old_indices_cpu,
299
- new_global_expert_indices=new_indices_cpu,
300
- expert_weights=expert_weights,
298
+ old_layer_indices=old_indices_cpu[layer_idx],
299
+ new_layer_indices=new_indices_cpu[layer_idx],
300
+ expert_weights=expert_weights[layer_idx],
301
expert_weights_buffer=expert_buffer,
302
ep_group=ep_group,
303
- layer=layer_idx,
303
cuda_stream=cuda_stream,
304
)
305
)
vllm/distributed/eplb/async_worker.py
+89
-18
@@ -11,13 +11,13 @@ from typing import TYPE_CHECKING
11
import torch
12
from torch.distributed import ProcessGroup
13
14
-from vllm.distributed.parallel_state import get_ep_group
14
+from vllm.distributed.parallel_state import get_eplb_group
15
from vllm.logger import init_logger
16
17
from .rebalance_execute import transfer_layer
18
19
if TYPE_CHECKING:
20
- from .eplb_state import EplbState
20
+ from .eplb_state import EplbModelState, EplbState
21
22
logger = init_logger(__name__)
23
@@ -27,8 +27,8 @@ def start_async_worker(
27
rank_mapping: dict[int, int] | None = None,
28
is_profile: bool = False,
29
) -> threading.Thread:
30
- ep_group = get_ep_group().device_group
31
- rank = ep_group.rank()
30
+ eplb_group = get_eplb_group().device_group
31
+ rank = eplb_group.rank()
32
device_index = state.cuda_device_index
33
assert state.is_async
34
@@ -42,7 +42,7 @@ def start_async_worker(
42
loop.run_until_complete(
43
transfer_run_periodically(
44
state=state,
45
- ep_group=ep_group,
45
+ eplb_group=eplb_group,
46
cuda_stream=cuda_stream,
47
is_profile=is_profile,
48
rank_mapping=rank_mapping,
@@ -58,9 +58,53 @@ def start_async_worker(
58
return thread
59
60
61
+def run_rebalance_experts(
62
+ model_state: "EplbModelState",
63
+ eplb_state: "EplbState",
64
+ physical_to_logical_map_cpu: torch.Tensor,
65
+) -> None:
66
+ assert model_state.eplb_stats is not None
67
+ eplb_stats = model_state.eplb_stats
68
+
69
+ # Wait for the main thread's all-reduce and clone to complete before
70
+ # accessing the global_expert_load_window tensor.
71
+ assert model_state.window_ready_event is not None
72
+ model_state.window_ready_event.wait()
73
+ model_state.window_ready_event = None
74
+
75
+ # Move the global expert load window to CPU for computation.
76
+ global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
77
+ # Compute new expert mappings for the model
78
+ (
79
+ new_physical_to_logical_map,
80
+ new_logical_to_physical_map,
81
+ new_logical_replica_count,
82
+ ) = eplb_state.policy.rebalance_experts(
83
+ global_expert_load_window,
84
+ eplb_stats.num_replicas,
85
+ eplb_stats.num_groups,
86
+ eplb_stats.num_nodes,
87
+ eplb_stats.num_gpus,
88
+ physical_to_logical_map_cpu,
89
+ )
90
+ assert new_physical_to_logical_map.device == torch.device("cpu")
91
+
92
+ model_state.new_physical_to_logical_map = new_physical_to_logical_map
93
+
94
+ max_slots = model_state.logical_to_physical_map.shape[-1]
95
+ padded_logical = torch.nn.functional.pad(
96
+ new_logical_to_physical_map,
97
+ (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
98
+ value=-1,
99
+ ).to(model_state.logical_to_physical_map.device)
100
+ new_replica = new_logical_replica_count.to(model_state.logical_replica_count.device)
101
+ model_state.new_logical_to_physical_map = padded_logical
102
+ model_state.new_logical_replica_count = new_replica
103
+
104
+
105
async def transfer_run_periodically(
106
state: "EplbState",
63
- ep_group: ProcessGroup,
107
+ eplb_group: ProcessGroup,
108
cuda_stream: torch.cuda.Stream,
109
is_profile: bool = False,
110
rank_mapping: dict[int, int] | None = None,
@@ -71,23 +115,51 @@ async def transfer_run_periodically(
115
116
assert state.is_async
117
for model_state in state.model_states.values():
118
+ rebalancing_algorithm_executed = False
119
+ physical_to_logical_map_cpu = None
120
current_num_layers = model_state.model.num_moe_layers
121
while (
122
model_state.rebalanced
123
and model_state.layer_to_transfer < current_num_layers
124
):
79
- if (
80
- not model_state.ep_buffer_ready
81
- and model_state.rebalanced
82
- and model_state.new_physical_to_logical_map is not None
83
- ):
84
- await asyncio.to_thread(model_state.buffer_lock.acquire)
125
+ if not model_state.ep_buffer_ready and model_state.rebalanced:
126
+ # Polling the lock directly in the async thread avoids
127
+ # the thread switch overhead of asyncio.to_thread.
128
+ # This is typically faster than offloading to a worker thread.
129
+ while not model_state.buffer_lock.acquire(blocking=False):
130
+ await asyncio.sleep(0)
131
try:
132
if model_state.layer_to_transfer >= current_num_layers:
133
break
134
+ if (
135
+ not rebalancing_algorithm_executed
136
+ or model_state.new_physical_to_logical_map is None
137
+ ):
138
+ # Move the physical_to_logical_map to CPU
139
+ # for rebalancing and transfer_layer.
140
+ physical_to_logical_map_cpu = (
141
+ model_state.physical_to_logical_map.cpu()
142
+ )
143
+ run_rebalance_experts(
144
+ model_state, state, physical_to_logical_map_cpu
145
+ )
146
+ rebalancing_algorithm_executed = True
147
+ logger.info(
148
+ "Async worker computed new indices for model %s",
149
+ model_state.model_name,
150
+ )
151
+
152
+ assert model_state.new_physical_to_logical_map is not None
153
+ assert physical_to_logical_map_cpu is not None
154
+
155
+ layer_idx = model_state.layer_to_transfer
156
+ old_layer_indices = physical_to_logical_map_cpu[layer_idx]
157
+ new_layer_indices = model_state.new_physical_to_logical_map[
158
+ layer_idx
159
+ ]
160
161
# Wait for the main thread to finish consuming the buffer
90
- # before overwriting it
162
+ # before initiating an EPLB transfer on another layer.
163
if model_state.buffer_consumed_event is not None:
164
cuda_stream.wait_event(model_state.buffer_consumed_event)
165
model_state.buffer_consumed_event = None
@@ -97,13 +169,12 @@ async def transfer_run_periodically(
169
model_state.is_received_locally,
170
model_state.recv_metadata,
171
) = await transfer_layer(
100
- old_global_expert_indices=model_state.physical_to_logical_map,
101
- new_global_expert_indices=model_state.new_physical_to_logical_map,
102
- expert_weights=model_state.model.expert_weights,
172
+ old_layer_indices=old_layer_indices,
173
+ new_layer_indices=new_layer_indices,
174
+ expert_weights=model_state.model.expert_weights[layer_idx],
175
expert_weights_buffer=model_state.expert_buffer,
104
- ep_group=ep_group,
176
+ ep_group=eplb_group,
177
is_profile=is_profile,
106
- layer=model_state.layer_to_transfer,
178
cuda_stream=cuda_stream,
179
rank_mapping=rank_mapping,
180
)
vllm/distributed/eplb/eplb_state.py
+88
-54
@@ -55,6 +55,35 @@ from .rebalance_execute import (
55
logger = init_logger(__name__)
56
57
58
+@dataclass
59
+class EplbStats:
60
+ """
61
+ Model stats used in EPLB rebalancing algorithm.
62
+ """
63
+
64
+ global_expert_load_window: torch.Tensor
65
+ """
66
+ Experts load window.
67
+ Shape: (window_size, num_moe_layers, num_physical_experts)
68
+ """
69
+ num_replicas: int
70
+ """
71
+ Number of physical experts.
72
+ """
73
+ num_groups: int
74
+ """
75
+ Number of expert groups.
76
+ """
77
+ num_nodes: int
78
+ """
79
+ Number of nodes.
80
+ """
81
+ num_gpus: int
82
+ """
83
+ Number of GPUs.
84
+ """
85
+
86
+
87
@dataclass
88
class EplbModelState:
89
"""EPLB metrics."""
@@ -156,6 +185,11 @@ class EplbModelState:
185
CUDA event recorded after the main thread finishes consuming the buffer.
186
The async worker waits on this before writing to the buffer again.
187
"""
188
+ window_ready_event: torch.cuda.Event | None
189
+ """
190
+ CUDA event recorded after all-reduce and clone on the main thread.
191
+ The async worker waits on this before accessing global_expert_load_window.
192
+ """
193
ep_buffer_ready: int
194
"""
195
The flag indicates whether the expert buffer is ready for transfer.
@@ -173,6 +207,10 @@ class EplbModelState:
207
"""
208
Whether the async EPLB needs to poll peers for buffer readiness.
209
"""
210
+ eplb_stats: EplbStats | None
211
+ """
212
+ EPLB stats for the model.
213
+ """
214
is_unchanged: np.ndarray
215
"""
216
intermediate variable between `move_to_buffer` and `move_to_workspace`.
@@ -508,10 +546,12 @@ class EplbState:
546
buffer_lock=threading.Lock(),
547
buffer_ready_event=None,
548
buffer_consumed_event=None,
549
+ window_ready_event=None,
550
ep_buffer_ready=0,
551
layer_to_transfer=0,
552
rebalanced=False,
553
pending_global_ready_check=False,
554
+ eplb_stats=None,
555
is_unchanged=np.array([]),
556
is_received_locally=np.array([]),
557
recv_metadata=RecvMetadata(
@@ -642,20 +682,6 @@ class EplbState:
682
ep_group=ep_group,
683
is_profile=is_profile,
684
)
645
- if (
646
- eplb_model_state.layer_to_transfer
647
- >= eplb_model_state.model.num_moe_layers
648
- ):
649
- self.post_eplb(eplb_model_state, is_profile)
650
- eplb_model_state.rebalanced = False
651
- eplb_model_state.layer_to_transfer = 0
652
- eplb_model_state.pending_global_ready_check = False
653
- logger.info(
654
- "finish async transfer for model %s rank %d layer %d",
655
- eplb_model_state.model_name,
656
- ep_group.rank(),
657
- eplb_model_state.model.num_moe_layers,
658
- )
685
686
if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval:
687
if self.is_async and any(
@@ -802,21 +828,21 @@ class EplbState:
828
for eplb_model_state, global_expert_load_window in zip(
829
self.model_states.values(), global_expert_load_windows
830
):
805
- # Get new expert mappings for the model
806
- (
807
- new_physical_to_logical_map,
808
- new_logical_to_physical_map,
809
- new_logical_replica_count,
810
- ) = self.policy.rebalance_experts(
811
- global_expert_load_window,
812
- num_replicas,
813
- num_groups,
814
- num_nodes,
815
- num_gpus,
816
- eplb_model_state.physical_to_logical_map,
817
- )
818
-
831
if not self.is_async or is_profile:
832
+ # Get new expert mappings for the model
833
+ (
834
+ new_physical_to_logical_map,
835
+ new_logical_to_physical_map,
836
+ new_logical_replica_count,
837
+ ) = self.policy.rebalance_experts(
838
+ global_expert_load_window,
839
+ num_replicas,
840
+ num_groups,
841
+ num_nodes,
842
+ num_gpus,
843
+ eplb_model_state.physical_to_logical_map,
844
+ )
845
+
846
# Update expert weights
847
rearrange_expert_weights_inplace(
848
eplb_model_state.physical_to_logical_map,
@@ -873,27 +899,25 @@ class EplbState:
899
gpu_elapsed,
900
)
901
else:
876
- max_slots = eplb_model_state.logical_to_physical_map.shape[-1]
877
- padded_logical = torch.nn.functional.pad(
878
- new_logical_to_physical_map,
879
- (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
880
- value=-1,
881
- ).to(eplb_model_state.logical_to_physical_map.device)
882
- new_replica = new_logical_replica_count.to(
883
- eplb_model_state.logical_replica_count.device
902
+ eplb_model_state.eplb_stats = EplbStats(
903
+ # We copy the tensor to snapshot the global_expert_load_window
904
+ # on the main thread so that async worker can access it safely
905
+ # while the main thread is running.
906
+ global_expert_load_window=global_expert_load_window.clone(),
907
+ num_replicas=num_replicas,
908
+ num_groups=num_groups,
909
+ num_nodes=num_nodes,
910
+ num_gpus=num_gpus,
911
)
885
-
886
- # Move map to cpu in advance
887
- eplb_model_state.new_physical_to_logical_map = (
888
- new_physical_to_logical_map.cpu()
889
- )
890
- eplb_model_state.new_logical_to_physical_map = padded_logical
891
- eplb_model_state.new_logical_replica_count = new_replica
912
+ # Record event after clone to signal async worker
913
+ # that load stats data is ready
914
+ sync_event = torch.cuda.Event()
915
+ sync_event.record()
916
+ eplb_model_state.window_ready_event = sync_event
917
918
eplb_model_state.rebalanced = True
919
eplb_model_state.layer_to_transfer = 0
920
eplb_model_state.pending_global_ready_check = True
896
-
921
# Signal async thread to start transferring layers
922
if self.is_async and (not is_profile):
923
self.rearrange_event.set()
@@ -925,11 +949,13 @@ class EplbState:
949
950
target_device = model_state.physical_to_logical_map.device
951
new_physical = model_state.new_physical_to_logical_map
952
+ # If the number of physical experts has changed, then the new map needs to
953
+ # be copied synchronously to avoid a race condition with the async worker
954
if model_state.physical_to_logical_map.shape[1] != new_physical.shape[1]:
955
model_state.physical_to_logical_map = new_physical.to(target_device)
956
else:
957
model_state.physical_to_logical_map[layer].copy_(
932
- new_physical[layer].to(target_device)
958
+ new_physical[layer].to(target_device, non_blocking=True)
959
)
960
961
logical_device = model_state.logical_to_physical_map.device
@@ -1004,11 +1030,9 @@ class EplbState:
1030
model_state.layer_to_transfer
1031
]
1032
expert_weights_buffer = model_state.expert_buffer
1007
- new_indices = (
1008
- model_state.new_physical_to_logical_map[model_state.layer_to_transfer]
1009
- .cpu()
1010
- .numpy()
1011
- )
1033
+ new_indices = model_state.new_physical_to_logical_map[
1034
+ model_state.layer_to_transfer
1035
+ ].numpy()
1036
move_from_buffer(
1037
expert_weights=expert_weights,
1038
expert_weights_buffers=expert_weights_buffer,
@@ -1019,7 +1043,7 @@ class EplbState:
1043
ep_rank=ep_group.rank(),
1044
)
1045
# Record event after consuming buffer to signal async thread
1022
- # that it's safe to overwrite the buffer
1046
+ # that it's safe to overwrite the intermediate buffer
1047
consumed_event = torch.cuda.Event()
1048
consumed_event.record()
1049
model_state.buffer_consumed_event = consumed_event
@@ -1034,6 +1058,18 @@ class EplbState:
1058
model_state.model_name,
1059
transferred_layer,
1060
)
1061
+ if model_state.layer_to_transfer >= model_state.model.num_moe_layers:
1062
+ self.post_eplb(model_state, is_profile)
1063
+ model_state.rebalanced = False
1064
+ model_state.layer_to_transfer = 0
1065
+ model_state.pending_global_ready_check = False
1066
+ logger.info(
1067
+ "finish async transfer for model %s rank %d layer %d",
1068
+ model_state.model_name,
1069
+ ep_group.rank(),
1070
+ model_state.model.num_moe_layers,
1071
+ )
1072
+
1073
finally:
1074
try:
1075
model_state.buffer_lock.release()
@@ -1048,9 +1084,7 @@ class EplbState:
1084
assert model_state.new_physical_to_logical_map is not None
1085
assert model_state.new_logical_to_physical_map is not None
1086
assert model_state.new_logical_replica_count is not None
1051
- if not is_profile:
1052
- for layer_idx in range(model_state.physical_to_logical_map.shape[0]):
1053
- self._update_layer_mapping_from_new(model_state, layer_idx)
1087
+
1088
model_state.new_physical_to_logical_map = None
1089
model_state.new_logical_to_physical_map = None
1090
model_state.new_logical_replica_count = None
vllm/distributed/eplb/rebalance_execute.py
+33
-26
@@ -434,13 +434,12 @@ def move_from_buffer(
434
435
436
async def transfer_layer(
437
- old_global_expert_indices: torch.Tensor,
438
- new_global_expert_indices: torch.Tensor,
439
- expert_weights: Sequence[Sequence[torch.Tensor]],
437
+ old_layer_indices: torch.Tensor,
438
+ new_layer_indices: torch.Tensor,
439
+ expert_weights: Sequence[torch.Tensor],
440
expert_weights_buffer: Sequence[torch.Tensor],
441
ep_group: ProcessGroup,
442
is_profile: bool = False,
443
- layer: int = 0,
443
cuda_stream: torch.cuda.Stream | None = None,
444
rank_mapping: dict[int, int] | None = None,
445
) -> MoveToBufferResult:
@@ -451,56 +450,64 @@ async def transfer_layer(
450
while keys are physical.
451
452
Args:
454
- old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
455
- new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
456
- expert_weights: A sequence of shape (num_moe_layers)(weight_count)
457
- of tensors of shape (num_local_physical_experts, hidden_size_i).
458
- For example, a linear layer may have up and down projection,
459
- so weight_count = 2. Each weight's hidden size can be different.
453
+ old_layer_indices: Shape (num_physical_experts,).
454
+ new_layer_indices: Shape (num_physical_experts,).
455
+ expert_weights: Iterable of weight tensors for this layer, each with shape
456
+ (num_local_physical_experts, hidden_size_i).
457
+ For example, a linear layer may have up and down projection.
458
+ expert_weights_buffer: Intermediate buffers (one per weight tensor).
459
ep_group: The device process group for expert parallelism.
460
is_profile (bool): If `True`, do not perform any actual weight copy.
461
This is used during profile run, where we only perform dummy
462
communications to reserve enough memory for the buffers.
463
+ cuda_stream: CUDA stream for async copies (can be None for sync mode).
464
+ rank_mapping: Optional rank mapping for elastic expert parallelism.
465
466
Returns:
466
- is_unchanged (np.ndarray): (1, num_local_experts), True where expert
467
+ is_unchanged (np.ndarray): (num_local_experts,), True where expert
468
is left unchanged.
468
- is_received_locally (np.ndarray): (1, num_local_experts), True where expert
469
+ is_received_locally (np.ndarray): (num_local_experts,), True where expert
470
can be received locally.
471
RecvMetadata: Metadata needed for completing remote weight transfers.
472
"""
473
ep_size = ep_group.size()
474
if rank_mapping is not None:
475
+ # Add a layer dimension for compatibility with mapping functions
476
+ old_layer_indices_2d = old_layer_indices.unsqueeze(0)
477
+ new_layer_indices_2d = new_layer_indices.unsqueeze(0)
478
+
479
if len(rank_mapping) == ep_group.size():
480
# scale down
476
- new_global_expert_indices = _map_new_expert_indices_with_rank_mapping(
477
- new_global_expert_indices,
481
+ new_layer_indices_2d = _map_new_expert_indices_with_rank_mapping(
482
+ new_layer_indices_2d,
483
rank_mapping,
484
)
485
else:
486
# scale up
482
- old_global_expert_indices = _map_old_expert_indices_with_rank_mapping(
483
- old_global_expert_indices,
487
+ old_layer_indices_2d = _map_old_expert_indices_with_rank_mapping(
488
+ old_layer_indices_2d,
489
rank_mapping,
490
ep_group.size(),
491
)
492
488
- assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
489
- num_moe_layers, num_physical_experts = old_global_expert_indices.shape
490
- assert len(expert_weights) == num_moe_layers
493
+ # Remove the layer dimension
494
+ old_layer_indices = old_layer_indices_2d.squeeze(0)
495
+ new_layer_indices = new_layer_indices_2d.squeeze(0)
496
+
497
+ assert old_layer_indices.shape == new_layer_indices.shape
498
+ num_physical_experts = old_layer_indices.shape[0]
499
assert len(expert_weights[0]) >= 1
492
- num_local_physical_experts = expert_weights[0][0].shape[0]
493
- assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
500
+ num_local_physical_experts = expert_weights[0].shape[0]
501
assert num_physical_experts == ep_size * num_local_physical_experts
502
496
- old_global_expert_indices_np = old_global_expert_indices.cpu().numpy()
497
- new_global_expert_indices_np = new_global_expert_indices.cpu().numpy()
503
+ old_layer_indices_np = old_layer_indices.cpu().numpy()
504
+ new_layer_indices_np = new_layer_indices.cpu().numpy()
505
506
is_unchanged, is_received_locally, recv_metadata = move_to_buffer(
507
num_local_experts=num_local_physical_experts,
501
- old_indices=old_global_expert_indices_np[layer],
502
- new_indices=new_global_expert_indices_np[layer],
503
- expert_weights=expert_weights[layer],
508
+ old_indices=old_layer_indices_np,
509
+ new_indices=new_layer_indices_np,
510
+ expert_weights=expert_weights,
511
expert_weights_buffers=expert_weights_buffer,
512
cuda_stream=cuda_stream,
513
ep_group=ep_group,
vllm/distributed/parallel_state.py
+38
-1
@@ -1143,6 +1143,18 @@ def get_ep_group() -> GroupCoordinator:
1143
return _EP
1144
1145
1146
+_EPLB: GroupCoordinator | None = None
1147
+
1148
+
1149
+def get_eplb_group() -> GroupCoordinator:
1150
+ assert _EPLB is not None, (
1151
+ "EPLB group is not initialized. "
1152
+ "EPLB group is only created for MoE models when EPLB is enabled. "
1153
+ "Ensure parallel_config.enable_eplb is True."
1154
+ )
1155
+ return _EPLB
1156
+
1157
+
1158
_PCP: GroupCoordinator | None = None
1159
1160
@@ -1440,12 +1452,29 @@ def initialize_model_parallel(
1452
_EP = init_model_parallel_group(
1453
group_ranks, get_world_group().local_rank, backend, group_name="ep"
1454
)
1455
+
1456
+ # Create EPLB group with the same ranks as EP if EPLB is enabled.
1457
+ # This is a separate process group to isolate EPLB communications
1458
+ # from MoE forward pass collectives and prevent deadlocks when
1459
+ # using torch.distributed in execution with torch.distributed in EPLB.
1460
+ global _EPLB
1461
+ assert _EPLB is None, "EPLB group is already initialized"
1462
+ if (
1463
+ config is not None
1464
+ and config.parallel_config is not None
1465
+ and config.parallel_config.enable_eplb
1466
+ ):
1467
+ # Reuse the same group_ranks from EP
1468
+ _EPLB = init_model_parallel_group(
1469
+ group_ranks, get_world_group().local_rank, backend, group_name="eplb"
1470
+ )
1471
# If no EP group needed, _EP remains None
1472
+ # If no EPLB group needed, _EPLB remains None
1473
1474
logger.info_once(
1475
"rank %s in world size %s is assigned as "
1476
"DP rank %s, PP rank %s, PCP rank %s, "
1448
- "TP rank %s, EP rank %s",
1477
+ "TP rank %s, EP rank %s, EPLB rank %s",
1478
rank,
1479
world_size,
1480
_DP.rank_in_group,
@@ -1453,6 +1482,7 @@ def initialize_model_parallel(
1482
_PCP.rank_in_group,
1483
_TP.rank_in_group,
1484
_EP.rank_in_group if _EP is not None else "N/A",
1485
+ _EPLB.rank_in_group if _EPLB is not None else "N/A",
1486
)
1487
1488
@@ -1514,6 +1544,8 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module):
1544
_DP.prepare_communication_buffer_for_model(model)
1545
if _EP is not None:
1546
_EP.prepare_communication_buffer_for_model(model)
1547
+ if _EPLB is not None:
1548
+ _EPLB.prepare_communication_buffer_for_model(model)
1549
1550
1551
def model_parallel_is_initialized():
@@ -1608,6 +1640,11 @@ def destroy_model_parallel():
1640
_EP.destroy()
1641
_EP = None
1642
1643
+ global _EPLB
1644
+ if _EPLB:
1645
+ _EPLB.destroy()
1646
+ _EPLB = None
1647
+
1648
1649
def destroy_distributed_environment():
1650
global _WORLD, _NODE_COUNT