"vllm/tokenizers/detokenizer_utils.py" did not exist on "d054da1992175787f936d18aead51bef663a0399"
communication_op.py 1.21 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from typing import Any, Dict, Optional, Union
4

5
import torch
6
import torch.distributed
7

8
from .parallel_state import get_tp_group
9
10


11
def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
12
13
    """All-reduce the input tensor across model parallel group."""
    return get_tp_group().all_reduce(input_)
14
15


16
17
def tensor_model_parallel_all_gather(input_: torch.Tensor,
                                     dim: int = -1) -> torch.Tensor:
18
    """All-gather the input tensor across model parallel group."""
19
    return get_tp_group().all_gather(input_, dim)
20
21


22
23
def tensor_model_parallel_gather(input_: torch.Tensor,
                                 dst: int = 0,
24
                                 dim: int = -1) -> Optional[torch.Tensor]:
25
26
    """Gather the input tensor across model parallel group."""
    return get_tp_group().gather(input_, dst, dim)
27
28


29
30
31
32
def broadcast_tensor_dict(tensor_dict: Optional[Dict[Any, Union[torch.Tensor,
                                                                Any]]] = None,
                          src: int = 0):
    if not torch.distributed.is_initialized():
33
        return tensor_dict
34
    return get_tp_group().broadcast_tensor_dict(tensor_dict, src)