Skip to content

vllm.distributed.eplb.policy.abstract

Classes:

AbstractEplbPolicy

Bases: ABC

Methods:

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(weight, num_replicas, num_groups, num_nodes, num_ranks, old_global_expert_indices=None) abstractmethod classmethod

Entry point for expert-parallelism load balancer.

Parameters:

  • weight

    (Tensor) –

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

  • num_replicas

    (int) –

    number of physical experts, must be a multiple of num_ranks

  • num_groups

    (int) –

    number of expert groups

  • num_nodes

    (int) –

    number of server nodes

  • num_ranks

    (int) –

    number of ranks, must be a multiple of num_nodes

  • old_global_expert_indices

    (Tensor | None, default: None ) –

    [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

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