Harry MellorandClaude Opus 5
8c1557a79c
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor <hej.mellor@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
import torch
|
|
|
|
|
|
class AbstractEplbPolicy(ABC):
|
|
@classmethod
|
|
@abstractmethod
|
|
def rebalance_experts(
|
|
cls,
|
|
weight: torch.Tensor,
|
|
num_replicas: int,
|
|
num_groups: int,
|
|
num_nodes: int,
|
|
num_ranks: int,
|
|
old_global_expert_indices: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Entry point for expert-parallelism load balancer.
|
|
|
|
Args:
|
|
weight: [layers, num_logical_experts], the load statistics
|
|
for all logical experts
|
|
num_replicas: number of physical experts, must be a multiple of
|
|
`num_ranks`
|
|
num_groups: number of expert groups
|
|
num_nodes: number of server nodes
|
|
num_ranks: number of ranks, must be a multiple of `num_nodes`
|
|
old_global_expert_indices: [layers, num_logical_experts], the old global
|
|
expert indices. Used to avoid unnecessary weight copying
|
|
for experts moving within one rank.
|
|
|
|
Returns:
|
|
physical_to_logical_map: [layers, num_replicas], the expert
|
|
index of each replica
|
|
|
|
"""
|
|
raise NotImplementedError
|