Skip to content

vllm.distributed.eplb.policy.abstract

AbstractEplbPolicy

Bases: ABC

Source code in vllm/distributed/eplb/policy/abstract.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
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.

        Parameters:
            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

rebalance_experts abstractmethod classmethod 

rebalance_experts(
    weight: Tensor,
    num_replicas: int,
    num_groups: int,
    num_nodes: int,
    num_ranks: int,
    old_global_expert_indices: Tensor | None = None,
) -> Tensor

Entry point for expert-parallelism load balancer.

Parameters:

Name Type Description Default
weight Tensor

[layers, num_logical_experts], the load statistics for all logical experts

 required
num_replicas int

number of physical experts, must be a multiple of num_ranks

 required
num_groups int

number of expert groups

 required
num_nodes int

number of server nodes

 required
num_ranks int

number of ranks, must be a multiple of num_nodes

 required
old_global_expert_indices Tensor | None

[layers, num_logical_experts], the old global expert indices. Used to avoid unnecessary weight copying for experts moving within one rank.

 None

Returns: physical_to_logical_map: [layers, num_replicas], the expert index of each replica

Source code in vllm/distributed/eplb/policy/abstract.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@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.

    Parameters:
        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