| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Expert parallelism load balancer (EPLB) metrics and states. |
| 4 | |
| 5 | # Glossary |
| 6 | |
| 7 | - **Logical Expert**: An expert that is part of the model's logical structure. |
| 8 | It holds a set of weights and is replicated across multiple physical |
| 9 | experts. |
| 10 | - **Redundant Expert**: To achieve load balancing, for some popular logical |
| 11 | experts, we create additional copies of the expert weights. During inference, |
| 12 | each of these copies can be routed to by the same set of tokens. |
| 13 | - **Physical Expert**: An expert that is instantiated on a specific device. |
| 14 | It is a replica of a logical expert and can be rearranged across devices. |
| 15 | I.e., one logical expert may have multiple sets of weights initialized on |
| 16 | different devices, and each of these sets is a physical expert. |
| 17 | - **Local Physical Expert**: A physical expert that is instantiated on the |
| 18 | current device. |
| 19 | |
| 20 | For example: DeepSeek-R1 has 256 logical experts, so each MoE layer |
| 21 | has 256 sets of linear layer weights in the model parameters. If we add 32 |
| 22 | redundant experts, DeepSeek-R1 will have 256 + 32 = 288 physical experts in |
| 23 | total. And when deploying, we'll have 288 sets of linear layer weights for each |
| 24 | MoE layer. If we have 32 EP ranks, then each GPU will hold 288 / 32 = 9 local |
| 25 | physical experts. |
| 26 | """ |
| 27 | |
| 28 | import threading |
| 29 | import time |
| 30 | from collections.abc import Sequence |
| 31 | from dataclasses import dataclass |
| 32 | |
| 33 | import torch |
| 34 | from torch.distributed import ProcessGroup, all_reduce |
| 35 | |
| 36 | from vllm.config import ModelConfig, ParallelConfig |
| 37 | from vllm.config.utils import compute_hash_cached |
| 38 | from vllm.distributed.parallel_state import ( |
| 39 | GroupCoordinator, |
| 40 | get_ep_group, |
| 41 | get_eplb_group, |
| 42 | get_node_count, |
| 43 | in_the_same_node_as, |
| 44 | ) |
| 45 | from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator |
| 46 | from vllm.distributed.utils import StatelessProcessGroup |
| 47 | from vllm.logger import init_logger |
| 48 | from vllm.model_executor.models.interfaces import MixtureOfExperts |
| 49 | from vllm.platforms import current_platform |
| 50 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 51 | from vllm.utils.torch_utils import PIN_MEMORY |
| 52 | |
| 53 | from .async_worker import start_async_worker |
| 54 | from .eplb_communicator import EplbCommunicator, create_eplb_communicator |
| 55 | from .eplb_utils import CpuGpuEvent |
| 56 | from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy |
| 57 | from .rebalance_execute import ( |
| 58 | AsyncEplbLayerResult, |
| 59 | move_from_buffer, |
| 60 | rearrange_expert_weights_inplace, |
| 61 | ) |
| 62 | |
| 63 | logger = init_logger(__name__) |
| 64 | |
| 65 | |
| 66 | def _compute_eplb_load_stats( |
| 67 | num_tokens_per_rank: torch.Tensor, |
| 68 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 69 | avg_tokens = num_tokens_per_rank.mean(dim=1).sum() |
| 70 | max_tokens = num_tokens_per_rank.max(dim=1).values.sum() |
| 71 | return avg_tokens, max_tokens |
| 72 | |
| 73 | |
| 74 | @dataclass |
| 75 | class EplbStats: |
| 76 | """Model stats used in EPLB rebalancing algorithm.""" |
| 77 | |
| 78 | global_expert_load_window: torch.Tensor |
| 79 | """ |
| 80 | Experts load window. |
| 81 | Shape: (window_size, num_moe_layers, num_physical_experts) |
| 82 | """ |
| 83 | num_replicas: int |
| 84 | """ |
| 85 | Number of physical experts. |
| 86 | """ |
| 87 | num_groups: int |
| 88 | """ |
| 89 | Number of expert groups. |
| 90 | """ |
| 91 | num_nodes: int |
| 92 | """ |
| 93 | Number of nodes. |
| 94 | """ |
| 95 | num_gpus: int |
| 96 | """ |
| 97 | Number of GPUs. |
| 98 | """ |
| 99 | |
| 100 | |
| 101 | @dataclass |
| 102 | class EplbModelState: |
| 103 | """EPLB metrics.""" |
| 104 | |
| 105 | physical_to_logical_map: torch.Tensor |
| 106 | """ |
| 107 | Mapping from physical experts to logical experts. |
| 108 | |
| 109 | Shape: (num_moe_layers, num_physical_experts) |
| 110 | |
| 111 | # Example |
| 112 | |
| 113 | For a 2-layer MoE model with 6 physical experts and 4 logical experts on 3 |
| 114 | EP ranks, the mapping could look like this: |
| 115 | |
| 116 | ``` |
| 117 | [[0, 1, 2, 3, 0, 1], |
| 118 | [0, 2, 0, 1, 0, 3]] |
| 119 | ``` |
| 120 | """ |
| 121 | physical_to_logical_map_buffer: torch.Tensor |
| 122 | """Maximum-capacity buffer backing ``physical_to_logical_map``.""" |
| 123 | logical_to_physical_map: torch.Tensor |
| 124 | """ |
| 125 | Mapping from logical experts to physical experts. |
| 126 | |
| 127 | This is a sparse matrix, where -1 indicates no mapping. |
| 128 | |
| 129 | Shape: (num_moe_layers, num_logical_experts, num_redundant_experts + 1) |
| 130 | |
| 131 | # Example |
| 132 | |
| 133 | For a 2-layer MoE model with 6 physical experts and 4 logical experts on 3 |
| 134 | EP ranks, the mapping could look like this: |
| 135 | |
| 136 | ``` |
| 137 | [[[0, 4, -1], |
| 138 | [1, 5, -1], |
| 139 | [2, -1, -1], |
| 140 | [3, -1, -1]], |
| 141 | [[0, 2, 4], |
| 142 | [3, -1, -1], |
| 143 | [1, -1, -1], |
| 144 | [5, -1, -1]]] |
| 145 | ``` |
| 146 | """ |
| 147 | logical_replica_count: torch.Tensor |
| 148 | """ |
| 149 | Number of replicas for each logical expert. |
| 150 | This is exactly the non-`-1` count in the `logical_to_physical_map`. |
| 151 | |
| 152 | Shape: (num_moe_layers, num_logical_experts) |
| 153 | |
| 154 | # Example |
| 155 | For a 2-layer MoE model with 6 physical experts and 4 logical experts on 3 |
| 156 | EP ranks, the count could look like this: |
| 157 | |
| 158 | ``` |
| 159 | [[2, 2, 1, 1], |
| 160 | [3, 1, 1, 1]] |
| 161 | """ |
| 162 | |
| 163 | expert_load_pass: torch.Tensor |
| 164 | """ |
| 165 | Expert load during this forward pass. |
| 166 | We use the token count each expert processes as the load. |
| 167 | |
| 168 | Shape: (num_moe_layers, num_physical_experts) |
| 169 | """ |
| 170 | expert_load_pass_buffer: torch.Tensor |
| 171 | """Maximum-capacity buffer backing ``expert_load_pass``.""" |
| 172 | expert_load_window: torch.Tensor |
| 173 | """ |
| 174 | A sliding window of expert load. |
| 175 | |
| 176 | Shape: (window_size, num_moe_layers, num_physical_experts) |
| 177 | |
| 178 | NOTE: The expert_load_view now records load for all physical experts |
| 179 | rather than just local experts. This ensures consistent load statistics |
| 180 | across different dispatch methods (naive all-to-all, DeepEP). |
| 181 | The recorded load will be multiplied by dp_size when using naive all-to-all |
| 182 | due to each DP rank contributing the same token set to the calculation. |
| 183 | See: |
| 184 | https://github.com/vllm-project/vllm/pull/22167#pullrequestreview-3086143856 |
| 185 | """ |
| 186 | last_expert_load: torch.Tensor |
| 187 | """Global logical-expert totals from the last reshuffle; equal demand initially.""" |
| 188 | model_name: str |
| 189 | model: MixtureOfExperts |
| 190 | expert_buffer: list[torch.Tensor] |
| 191 | """ |
| 192 | The buffer to store the expert weights during transfer. |
| 193 | """ |
| 194 | rebalanced: bool |
| 195 | """ |
| 196 | This flag is only used when running Async EPLB. It is set to True by the main thread |
| 197 | after the new expert maps have been computed. This indicates that the async worker |
| 198 | should start transferring weights. move_to_workspace sets this flag to False when |
| 199 | all weights have been transferred and the new map has been successfully committed. |
| 200 | |
| 201 | rebalanced relies on the GIL to synchronize access between the main thread and |
| 202 | the async worker. |
| 203 | """ |
| 204 | eplb_stats: EplbStats | None |
| 205 | """ |
| 206 | EPLB stats for the model. |
| 207 | """ |
| 208 | device_index: int | None |
| 209 | """ |
| 210 | Device index for the async EPLB worker thread. |
| 211 | """ |
| 212 | communicator: EplbCommunicator |
| 213 | """ |
| 214 | The communicator for expert weight transfers. |
| 215 | """ |
| 216 | pending_result: AsyncEplbLayerResult | None = None |
| 217 | """ |
| 218 | Set by the async worker after all writes to expert_buffer are done. Consumed |
| 219 | and reset to None by the main thread in move_to_workspace() after the contents of |
| 220 | expert_buffer have been transferred out. At most one result is pending at a time. |
| 221 | |
| 222 | pending_result relies on the GIL to synchronize access between the main thread and |
| 223 | the async worker. |
| 224 | """ |
| 225 | num_unpadded_tokens_tensors: list[torch.Tensor] | None = None |
| 226 | """ |
| 227 | Per-ubatch scalar int32 tensors holding the number of real (non-padding) |
| 228 | tokens. Allocated once in :meth:`EplbState.add_model` so that device |
| 229 | pointers remain stable across CUDA-graph replays. The router kernel |
| 230 | indexes this list with ``dbo_current_ubatch_id()``. |
| 231 | """ |
| 232 | |
| 233 | |
| 234 | class EplbState: |
| 235 | """EplbState of each expert parallel model. Key is the model config hash.""" |
| 236 | |
| 237 | def __init__(self, parallel_config: ParallelConfig, device: torch.device): |
| 238 | self.parallel_config = parallel_config |
| 239 | self.device = device |
| 240 | self.model_states: dict[str, EplbModelState] = {} |
| 241 | self.policy: type[AbstractEplbPolicy] = DefaultEplbPolicy |
| 242 | """ |
| 243 | Selected EPLB algorithm class |
| 244 | """ |
| 245 | self.expert_load_window_step: int = 0 |
| 246 | """ |
| 247 | Current step in the sliding window. |
| 248 | |
| 249 | Different from `expert_rearrangement_step`, |
| 250 | each EP rank may have its own `expert_load_window_step`. |
| 251 | """ |
| 252 | self.expert_load_window_size: int = 0 |
| 253 | """ |
| 254 | Size of the expert load sliding window. |
| 255 | This is a constant and is taken from the config. |
| 256 | """ |
| 257 | self.expert_rearrangement_step: int = 0 |
| 258 | """ |
| 259 | Steps after last rearrangement. |
| 260 | Will trigger a rearrangement if it exceeds the threshold. |
| 261 | |
| 262 | NOTE: Keep in mind that all EP ranks need to have the same |
| 263 | `expert_rearrangement_step` value to ensure synchronization. |
| 264 | Otherwise, the rearrangement will hang at collective |
| 265 | communication calls. |
| 266 | """ |
| 267 | self.expert_rearrangement_step_interval: int = 0 |
| 268 | """ |
| 269 | Interval for expert rearrangement steps. |
| 270 | This is a constant and is taken from the config. |
| 271 | """ |
| 272 | self.should_record_tensor: torch.Tensor | None = None |
| 273 | """ |
| 274 | Shared scalar bool tensor for all layers. Every |
| 275 | :class:`EplbLayerState` holds a reference to the **same** object so |
| 276 | a single ``.fill_()`` updates all layers at once. Allocated on the |
| 277 | first call to :meth:`_propagate_shared_tensors`. |
| 278 | """ |
| 279 | self.is_async: bool = False |
| 280 | """ |
| 281 | The flag indicates whether the EPLB is running in async mode. |
| 282 | """ |
| 283 | self.rearrange_event: CpuGpuEvent = CpuGpuEvent() |
| 284 | """ |
| 285 | Event to signal when a new rearrangement is needed for the async thread. |
| 286 | """ |
| 287 | self.async_worker: threading.Thread | None = None |
| 288 | """ |
| 289 | Background thread handling async transfers. |
| 290 | """ |
| 291 | self.device_index: int | None = None |
| 292 | """ |
| 293 | Device index for the async EPLB worker thread. |
| 294 | """ |
| 295 | if self.device.type in ("cuda", "xpu"): |
| 296 | self.device_index = self.device.index |
| 297 | if self.device_index is None and torch.accelerator.is_available(): |
| 298 | self.device_index = torch.accelerator.current_device_index() |
| 299 | |
| 300 | @staticmethod |
| 301 | def build_initial_global_physical_to_logical_map( |
| 302 | num_routed_experts: int, |
| 303 | num_redundant_experts: int, |
| 304 | ) -> Sequence[int]: |
| 305 | """Build an initial expert arrangement using the following structure: |
| 306 | [original routed experts, redundant experts] |
| 307 | |
| 308 | Returns: |
| 309 | physical_to_logical_map (Sequence[int]): A list of integers, |
| 310 | where each integer is the index of the logical expert |
| 311 | that the corresponding physical expert maps to. |
| 312 | |
| 313 | """ |
| 314 | global_physical_to_logical_map = list(range(num_routed_experts)) |
| 315 | global_physical_to_logical_map += [ |
| 316 | i % num_routed_experts for i in range(num_redundant_experts) |
| 317 | ] |
| 318 | return global_physical_to_logical_map |
| 319 | |
| 320 | def validate_ep_configuration(self, new_model: MixtureOfExperts): |
| 321 | """Validate that the expert parallel configuration of |
| 322 | the new model is the same as the existing models. |
| 323 | """ |
| 324 | if len(self.model_states) > 0: |
| 325 | model = next(iter(self.model_states.values())).model |
| 326 | if ( |
| 327 | model.num_routed_experts != new_model.num_routed_experts |
| 328 | or model.num_redundant_experts != new_model.num_redundant_experts |
| 329 | or model.num_physical_experts != new_model.num_physical_experts |
| 330 | or model.num_logical_experts != new_model.num_logical_experts |
| 331 | or model.num_expert_groups != new_model.num_expert_groups |
| 332 | ): |
| 333 | raise RuntimeError( |
| 334 | "Model: {} " |
| 335 | "with config {} " |
| 336 | "{} {} {} {} " |
| 337 | "mismatch with new model {} " |
| 338 | "with config {} " |
| 339 | "{} {} {} {}".format( |
| 340 | type(model), |
| 341 | model.num_routed_experts, |
| 342 | model.num_redundant_experts, |
| 343 | model.num_physical_experts, |
| 344 | model.num_logical_experts, |
| 345 | model.num_expert_groups, |
| 346 | type(new_model), |
| 347 | new_model.num_routed_experts, |
| 348 | new_model.num_redundant_experts, |
| 349 | new_model.num_physical_experts, |
| 350 | new_model.num_logical_experts, |
| 351 | new_model.num_expert_groups, |
| 352 | ) |
| 353 | ) |
| 354 | |
| 355 | def add_model( |
| 356 | self, |
| 357 | model: MixtureOfExperts, |
| 358 | model_config: ModelConfig, |
| 359 | ): |
| 360 | """Build the initial EPLB state.""" |
| 361 | self.validate_ep_configuration(model) |
| 362 | self.is_async = self.parallel_config.eplb_config.use_async |
| 363 | |
| 364 | physical_to_logical_map_list = ( |
| 365 | EplbState.build_initial_global_physical_to_logical_map( |
| 366 | model.num_routed_experts, |
| 367 | model.num_redundant_experts, |
| 368 | ) |
| 369 | ) |
| 370 | physical_to_logical_map = torch.tensor( |
| 371 | physical_to_logical_map_list, |
| 372 | device=self.device, |
| 373 | ) |
| 374 | physical_expert_capacity = model.num_physical_experts |
| 375 | if self.parallel_config.enable_elastic_ep: |
| 376 | physical_expert_capacity = ( |
| 377 | model.num_local_physical_experts |
| 378 | * self.parallel_config.elastic_ep_max_dp_size |
| 379 | * get_ep_group().world_size |
| 380 | // self.parallel_config.data_parallel_size |
| 381 | ) |
| 382 | physical_to_logical_map = torch.nn.functional.pad( |
| 383 | physical_to_logical_map, |
| 384 | (0, physical_expert_capacity - model.num_physical_experts), |
| 385 | value=-1, |
| 386 | ) |
| 387 | # Assuming 8 GPUs per node, this supports up to |
| 388 | # (1023 + 1) / 8 = 128 nodes for now. |
| 389 | # TODO(rui): make this configurable |
| 390 | MAX_EXPERT_REDUNDANCY = 1023 |
| 391 | assert model.num_redundant_experts <= MAX_EXPERT_REDUNDANCY, ( |
| 392 | f"num_redundant_experts {model.num_redundant_experts} " |
| 393 | f"must be less than or equal to {MAX_EXPERT_REDUNDANCY}" |
| 394 | ) |
| 395 | max_slots_per_logical_expert = MAX_EXPERT_REDUNDANCY + 1 |
| 396 | logical_to_physical_map = torch.full( |
| 397 | (model.num_logical_experts, max_slots_per_logical_expert), |
| 398 | -1, |
| 399 | device=self.device, |
| 400 | ) |
| 401 | logical_replica_count = torch.zeros( |
| 402 | (model.num_logical_experts,), |
| 403 | device=self.device, |
| 404 | dtype=torch.long, |
| 405 | ) |
| 406 | |
| 407 | for i in range(model.num_physical_experts): |
| 408 | logical_idx = physical_to_logical_map[i] |
| 409 | logical_to_physical_map[logical_idx, logical_replica_count[logical_idx]] = i |
| 410 | logical_replica_count[logical_idx] += 1 |
| 411 | |
| 412 | # Duplicate initial mapping for all layers |
| 413 | physical_to_logical_map = ( |
| 414 | physical_to_logical_map.unsqueeze(0) |
| 415 | .expand( |
| 416 | model.num_moe_layers, |
| 417 | -1, |
| 418 | ) |
| 419 | .contiguous() |
| 420 | ) |
| 421 | physical_to_logical_map_buffer = physical_to_logical_map |
| 422 | physical_to_logical_map = physical_to_logical_map_buffer[ |
| 423 | :, : model.num_physical_experts |
| 424 | ] |
| 425 | logical_to_physical_map = ( |
| 426 | logical_to_physical_map.unsqueeze(0) |
| 427 | .expand( |
| 428 | model.num_moe_layers, |
| 429 | -1, |
| 430 | -1, |
| 431 | ) |
| 432 | .contiguous() |
| 433 | ) |
| 434 | logical_replica_count = ( |
| 435 | logical_replica_count.unsqueeze(0) |
| 436 | .expand( |
| 437 | model.num_moe_layers, |
| 438 | -1, |
| 439 | ) |
| 440 | .contiguous() |
| 441 | ) |
| 442 | |
| 443 | expert_load_pass_buffer = torch.zeros( |
| 444 | (model.num_moe_layers, physical_expert_capacity), |
| 445 | dtype=torch.int32, |
| 446 | device=self.device, |
| 447 | ) |
| 448 | expert_load_pass = expert_load_pass_buffer[:, : model.num_physical_experts] |
| 449 | self.expert_load_window_size = self.parallel_config.eplb_config.window_size |
| 450 | expert_load_window = torch.zeros( |
| 451 | ( |
| 452 | self.expert_load_window_size, |
| 453 | model.num_moe_layers, |
| 454 | model.num_physical_experts, |
| 455 | ), |
| 456 | dtype=torch.int32, |
| 457 | device=self.device, |
| 458 | ) |
| 459 | |
| 460 | # Set the initial progress of rearrangement to 3/4 |
| 461 | eplb_step_interval = self.parallel_config.eplb_config.step_interval |
| 462 | self.expert_rearrangement_step = max( |
| 463 | 0, eplb_step_interval - eplb_step_interval // 4 |
| 464 | ) |
| 465 | self.expert_rearrangement_step_interval = eplb_step_interval |
| 466 | |
| 467 | policy_type = self.parallel_config.eplb_config.policy |
| 468 | self.policy = EPLB_POLICIES[policy_type] |
| 469 | logger.debug("Selected EPLB policy: %s", policy_type) |
| 470 | |
| 471 | # num_ubatches is 0 when DBO is disabled. |
| 472 | num_ubatches = max(1, self.parallel_config.num_ubatches) |
| 473 | num_unpadded_tokens_tensors = [ |
| 474 | torch.tensor(0, dtype=torch.int32, device=self.device) |
| 475 | for _ in range(num_ubatches) |
| 476 | ] |
| 477 | |
| 478 | model.set_eplb_state( |
| 479 | expert_load_pass_buffer, |
| 480 | logical_to_physical_map, |
| 481 | logical_replica_count, |
| 482 | ) |
| 483 | self._propagate_shared_tensors(model, num_unpadded_tokens_tensors) |
| 484 | expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]] |
| 485 | |
| 486 | assert self.parallel_config.eplb_config.communicator is not None, ( |
| 487 | "EPLB communicator backend must be set by ParallelConfig" |
| 488 | ) |
| 489 | communicator = create_eplb_communicator( |
| 490 | group_coordinator=get_eplb_group(), |
| 491 | backend=self.parallel_config.eplb_config.communicator, |
| 492 | expert_weights=model.expert_weights, |
| 493 | expert_buffer=expert_buffer, |
| 494 | ) |
| 495 | |
| 496 | model_state = EplbModelState( |
| 497 | physical_to_logical_map=physical_to_logical_map, |
| 498 | physical_to_logical_map_buffer=physical_to_logical_map_buffer, |
| 499 | logical_to_physical_map=logical_to_physical_map, |
| 500 | logical_replica_count=logical_replica_count, |
| 501 | expert_load_pass=expert_load_pass, |
| 502 | expert_load_pass_buffer=expert_load_pass_buffer, |
| 503 | expert_load_window=expert_load_window, |
| 504 | last_expert_load=torch.ones( |
| 505 | (model.num_moe_layers, model.num_logical_experts), |
| 506 | dtype=torch.int64, |
| 507 | device=self.device, |
| 508 | ), |
| 509 | model_name=model_config.model, |
| 510 | model=model, |
| 511 | expert_buffer=expert_buffer, |
| 512 | rebalanced=False, |
| 513 | eplb_stats=None, |
| 514 | device_index=self.device_index, |
| 515 | communicator=communicator, |
| 516 | num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, |
| 517 | ) |
| 518 | self.model_states[model_config.compute_hash()] = model_state |
| 519 | |
| 520 | def prepare_forward( |
| 521 | self, |
| 522 | model_config: ModelConfig, |
| 523 | num_unpadded_tokens: int, |
| 524 | ubatch_slices: list | None = None, |
| 525 | ) -> None: |
| 526 | """Fill the per-[u]batch ``num_unpadded_tokens`` tensors before a |
| 527 | forward pass. |
| 528 | |
| 529 | Args: |
| 530 | model_config: Identifies which ``EplbModelState`` to update. |
| 531 | num_unpadded_tokens: Total number of real (non-padding) tokens |
| 532 | in the batch. |
| 533 | ubatch_slices: When DBO is active, a list of |
| 534 | ``UBatchSlice`` objects describing each micro-batch's |
| 535 | token range. When ``None``, only ``tensors[0]`` is filled. |
| 536 | |
| 537 | """ |
| 538 | model_state = self.model_states.get(compute_hash_cached(model_config)) |
| 539 | if model_state is None or model_state.num_unpadded_tokens_tensors is None: |
| 540 | return |
| 541 | tensors = model_state.num_unpadded_tokens_tensors |
| 542 | if ubatch_slices is None: |
| 543 | tensors[0].fill_(num_unpadded_tokens) |
| 544 | else: |
| 545 | for i, ubatch_slice in enumerate(ubatch_slices): |
| 546 | ts = ubatch_slice.token_slice |
| 547 | # Real tokens in this ubatch: clamp the global count into |
| 548 | # the slice range so partially-filled ubatches get the |
| 549 | # correct count. |
| 550 | val = max(0, min(num_unpadded_tokens, ts.stop) - ts.start) |
| 551 | tensors[i].fill_(val) |
| 552 | |
| 553 | def step( |
| 554 | self, |
| 555 | is_dummy: bool = False, |
| 556 | is_profile: bool = False, |
| 557 | log_stats: bool = False, |
| 558 | ) -> None: |
| 559 | """Step the EPLB state. |
| 560 | |
| 561 | Args: |
| 562 | is_dummy (bool): If `True`, this is a dummy step and the load |
| 563 | metrics recorded in this forward pass will not count. |
| 564 | Defaults to `False`. |
| 565 | is_profile (bool): If `True`, perform a dummy rearrangement |
| 566 | with maximum communication cost. This is used in |
| 567 | `profile_run` to reserve enough memory |
| 568 | for the communication buffer. |
| 569 | log_stats (bool): If `True`, log the expert load metrics. |
| 570 | |
| 571 | # Stats |
| 572 | The metrics are all summed up across layers. |
| 573 | - `avg_tokens`: The average load across ranks. |
| 574 | - `max_tokens`: The maximum load across ranks. |
| 575 | - `balancedness`: The ratio of average load to maximum load. |
| 576 | |
| 577 | """ |
| 578 | ep_group = get_ep_group().device_group |
| 579 | if is_profile: |
| 580 | self.rearrange(is_profile=True) |
| 581 | return |
| 582 | |
| 583 | if is_dummy: |
| 584 | # Do not record load metrics for dummy steps |
| 585 | for eplb_model_state in self.model_states.values(): |
| 586 | eplb_model_state.expert_load_pass.zero_() |
| 587 | |
| 588 | if ( |
| 589 | log_stats |
| 590 | and self.expert_rearrangement_step |
| 591 | % self.parallel_config.eplb_config.log_balancedness_interval |
| 592 | == 0 |
| 593 | ): |
| 594 | # Sync the expert load pass for each model (main and drafter). |
| 595 | # expert_load_pass: (num_moe_layers, num_physical_experts) |
| 596 | expert_load_pass_list = self._sync_load_pass() |
| 597 | ep_group = get_ep_group().device_group |
| 598 | for expert_load_pass, eplb_model_state in zip( |
| 599 | expert_load_pass_list, self.model_states.values() |
| 600 | ): |
| 601 | # num_tokens_per_rank: (num_moe_layers, num_ranks) |
| 602 | num_tokens_per_rank = ( |
| 603 | expert_load_pass.reshape( |
| 604 | expert_load_pass.shape[0], ep_group.size(), -1 |
| 605 | ) |
| 606 | .sum(dim=-1) |
| 607 | .float() |
| 608 | ) |
| 609 | |
| 610 | # Compute balancedness ratio: |
| 611 | # for each layer: |
| 612 | # (mean load across ranks) / (max load across ranks) |
| 613 | avg_tokens_tensor, max_tokens_tensor = _compute_eplb_load_stats( |
| 614 | num_tokens_per_rank |
| 615 | ) |
| 616 | |
| 617 | # This gpu/cpu sync only happens with expert stat logging enabled. |
| 618 | with gpu_sync_allowed(): |
| 619 | tokens_tensors: list[float] = torch.stack( |
| 620 | [avg_tokens_tensor, max_tokens_tensor] |
| 621 | ).tolist() |
| 622 | avg_tokens, max_tokens = tokens_tensors |
| 623 | balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 |
| 624 | |
| 625 | if ep_group.rank() == 0: |
| 626 | logger.info( |
| 627 | "EPLB step: %d for model %s: avg_tokens=%.2f, " |
| 628 | "max_tokens=%d, balancedness=%.4f, " |
| 629 | "steps until the next rearrangement: %d", |
| 630 | self.expert_rearrangement_step, |
| 631 | eplb_model_state.model_name, |
| 632 | avg_tokens, |
| 633 | max_tokens, |
| 634 | balancedness, |
| 635 | self.expert_rearrangement_step_interval |
| 636 | - self.expert_rearrangement_step, |
| 637 | ) |
| 638 | |
| 639 | # Update the expert load sliding window |
| 640 | if not is_dummy: |
| 641 | should_record = self._should_record_current_step(log_stats=log_stats) |
| 642 | for eplb_model_state in self.model_states.values(): |
| 643 | if should_record: |
| 644 | eplb_model_state.expert_load_window[ |
| 645 | self.expert_load_window_step |
| 646 | ].copy_(eplb_model_state.expert_load_pass) |
| 647 | eplb_model_state.expert_load_pass.zero_() |
| 648 | |
| 649 | if should_record: |
| 650 | self.expert_load_window_step += 1 |
| 651 | if self.expert_load_window_step >= self.expert_load_window_size: |
| 652 | self.expert_load_window_step = 0 |
| 653 | |
| 654 | # Step the expert rearrangement step |
| 655 | # Note that even if this is a dummy step, we still increment the |
| 656 | # rearrangement step and perform rearrangement to ensure all ranks are |
| 657 | # performing collective communication. |
| 658 | self.expert_rearrangement_step += 1 |
| 659 | |
| 660 | if self.is_async: |
| 661 | # Run _move_to_workspace if all ranks have finished transferring the |
| 662 | # new weights to the intermediate buffer |
| 663 | for eplb_model_state in self.model_states.values(): |
| 664 | # rebalanced must remain consistent amongst all ranks otherwise the |
| 665 | # all_reduce in _all_ranks_result_ready will hang |
| 666 | if eplb_model_state.rebalanced and self._all_ranks_result_ready( |
| 667 | eplb_model_state |
| 668 | ): |
| 669 | _move_to_workspace( |
| 670 | model_state=eplb_model_state, |
| 671 | ep_rank=ep_group.rank(), |
| 672 | ) |
| 673 | |
| 674 | if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval: |
| 675 | if self.is_async and any( |
| 676 | eplb_model_state.rebalanced |
| 677 | for eplb_model_state in self.model_states.values() |
| 678 | ): |
| 679 | # Still performing asynchronous rearrangement; update |
| 680 | # should_record (step > step_interval, so always True) and |
| 681 | # bail out before the step counter is reset. |
| 682 | self._update_layer_should_record(log_stats=log_stats) |
| 683 | return |
| 684 | self.expert_rearrangement_step = 0 |
| 685 | self.rearrange() |
| 686 | |
| 687 | self._update_layer_should_record(log_stats=log_stats) |
| 688 | |
| 689 | def _should_record_current_step(self, log_stats: bool = False) -> bool: |
| 690 | """Return whether expert-load recording should be enabled this step. |
| 691 | |
| 692 | Recording is enabled when we are close to either: |
| 693 | 1) The next rearrangement step, so the sliding window is ready. |
| 694 | 2) The next balancedness logging step, when log_stats is enabled. |
| 695 | """ |
| 696 | steps_remaining = ( |
| 697 | self.expert_rearrangement_step_interval - self.expert_rearrangement_step |
| 698 | ) |
| 699 | should_record_for_rearrange = steps_remaining <= self.expert_load_window_size |
| 700 | |
| 701 | if not log_stats: |
| 702 | return should_record_for_rearrange |
| 703 | |
| 704 | log_interval = self.parallel_config.eplb_config.log_balancedness_interval |
| 705 | steps_until_next_log = ( |
| 706 | log_interval - (self.expert_rearrangement_step % log_interval) |
| 707 | ) % log_interval |
| 708 | should_record_for_log = steps_until_next_log <= self.expert_load_window_size |
| 709 | return should_record_for_rearrange or should_record_for_log |
| 710 | |
| 711 | def _update_layer_should_record(self, log_stats: bool = False) -> None: |
| 712 | """Update the shared ``should_record_tensor`` for all layers.""" |
| 713 | if self.should_record_tensor is not None: |
| 714 | self.should_record_tensor.fill_( |
| 715 | self._should_record_current_step(log_stats=log_stats) |
| 716 | ) |
| 717 | |
| 718 | def _propagate_shared_tensors( |
| 719 | self, |
| 720 | model: "MixtureOfExperts", # type: ignore[name-defined] |
| 721 | num_unpadded_tokens_tensors: list[torch.Tensor], |
| 722 | ) -> None: |
| 723 | """Propagate shared tensors to every :class:`EplbLayerState`. |
| 724 | |
| 725 | Allocates ``should_record_tensor`` on the first call and then |
| 726 | assigns both it and ``num_unpadded_tokens_tensors`` to every |
| 727 | MoE layer's :class:`EplbLayerState`. All layers reference the |
| 728 | **same** objects so a single update is visible everywhere. |
| 729 | |
| 730 | Must be called after :meth:`model.set_eplb_state` so that each |
| 731 | layer's ``eplb_state`` is already populated. |
| 732 | """ |
| 733 | layer_states = [ |
| 734 | layer.eplb_state |
| 735 | for layer in model.moe_layers |
| 736 | if hasattr(layer, "eplb_state") |
| 737 | and isinstance(layer.eplb_state, EplbLayerState) |
| 738 | ] |
| 739 | |
| 740 | if self.should_record_tensor is None and layer_states: |
| 741 | self.should_record_tensor = torch.ones( |
| 742 | (), dtype=torch.bool, device=self.device |
| 743 | ) |
| 744 | |
| 745 | for ls in layer_states: |
| 746 | if ls is not None: |
| 747 | ls.should_record_tensor = self.should_record_tensor |
| 748 | ls.num_unpadded_tokens_tensors = num_unpadded_tokens_tensors |
| 749 | |
| 750 | def rearrange( |
| 751 | self, |
| 752 | is_profile: bool = False, |
| 753 | rank_mapping: dict[int, int] | None = None, |
| 754 | use_last_expert_load: bool = False, |
| 755 | ) -> torch.Tensor | None: |
| 756 | """Rearrange the experts according to the current load. |
| 757 | |
| 758 | Args: |
| 759 | is_profile (bool): If `True`, perform a dummy rearrangement. |
| 760 | This is used in `profile_run` to reserve enough memory, |
| 761 | no memory movement will be performed. Default is False. |
| 762 | rank_mapping (dict[int, int] | None): The rank mapping |
| 763 | when scaling is done in EEP. |
| 764 | use_last_expert_load: Use the last reshuffle's logical-expert totals. |
| 765 | |
| 766 | """ |
| 767 | ep_group = get_ep_group().device_group |
| 768 | ep_rank = ep_group.rank() |
| 769 | |
| 770 | start_event = None |
| 771 | end_event = None |
| 772 | is_main_rank = ep_rank == 0 |
| 773 | if is_main_rank: |
| 774 | if not self.is_async or is_profile: |
| 775 | start_event = torch.Event(enable_timing=True) |
| 776 | end_event = torch.Event(enable_timing=True) |
| 777 | start_event.record() |
| 778 | logger.info( |
| 779 | "Rearranging experts %s %s...", |
| 780 | "(async mode)" if self.is_async else "sync mode", |
| 781 | "(profile)" if is_profile else "", |
| 782 | ) |
| 783 | |
| 784 | # Map the physical expert load to global logical experts |
| 785 | global_expert_load_windows = [] |
| 786 | for eplb_model_state in self.model_states.values(): |
| 787 | if use_last_expert_load: |
| 788 | # Distribute rank 0's saved counts through the existing all-reduce. |
| 789 | saved_load = eplb_model_state.last_expert_load |
| 790 | global_expert_load_windows.append( |
| 791 | saved_load.clone() if is_main_rank else torch.zeros_like(saved_load) |
| 792 | ) |
| 793 | continue |
| 794 | expert_load_window = eplb_model_state.expert_load_window |
| 795 | physical_to_logical = eplb_model_state.physical_to_logical_map |
| 796 | invalid_idx = eplb_model_state.model.num_logical_experts |
| 797 | logical_expert_load_window = torch.zeros( |
| 798 | self.expert_load_window_size, |
| 799 | eplb_model_state.model.num_moe_layers, |
| 800 | invalid_idx + 1, |
| 801 | dtype=eplb_model_state.expert_load_window.dtype, |
| 802 | device=eplb_model_state.expert_load_window.device, |
| 803 | ) |
| 804 | logical_expert_load_window.scatter_add_( |
| 805 | dim=-1, |
| 806 | index=physical_to_logical.masked_fill( |
| 807 | physical_to_logical < 0, invalid_idx |
| 808 | ) |
| 809 | .unsqueeze(0) |
| 810 | .expand_as(expert_load_window) |
| 811 | .long(), |
| 812 | src=expert_load_window, |
| 813 | ) |
| 814 | |
| 815 | global_expert_load_window = logical_expert_load_window[..., :-1].sum(dim=0) |
| 816 | global_expert_load_windows.append(global_expert_load_window) |
| 817 | # Perform all-reduce to get the expert load across all ranks for each model |
| 818 | global_expert_load_windows = self._allreduce_list(global_expert_load_windows) |
| 819 | |
| 820 | # TODO(bowen): Treat differently for prefill and decode nodes |
| 821 | eplb_model_state = next(iter(self.model_states.values())) |
| 822 | model = eplb_model_state.model |
| 823 | num_replicas = model.num_physical_experts |
| 824 | num_groups = model.num_expert_groups |
| 825 | |
| 826 | if rank_mapping is not None and len(rank_mapping) == ep_group.size(): |
| 827 | # NOTE(yongji): scale down, we need to rebalance the experts on |
| 828 | # remaining GPUs, transfer the experts while we haven't shutdown |
| 829 | # the GPUs to be released. |
| 830 | coordinator = get_ep_group() |
| 831 | assert isinstance(coordinator, StatelessGroupCoordinator) |
| 832 | tcp_store_group = coordinator.tcp_store_group |
| 833 | num_nodes = _node_count_with_rank_mapping(tcp_store_group, rank_mapping) |
| 834 | num_gpus = sum(new_rank != -1 for new_rank in rank_mapping.values()) |
| 835 | num_replicas = ( |
| 836 | num_replicas // ep_group.size() * num_gpus |
| 837 | ) # handle num replicas change |
| 838 | else: |
| 839 | num_nodes = get_node_count() |
| 840 | num_gpus = ep_group.size() |
| 841 | |
| 842 | if num_gpus % num_nodes != 0: |
| 843 | num_nodes = 1 |
| 844 | logger.warning_once( |
| 845 | f"num_gpus % num_nodes != 0, " |
| 846 | "not using hierarchical rearrangement algorithm.\n" |
| 847 | f"{num_gpus=}, {num_nodes=}" |
| 848 | ) |
| 849 | |
| 850 | # Get new expert mappings |
| 851 | for eplb_model_state, global_expert_load_window in zip( |
| 852 | self.model_states.values(), global_expert_load_windows |
| 853 | ): |
| 854 | if not is_profile: |
| 855 | eplb_model_state.last_expert_load = global_expert_load_window |
| 856 | if not self.is_async or is_profile: |
| 857 | # Get new expert mappings for the model. The policy runs on the |
| 858 | # host, so the load window and current map have to come back. |
| 859 | with gpu_sync_allowed(): |
| 860 | new_physical_to_logical_map = self.policy.rebalance_experts( |
| 861 | global_expert_load_window.cpu(), |
| 862 | num_replicas, |
| 863 | num_groups, |
| 864 | num_nodes, |
| 865 | num_gpus, |
| 866 | eplb_model_state.physical_to_logical_map.cpu(), |
| 867 | ) |
| 868 | |
| 869 | skip_rearrange = False |
| 870 | if ( |
| 871 | current_platform.is_rocm() |
| 872 | and not is_profile |
| 873 | and rank_mapping is None |
| 874 | and bool((eplb_model_state.physical_to_logical_map >= 0).all()) |
| 875 | ): |
| 876 | logical_loads = global_expert_load_window.float() |
| 877 | ep_size = ep_group.size() |
| 878 | |
| 879 | def rank_load_imbalance( |
| 880 | mapping: torch.Tensor, |
| 881 | logical_loads: torch.Tensor = logical_loads, |
| 882 | ep_size: int = ep_size, |
| 883 | ) -> float: |
| 884 | mapping = mapping.to( |
| 885 | device=logical_loads.device, |
| 886 | dtype=torch.long, |
| 887 | ) |
| 888 | replica_counts = torch.zeros_like(logical_loads) |
| 889 | replica_counts.scatter_add_( |
| 890 | dim=1, |
| 891 | index=mapping, |
| 892 | src=torch.ones_like(mapping, dtype=logical_loads.dtype), |
| 893 | ) |
| 894 | loads_per_replica = torch.gather( |
| 895 | logical_loads / replica_counts.clamp_min(1), |
| 896 | dim=1, |
| 897 | index=mapping, |
| 898 | ) |
| 899 | loads_per_rank = loads_per_replica.reshape( |
| 900 | logical_loads.shape[0], ep_size, -1 |
| 901 | ).sum(dim=(0, 2)) |
| 902 | mean_load = loads_per_rank.mean() |
| 903 | if mean_load == 0: |
| 904 | return 1.0 |
| 905 | return (loads_per_rank.max() / mean_load).item() |
| 906 | |
| 907 | current_imbalance = rank_load_imbalance( |
| 908 | eplb_model_state.physical_to_logical_map |
| 909 | ) |
| 910 | proposed_imbalance = rank_load_imbalance( |
| 911 | new_physical_to_logical_map |
| 912 | ) |
| 913 | relative_improvement = ( |
| 914 | current_imbalance - proposed_imbalance |
| 915 | ) / current_imbalance |
| 916 | skip_rearrange = relative_improvement < 0.05 |
| 917 | if skip_rearrange and is_main_rank: |
| 918 | logger.info( |
| 919 | "[EPLB] Skip rearrange: imbalance %.4f -> " |
| 920 | "%.4f (no material gain)", |
| 921 | current_imbalance, |
| 922 | proposed_imbalance, |
| 923 | ) |
| 924 | |
| 925 | if not skip_rearrange: |
| 926 | # Update expert weights |
| 927 | rearrange_expert_weights_inplace( |
| 928 | eplb_model_state.physical_to_logical_map, |
| 929 | new_physical_to_logical_map, |
| 930 | eplb_model_state.model.expert_weights, |
| 931 | eplb_model_state.expert_buffer, |
| 932 | ep_group, |
| 933 | eplb_model_state.communicator, |
| 934 | is_profile, |
| 935 | rank_mapping, |
| 936 | ) |
| 937 | |
| 938 | if not is_profile: |
| 939 | _commit_eplb_maps( |
| 940 | eplb_model_state, |
| 941 | new_physical_to_logical_map=new_physical_to_logical_map, |
| 942 | ) |
| 943 | |
| 944 | if is_main_rank: |
| 945 | assert start_event is not None |
| 946 | assert end_event is not None |
| 947 | end_event.record() |
| 948 | end_event.synchronize() |
| 949 | gpu_elapsed = start_event.elapsed_time(end_event) / 1000.0 |
| 950 | logger.info( |
| 951 | "Rearranged experts %s in %.2f s.", |
| 952 | " (profile) " if is_profile else " ", |
| 953 | gpu_elapsed, |
| 954 | ) |
| 955 | else: |
| 956 | eplb_model_state.eplb_stats = EplbStats( |
| 957 | # We copy the tensor to snapshot the global_expert_load_window |
| 958 | # on the main thread so that async worker can access it safely |
| 959 | # while the main thread is running. |
| 960 | global_expert_load_window=global_expert_load_window.clone(), |
| 961 | num_replicas=num_replicas, |
| 962 | num_groups=num_groups, |
| 963 | num_nodes=num_nodes, |
| 964 | num_gpus=num_gpus, |
| 965 | ) |
| 966 | eplb_model_state.rebalanced = True |
| 967 | # Signal async thread to start transferring layers |
| 968 | if self.is_async and (not is_profile): |
| 969 | self.rearrange_event.record() |
| 970 | return None |
| 971 | |
| 972 | def start_async_loop( |
| 973 | self, |
| 974 | rank_mapping: dict[int, int] | None = None, |
| 975 | is_profile: bool = False, |
| 976 | ): |
| 977 | if not self.is_async: |
| 978 | return |
| 979 | if self.async_worker is None: |
| 980 | self.async_worker = start_async_worker( |
| 981 | self, |
| 982 | is_profile=is_profile, |
| 983 | ) |
| 984 | |
| 985 | def drain_async(self) -> None: |
| 986 | """Drain in-flight async EPLB by consuming all remaining layer results. |
| 987 | |
| 988 | Each pending result is acknowledged (consumed_event recorded) so the |
| 989 | async worker can proceed, but the transferred weights are intentionally |
| 990 | NOT applied — a full rearrange is expected to follow. |
| 991 | |
| 992 | Ranks are kept in lockstep via _all_ranks_result_ready (all_reduce |
| 993 | on the EP CPU group). The async worker's coordinated-stop collectives |
| 994 | use the separate EPLB group, so the two sets of collectives do not |
| 995 | interfere. |
| 996 | |
| 997 | No-op when no async cycle is in progress (rebalanced=False). |
| 998 | """ |
| 999 | if not self.is_async: |
| 1000 | return |
| 1001 | for model_key, ms in self.model_states.items(): |
| 1002 | needs_drain = ms.rebalanced |
| 1003 | if needs_drain: |
| 1004 | logger.info( |
| 1005 | "Draining async EPLB worker for model %s", |
| 1006 | model_key, |
| 1007 | ) |
| 1008 | while ms.rebalanced: |
| 1009 | if self._all_ranks_result_ready(ms): |
| 1010 | result = ms.pending_result |
| 1011 | assert result is not None |
| 1012 | if result.layer_idx == ms.model.num_moe_layers - 1: |
| 1013 | ms.rebalanced = False |
| 1014 | ms.pending_result = None |
| 1015 | result.consumed_event.record() |
| 1016 | else: |
| 1017 | time.sleep(0.001) |
| 1018 | if needs_drain: |
| 1019 | logger.info( |
| 1020 | "Async EPLB worker drained for model %s", |
| 1021 | model_key, |
| 1022 | ) |
| 1023 | |
| 1024 | def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool: |
| 1025 | parallel_state = get_ep_group() |
| 1026 | has_result = int(model_state.pending_result is not None) |
| 1027 | |
| 1028 | cpu_group = getattr(parallel_state, "cpu_group", None) |
| 1029 | if cpu_group is not None and cpu_group.size() > 1: |
| 1030 | flag = torch.tensor((has_result,), dtype=torch.int32, device="cpu") |
| 1031 | all_reduce(flag, group=cpu_group) |
| 1032 | return int(flag.item()) == cpu_group.size() |
| 1033 | |
| 1034 | device_group = parallel_state.device_group |
| 1035 | if device_group.size() <= 1: |
| 1036 | return bool(has_result) |
| 1037 | |
| 1038 | device = getattr( |
| 1039 | parallel_state, "device", model_state.physical_to_logical_map.device |
| 1040 | ) |
| 1041 | flag = torch.tensor((has_result,), dtype=torch.int32, device=device) |
| 1042 | all_reduce(flag, group=device_group) |
| 1043 | return int(flag.item()) == device_group.size() |
| 1044 | |
| 1045 | def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]: |
| 1046 | """All-reduce a list of tensors.""" |
| 1047 | ep_group = get_ep_group().device_group |
| 1048 | if len(tensor_list) == 1: |
| 1049 | all_reduce(tensor_list[0], group=ep_group) |
| 1050 | return tensor_list |
| 1051 | assert all(t.dim() == 2 for t in tensor_list), "All tensors must be 2D." |
| 1052 | assert all(t.shape[1] == tensor_list[0].shape[1] for t in tensor_list), ( |
| 1053 | "All tensors must have the same shape[1]." |
| 1054 | ) |
| 1055 | # Concatenate, all_reduce, then unpack to original shapes. |
| 1056 | # We assume all tensors are 2D and shape[1] (num_physical_experts) |
| 1057 | # is the same across all models. |
| 1058 | shapes = [t.shape for t in tensor_list] |
| 1059 | concat_tensor = torch.cat(tensor_list, dim=0) |
| 1060 | |
| 1061 | all_reduce(concat_tensor, group=ep_group) |
| 1062 | |
| 1063 | all_reduce_list = [] |
| 1064 | offset = 0 |
| 1065 | for shape in shapes: |
| 1066 | all_reduce_list.append(concat_tensor[offset : offset + shape[0], :]) |
| 1067 | offset += shape[0] |
| 1068 | return all_reduce_list |
| 1069 | |
| 1070 | def _sync_load_pass(self) -> list[torch.Tensor]: |
| 1071 | """Sync the expert load pass across all ranks for log stats. |
| 1072 | Doesn't update the expert load pass in eplb_model_state. |
| 1073 | """ |
| 1074 | load_pass_list = [] |
| 1075 | for eplb_model_state in self.model_states.values(): |
| 1076 | load_pass_list.append(eplb_model_state.expert_load_pass.clone()) |
| 1077 | return self._allreduce_list(load_pass_list) |
| 1078 | |
| 1079 | @classmethod |
| 1080 | def from_mapping( |
| 1081 | cls, |
| 1082 | model: MixtureOfExperts, |
| 1083 | model_config: ModelConfig, |
| 1084 | device: torch.device, |
| 1085 | parallel_config: ParallelConfig, |
| 1086 | expanded_physical_to_logical: torch.Tensor, |
| 1087 | ) -> "EplbState": |
| 1088 | eplb_state = cls( |
| 1089 | parallel_config=parallel_config, |
| 1090 | device=device, |
| 1091 | ) |
| 1092 | eplb_state.add_model( |
| 1093 | model=model, |
| 1094 | model_config=model_config, |
| 1095 | ) |
| 1096 | eplb_state.update_mapping( |
| 1097 | model_config, |
| 1098 | expanded_physical_to_logical, |
| 1099 | ) |
| 1100 | |
| 1101 | return eplb_state |
| 1102 | |
| 1103 | def update_mapping( |
| 1104 | self, |
| 1105 | model_config: ModelConfig, |
| 1106 | expanded_physical_to_logical: torch.Tensor, |
| 1107 | ) -> None: |
| 1108 | eplb_model_state = self.model_states[model_config.compute_hash()] |
| 1109 | eplb_model_state.physical_to_logical_map_buffer.copy_( |
| 1110 | expanded_physical_to_logical |
| 1111 | ) |
| 1112 | eplb_model_state.expert_load_pass.zero_() |
| 1113 | physical_to_logical_map = eplb_model_state.physical_to_logical_map |
| 1114 | |
| 1115 | (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps( |
| 1116 | physical_to_logical_map.cpu(), |
| 1117 | eplb_model_state.model.num_logical_experts, |
| 1118 | ) |
| 1119 | |
| 1120 | max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1] |
| 1121 | num_replicas = logical_to_physical_map_cpu.shape[-1] |
| 1122 | logical_to_physical_map = torch.nn.functional.pad( |
| 1123 | logical_to_physical_map_cpu, |
| 1124 | ( |
| 1125 | 0, |
| 1126 | max_num_replicas - num_replicas, |
| 1127 | ), |
| 1128 | value=-1, |
| 1129 | ).to(self.device) |
| 1130 | logical_replica_count = logical_replica_count_cpu.to(self.device) |
| 1131 | |
| 1132 | eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map) |
| 1133 | eplb_model_state.logical_replica_count.copy_(logical_replica_count) |
| 1134 | |
| 1135 | def reconfigure_physical_expert_slots( |
| 1136 | self, |
| 1137 | model_config: ModelConfig, |
| 1138 | num_physical_experts: int, |
| 1139 | ) -> None: |
| 1140 | """Replace physical_to_logical_map and expert_load_pass with views |
| 1141 | covering the new active size. |
| 1142 | """ |
| 1143 | model_state = self.model_states[model_config.compute_hash()] |
| 1144 | old_num_physical_experts = model_state.model.num_physical_experts |
| 1145 | first_slot, last_slot = sorted((old_num_physical_experts, num_physical_experts)) |
| 1146 | physical_to_logical_map_buffer = model_state.physical_to_logical_map_buffer |
| 1147 | expert_load_pass_buffer = model_state.expert_load_pass_buffer |
| 1148 | assert last_slot <= physical_to_logical_map_buffer.shape[1] |
| 1149 | expert_slots = slice(first_slot, last_slot) |
| 1150 | physical_to_logical_map_buffer[:, expert_slots].fill_(-1) |
| 1151 | expert_load_pass_buffer[:, expert_slots].zero_() |
| 1152 | model_state.physical_to_logical_map = physical_to_logical_map_buffer[ |
| 1153 | :, :num_physical_experts |
| 1154 | ] |
| 1155 | model_state.expert_load_pass = expert_load_pass_buffer[:, :num_physical_experts] |
| 1156 | |
| 1157 | pad_size = num_physical_experts - model_state.expert_load_window.shape[-1] |
| 1158 | model_state.expert_load_window = torch.nn.functional.pad( |
| 1159 | model_state.expert_load_window, (0, pad_size) |
| 1160 | ) |
| 1161 | |
| 1162 | def create_communicator( |
| 1163 | self, model_config: ModelConfig, group_coordinator: GroupCoordinator |
| 1164 | ) -> EplbCommunicator: |
| 1165 | model_state = self.model_states[model_config.compute_hash()] |
| 1166 | backend = self.parallel_config.eplb_config.communicator |
| 1167 | assert backend is not None |
| 1168 | return create_eplb_communicator( |
| 1169 | group_coordinator, |
| 1170 | backend, |
| 1171 | model_state.model.expert_weights, |
| 1172 | model_state.expert_buffer, |
| 1173 | ) |
| 1174 | |
| 1175 | def update_communicator( |
| 1176 | self, |
| 1177 | model_config: ModelConfig, |
| 1178 | communicator: EplbCommunicator, |
| 1179 | ) -> None: |
| 1180 | self.model_states[model_config.compute_hash()].communicator = communicator |
| 1181 | |
| 1182 | |
| 1183 | @dataclass |
| 1184 | class EplbLayerState: |
| 1185 | """Runtime EPLB data stored in the MoE layer.""" |
| 1186 | |
| 1187 | expert_load_view: torch.Tensor | None = None |
| 1188 | logical_to_physical_map: torch.Tensor | None = None |
| 1189 | logical_replica_count: torch.Tensor | None = None |
| 1190 | should_record_tensor: torch.Tensor | None = None |
| 1191 | """ |
| 1192 | Shared scalar bool tensor controlling whether to accumulate expert load |
| 1193 | metrics during this forward pass. All layers reference the **same** |
| 1194 | tensor object, which is owned and updated by :class:`EplbState`. |
| 1195 | |
| 1196 | Set to ``False`` for the first ``step_interval - window_size`` steps of |
| 1197 | each rearrangement period: those steps would be overwritten in the |
| 1198 | sliding window before the next rearrangement, so recording them wastes |
| 1199 | GPU work. |
| 1200 | """ |
| 1201 | num_unpadded_tokens_tensors: list[torch.Tensor] | None = None |
| 1202 | """ |
| 1203 | Reference to the parent :class:`EplbModelState`'s tensor list so the |
| 1204 | router can read the correct per-[u]batch unpadded token count. |
| 1205 | """ |
| 1206 | |
| 1207 | def set_layer_state( |
| 1208 | self, |
| 1209 | moe_layer_idx: int, |
| 1210 | expert_load_view: torch.Tensor, |
| 1211 | logical_to_physical_map: torch.Tensor, |
| 1212 | logical_replica_count: torch.Tensor, |
| 1213 | ) -> None: |
| 1214 | self.expert_load_view = expert_load_view[moe_layer_idx] |
| 1215 | self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx] |
| 1216 | self.logical_replica_count = logical_replica_count[moe_layer_idx] |
| 1217 | |
| 1218 | |
| 1219 | def _node_count_with_rank_mapping( |
| 1220 | pg: ProcessGroup | StatelessProcessGroup, |
| 1221 | rank_mapping: dict[int, int], |
| 1222 | ) -> int: |
| 1223 | if isinstance(pg, ProcessGroup): |
| 1224 | world_size = torch.distributed.get_world_size(group=pg) |
| 1225 | else: |
| 1226 | world_size = pg.world_size |
| 1227 | |
| 1228 | if world_size == 1: |
| 1229 | return 1 |
| 1230 | |
| 1231 | # Build node assignment map |
| 1232 | node_assignment = [0] * world_size # rank -> node_id |
| 1233 | next_node_id = 0 |
| 1234 | |
| 1235 | for current_rank in range(world_size): |
| 1236 | if node_assignment[current_rank] != 0: |
| 1237 | continue # Already assigned to a node |
| 1238 | |
| 1239 | assert current_rank in rank_mapping |
| 1240 | if rank_mapping[current_rank] == -1: |
| 1241 | continue # Pending shutdown |
| 1242 | |
| 1243 | # Assign current rank to a new node |
| 1244 | next_node_id += 1 |
| 1245 | node_assignment[current_rank] = next_node_id |
| 1246 | |
| 1247 | # Find all ranks on the same node as current_rank |
| 1248 | same_node_flags = in_the_same_node_as(pg, current_rank) |
| 1249 | for other_rank, is_same_node in enumerate(same_node_flags): |
| 1250 | if is_same_node and node_assignment[other_rank] == 0: |
| 1251 | node_assignment[other_rank] = next_node_id |
| 1252 | |
| 1253 | return next_node_id |
| 1254 | |
| 1255 | |
| 1256 | def compute_logical_maps( |
| 1257 | physical_to_logical_map: torch.Tensor, |
| 1258 | num_logical_experts: int, |
| 1259 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1260 | """Derive logical_to_physical_map and logical_replica_count from |
| 1261 | physical_to_logical_map. |
| 1262 | |
| 1263 | Args: |
| 1264 | physical_to_logical_map: [num_layers, num_physical_experts], logical |
| 1265 | expert index for each physical expert slot |
| 1266 | num_logical_experts: total number of logical experts |
| 1267 | |
| 1268 | Returns: |
| 1269 | logical_to_physical_map: [num_layers, num_logical_experts, max_replicas], |
| 1270 | physical slots per logical expert; -1 where unused |
| 1271 | logical_replica_count: [num_layers, num_logical_experts], number of |
| 1272 | physical replicas per logical expert |
| 1273 | |
| 1274 | """ |
| 1275 | device = physical_to_logical_map.device |
| 1276 | assert physical_to_logical_map.device.type == "cpu" |
| 1277 | |
| 1278 | dtype = physical_to_logical_map.dtype |
| 1279 | |
| 1280 | # If computing maps for a single layer, unsqueeze a single element layer dimension |
| 1281 | per_layer = physical_to_logical_map.dim() == 1 |
| 1282 | physical_to_logical_map_view = physical_to_logical_map |
| 1283 | if per_layer: |
| 1284 | physical_to_logical_map_view = physical_to_logical_map.unsqueeze(0) |
| 1285 | assert len(physical_to_logical_map_view.shape) == 2 |
| 1286 | num_layers, num_physical = physical_to_logical_map_view.shape |
| 1287 | |
| 1288 | valid_mask = physical_to_logical_map_view >= 0 |
| 1289 | logical_replica_count = torch.zeros( |
| 1290 | num_layers, |
| 1291 | num_logical_experts, |
| 1292 | dtype=dtype, |
| 1293 | device=device, |
| 1294 | ) |
| 1295 | logical_replica_count.scatter_add_( |
| 1296 | 1, |
| 1297 | physical_to_logical_map_view.clamp(min=0), |
| 1298 | valid_mask.to(dtype), |
| 1299 | ) |
| 1300 | |
| 1301 | max_replicas = int(logical_replica_count.max().item()) |
| 1302 | logical_to_physical_map_out = torch.full( |
| 1303 | (num_layers, num_logical_experts, max_replicas), |
| 1304 | -1, |
| 1305 | dtype=dtype, |
| 1306 | device=device, |
| 1307 | ) |
| 1308 | |
| 1309 | running_count = torch.zeros_like(logical_replica_count) |
| 1310 | layer_indices = torch.arange(num_layers, device=device) |
| 1311 | for phys_idx in range(num_physical): |
| 1312 | # Logical expert at physical slot phys_idx for each layer |
| 1313 | logical_expert_ids = physical_to_logical_map_view[:, phys_idx] # [num_layers] |
| 1314 | |
| 1315 | # Scale up will set the logical expert ids to -1 for all new physical experts. |
| 1316 | # Only consider "valid" experts when setting up the logical_to_physical map. |
| 1317 | valid_expert_mask = logical_expert_ids >= 0 |
| 1318 | if not valid_expert_mask.any(): |
| 1319 | continue |
| 1320 | valid_layers = layer_indices[valid_expert_mask] |
| 1321 | valid_experts = logical_expert_ids[valid_expert_mask] |
| 1322 | |
| 1323 | # Use the current running count as the replica index, then increment it. |
| 1324 | replica_idx = running_count[valid_layers, valid_experts] |
| 1325 | logical_to_physical_map_out[valid_layers, valid_experts, replica_idx] = phys_idx |
| 1326 | running_count[valid_layers, valid_experts] += 1 |
| 1327 | |
| 1328 | # If computing maps for a single layer, squeeze out the extra layer dimension |
| 1329 | # before returning |
| 1330 | if per_layer: |
| 1331 | return logical_to_physical_map_out.squeeze(0), logical_replica_count.squeeze(0) |
| 1332 | return logical_to_physical_map_out, logical_replica_count |
| 1333 | |
| 1334 | |
| 1335 | def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: |
| 1336 | src_padding = dst.shape[-1] - src.shape[-1] |
| 1337 | assert src_padding >= 0 |
| 1338 | new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1) |
| 1339 | # The map is committed from the host once per layer per rearrangement. |
| 1340 | with gpu_sync_allowed(): |
| 1341 | dst.copy_(new_src) |
| 1342 | |
| 1343 | |
| 1344 | def _commit_eplb_maps_for_layer( |
| 1345 | model_state: EplbModelState, |
| 1346 | new_physical_to_logical_map: torch.Tensor, |
| 1347 | layer: int, |
| 1348 | ) -> None: |
| 1349 | """Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB |
| 1350 | when running async EPLB. Copies all of the new_* maps into model_state. After this |
| 1351 | function completes, the new mappings will become the current mappings and will be |
| 1352 | visible to the model. |
| 1353 | """ |
| 1354 | # Commit physical_to_logical_map |
| 1355 | src = new_physical_to_logical_map |
| 1356 | dst = model_state.physical_to_logical_map[layer] |
| 1357 | assert src.shape == dst.shape, ( |
| 1358 | "The number of physical experts must stay the same while running Async EPLB. " |
| 1359 | f"Current number of physical experts: {dst.shape[0]}. New number of physical " |
| 1360 | f"experts {src.shape[0]}." |
| 1361 | ) |
| 1362 | if PIN_MEMORY and src.is_cpu: |
| 1363 | src = src.new_empty(src.shape, pin_memory=True).copy_(src) |
| 1364 | dst.copy_(src, non_blocking=True) |
| 1365 | |
| 1366 | num_logical_experts = model_state.logical_to_physical_map.shape[1] |
| 1367 | new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts) |
| 1368 | # Commit logical_to_physical_map |
| 1369 | _pad_out_tensor(src=new_logical, dst=model_state.logical_to_physical_map[layer]) |
| 1370 | |
| 1371 | # Commit logical_replica_count |
| 1372 | src = new_replica_count |
| 1373 | dst = model_state.logical_replica_count[layer] |
| 1374 | assert src.shape == dst.shape |
| 1375 | if PIN_MEMORY: |
| 1376 | src = src.pin_memory() |
| 1377 | dst.copy_(src, non_blocking=True) |
| 1378 | |
| 1379 | |
| 1380 | def _commit_eplb_maps( |
| 1381 | model_state: EplbModelState, |
| 1382 | new_physical_to_logical_map: torch.Tensor, |
| 1383 | ) -> None: |
| 1384 | """Copies all of the new_* maps into model_state. After this function completes, |
| 1385 | the new mappings will become the current mappings and will be visible to the |
| 1386 | model. |
| 1387 | """ |
| 1388 | # Commit physical_to_logical_map |
| 1389 | src = new_physical_to_logical_map |
| 1390 | dst = model_state.physical_to_logical_map |
| 1391 | |
| 1392 | if PIN_MEMORY and src.is_cpu: |
| 1393 | src = src.new_empty(src.shape, pin_memory=True).copy_(src) |
| 1394 | # When the number of physical experts changes, refresh the active map view to |
| 1395 | # cover the new size before copying it. |
| 1396 | if src.shape[1] != dst.shape[1]: |
| 1397 | dst = model_state.physical_to_logical_map_buffer[:, : src.shape[1]] |
| 1398 | assert dst.shape == src.shape |
| 1399 | model_state.physical_to_logical_map = dst |
| 1400 | dst.copy_(src, non_blocking=True) |
| 1401 | |
| 1402 | num_logical_experts = model_state.logical_to_physical_map.shape[1] |
| 1403 | new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts) |
| 1404 | # Commit logical_to_physical_map |
| 1405 | _pad_out_tensor( |
| 1406 | src=new_logical, |
| 1407 | dst=model_state.logical_to_physical_map, |
| 1408 | ) |
| 1409 | |
| 1410 | # Commit logical_replica_count |
| 1411 | src = new_replica_count |
| 1412 | dst = model_state.logical_replica_count |
| 1413 | if PIN_MEMORY: |
| 1414 | src = src.pin_memory() |
| 1415 | dst.copy_(src, non_blocking=True) |
| 1416 | |
| 1417 | |
| 1418 | def _move_to_workspace( |
| 1419 | model_state: EplbModelState, |
| 1420 | ep_rank: int, |
| 1421 | ) -> None: |
| 1422 | result = model_state.pending_result |
| 1423 | assert result is not None |
| 1424 | move_from_buffer( |
| 1425 | expert_weights=model_state.model.expert_weights[result.layer_idx], |
| 1426 | expert_weights_buffers=model_state.expert_buffer, |
| 1427 | transfer_metadata=result.transfer_metadata, |
| 1428 | new_indices=result.new_physical_to_logical_map.numpy(), |
| 1429 | ep_rank=ep_rank, |
| 1430 | ) |
| 1431 | |
| 1432 | _commit_eplb_maps_for_layer( |
| 1433 | model_state, |
| 1434 | new_physical_to_logical_map=result.new_physical_to_logical_map, |
| 1435 | layer=result.layer_idx, |
| 1436 | ) |
| 1437 | |
| 1438 | if result.layer_idx == model_state.model.num_moe_layers - 1: |
| 1439 | model_state.rebalanced = False |
| 1440 | |
| 1441 | # Reset pending_result before unblocking the async worker |
| 1442 | model_state.pending_result = None |
| 1443 | result.consumed_event.record() |