| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.distributed import ( |
| 7 | get_tensor_model_parallel_rank, |
| 8 | get_tensor_model_parallel_world_size, |
| 9 | get_tp_group, |
| 10 | tensor_model_parallel_all_gather, |
| 11 | tensor_model_parallel_reduce_scatter, |
| 12 | ) |
| 13 | |
| 14 | |
| 15 | def _custom_collective(name: str, x: torch.Tensor) -> torch.Tensor | None: |
| 16 | device_communicator = get_tp_group().device_communicator |
| 17 | if device_communicator is None: |
| 18 | return None |
| 19 | collective = getattr(device_communicator, name, None) |
| 20 | return None if collective is None else collective(x) |
| 21 | |
| 22 | |
| 23 | def sp_all_gather(x: torch.Tensor) -> torch.Tensor: |
| 24 | output = _custom_collective("custom_all_gather", x) |
| 25 | if output is not None: |
| 26 | return output |
| 27 | return tensor_model_parallel_all_gather(x, 0) |
| 28 | |
| 29 | |
| 30 | def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor: |
| 31 | assert x.ndim == 2 |
| 32 | tp_size = get_tensor_model_parallel_world_size() |
| 33 | sp_pad = (-x.shape[0]) % tp_size |
| 34 | if sp_pad > 0: |
| 35 | x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) |
| 36 | output = _custom_collective("custom_reduce_scatter", x) |
| 37 | if output is not None: |
| 38 | return output |
| 39 | return tensor_model_parallel_reduce_scatter(x, 0) |
| 40 | |
| 41 | |
| 42 | def sp_shard(x: torch.Tensor) -> torch.Tensor: |
| 43 | tp_size = get_tensor_model_parallel_world_size() |
| 44 | tp_rank = get_tensor_model_parallel_rank() |
| 45 | sp_pad = (-x.shape[0]) % tp_size |
| 46 | if sp_pad > 0: |
| 47 | pad = (0, 0) * (x.ndim - 1) + (0, sp_pad) |
| 48 | x = torch.nn.functional.pad(x, pad) |
| 49 | chunk = x.shape[0] // tp_size |
| 50 | return x[tp_rank * chunk : (tp_rank + 1) * chunk] |
| 51 | |
| 52 | |
| 53 | def sp_padding_mask( |
| 54 | is_padding: torch.Tensor | None, |
| 55 | hidden_states: torch.Tensor, |
| 56 | ) -> torch.Tensor: |
| 57 | num_tokens = hidden_states.shape[0] |
| 58 | if is_padding is None: |
| 59 | is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool) |
| 60 | assert is_padding.shape[0] == num_tokens |
| 61 | |
| 62 | tp_size = get_tensor_model_parallel_world_size() |
| 63 | sp_pad = (-num_tokens) % tp_size |
| 64 | if sp_pad > 0: |
| 65 | is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True) |
| 66 | chunk = is_padding.shape[0] // tp_size |
| 67 | tp_rank = get_tensor_model_parallel_rank() |
| 68 | return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk] |