Skip to content

vllm_gaudi.distributed.device_communicators.hpu_communicator

HpuCommunicator

Bases: DeviceCommunicatorBase

Source code in vllm_gaudi/distributed/device_communicators/hpu_communicator.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class HpuCommunicator(DeviceCommunicatorBase):

    def __init__(self,
                 cpu_group: ProcessGroup,
                 device: Optional[torch.device] = None,
                 device_group: Optional[ProcessGroup] = None,
                 unique_name: str = ""):
        super().__init__(cpu_group, device, device_group, unique_name)

        self.dp_group: Optional[GroupCoordinator] = None
        self.dp_rank = 0
        self.dp_world_size = 1
        # assume EP is enabled along with DP
        if "ep" in unique_name:
            self.dp_group = get_dp_group()
            self.dp_rank = self.dp_group.rank_in_group
            self.dp_world_size = self.dp_group.world_size
            self.tp_group = get_tp_group()
        self.world_size = dist.get_world_size(group=self.cpu_group)
        self.rank = dist.get_rank(group=self.cpu_group)

    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
        # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge
        # occurring when PT_HPU_ENABLE_LAZY_COLLECTIVES=true env var is used
        # (which is required for tensor parallel HPUGraph inference)
        htorch.core.mark_step()
        dist.all_reduce(input_, group=self.device_group)
        return input_

    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
        world_size = self.world_size
        if dim < 0:
            # Convert negative dim to positive.
            dim += input_.dim()
        input_size = input_.size()
        # Allocate output tensor.
        output_tensor = torch.empty((world_size, ) + input_size, dtype=input_.dtype, device=input_.device)
        # All-gather.
        htorch.core.mark_step()
        dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
        # Reshape
        output_tensor = output_tensor.movedim(0, dim)
        output_tensor = output_tensor.reshape(input_size[:dim] + (world_size * input_size[dim], ) +
                                              input_size[dim + 1:])
        return output_tensor

    def dispatch(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        is_sequence_parallel: bool = False,
        extra_tensors: list[torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if extra_tensors is not None:
            raise NotImplementedError("extra_tensors is not supported for HPU")
        # Use dispatch_tensor in the plugin FusedMoEMethod for better performance
        return hidden_states, router_logits

    def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False) -> torch.Tensor:
        if htorch.utils.internal.is_lazy():
            htorch.core.mark_step()
        assert self.dp_group is not None
        assert hidden_states.dim() == 2, "Input hidden states must be 2D"

        dp_metadata = get_hpu_dp_metadata()
        if dp_metadata is not None:
            local_hidden_states = dp_metadata.local_hidden_states
        else:
            local_num_tokens = hidden_states.size(0) // self.world_size if is_sequence_parallel else hidden_states.size(
                0) // self.dp_world_size
            local_hidden_states = torch.empty((local_num_tokens, hidden_states.size(-1)),
                                              device=hidden_states.device,
                                              dtype=hidden_states.dtype)

        torch.distributed.reduce_scatter_tensor(
            local_hidden_states,
            hidden_states,
            group=get_ep_group().device_group if is_sequence_parallel else self.dp_group.device_group)
        hidden_states = local_hidden_states
        return hidden_states

dp_group instance-attribute 

dp_group: Optional[GroupCoordinator] = None

dp_rank instance-attribute 

dp_rank = 0

dp_world_size instance-attribute 

dp_world_size = 1

rank instance-attribute 

rank = get_rank(group=cpu_group)

tp_group instance-attribute 

tp_group = get_tp_group()

world_size instance-attribute 

world_size = get_world_size(group=cpu_group)

__init__ 

__init__(
    cpu_group: ProcessGroup,
    device: Optional[device] = None,
    device_group: Optional[ProcessGroup] = None,
    unique_name: str = "",
)
Source code in vllm_gaudi/distributed/device_communicators/hpu_communicator.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def __init__(self,
             cpu_group: ProcessGroup,
             device: Optional[torch.device] = None,
             device_group: Optional[ProcessGroup] = None,
             unique_name: str = ""):
    super().__init__(cpu_group, device, device_group, unique_name)

    self.dp_group: Optional[GroupCoordinator] = None
    self.dp_rank = 0
    self.dp_world_size = 1
    # assume EP is enabled along with DP
    if "ep" in unique_name:
        self.dp_group = get_dp_group()
        self.dp_rank = self.dp_group.rank_in_group
        self.dp_world_size = self.dp_group.world_size
        self.tp_group = get_tp_group()
    self.world_size = dist.get_world_size(group=self.cpu_group)
    self.rank = dist.get_rank(group=self.cpu_group)

all_gather 

all_gather(input_: Tensor, dim: int = -1) -> Tensor
Source code in vllm_gaudi/distributed/device_communicators/hpu_communicator.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
    world_size = self.world_size
    if dim < 0:
        # Convert negative dim to positive.
        dim += input_.dim()
    input_size = input_.size()
    # Allocate output tensor.
    output_tensor = torch.empty((world_size, ) + input_size, dtype=input_.dtype, device=input_.device)
    # All-gather.
    htorch.core.mark_step()
    dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
    # Reshape
    output_tensor = output_tensor.movedim(0, dim)
    output_tensor = output_tensor.reshape(input_size[:dim] + (world_size * input_size[dim], ) +
                                          input_size[dim + 1:])
    return output_tensor

all_reduce 

all_reduce(input_: Tensor) -> Tensor
Source code in vllm_gaudi/distributed/device_communicators/hpu_communicator.py 
38
39
40
41
42
43
44
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
    # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge
    # occurring when PT_HPU_ENABLE_LAZY_COLLECTIVES=true env var is used
    # (which is required for tensor parallel HPUGraph inference)
    htorch.core.mark_step()
    dist.all_reduce(input_, group=self.device_group)
    return input_

combine 

combine(
    hidden_states: Tensor,
    is_sequence_parallel: bool = False,
) -> Tensor
Source code in vllm_gaudi/distributed/device_communicators/hpu_communicator.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False) -> torch.Tensor:
    if htorch.utils.internal.is_lazy():
        htorch.core.mark_step()
    assert self.dp_group is not None
    assert hidden_states.dim() == 2, "Input hidden states must be 2D"

    dp_metadata = get_hpu_dp_metadata()
    if dp_metadata is not None:
        local_hidden_states = dp_metadata.local_hidden_states
    else:
        local_num_tokens = hidden_states.size(0) // self.world_size if is_sequence_parallel else hidden_states.size(
            0) // self.dp_world_size
        local_hidden_states = torch.empty((local_num_tokens, hidden_states.size(-1)),
                                          device=hidden_states.device,
                                          dtype=hidden_states.dtype)

    torch.distributed.reduce_scatter_tensor(
        local_hidden_states,
        hidden_states,
        group=get_ep_group().device_group if is_sequence_parallel else self.dp_group.device_group)
    hidden_states = local_hidden_states
    return hidden_states

dispatch 

dispatch(
    hidden_states: Tensor,
    router_logits: Tensor,
    is_sequence_parallel: bool = False,
    extra_tensors: list[Tensor] | None = None,
) -> tuple[Tensor, Tensor]
Source code in vllm_gaudi/distributed/device_communicators/hpu_communicator.py 
63
64
65
66
67
68
69
70
71
72
73
def dispatch(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
    is_sequence_parallel: bool = False,
    extra_tensors: list[torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    if extra_tensors is not None:
        raise NotImplementedError("extra_tensors is not supported for HPU")
    # Use dispatch_tensor in the plugin FusedMoEMethod for better performance
    return hidden_states, router_logits