mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Derived from vLLM; modified by DeepSeek.
|
|
|
|
from typing import Any, NamedTuple
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from deepinfer.model_executor.parallel_utils.parallel_state import PG, cp_broadcast_object
|
|
|
|
|
|
def gather(
|
|
input_: torch.Tensor,
|
|
group: dist.ProcessGroup,
|
|
dst: int = 0,
|
|
dim: int = -1,
|
|
) -> torch.Tensor | None:
|
|
"""Gather input_ from all ranks of group and concat along dim; returns None on non-dst ranks."""
|
|
if group.size() == 1:
|
|
return input_
|
|
assert -input_.dim() <= dim < input_.dim(), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
|
|
is_dst = group.rank() == dst
|
|
gather_list = [torch.empty_like(input_) for _ in range(group.size())] if is_dst else None
|
|
dist.gather(input_, gather_list, dst=dist.get_global_rank(group, dst), group=group)
|
|
return torch.cat(gather_list, dim=dim) if is_dst else None
|
|
|
|
|
|
def cp_broadcast(input_: torch.Tensor, src: int = 0) -> torch.Tensor:
|
|
"""Broadcast input_ in place across the CP group."""
|
|
cp_group = PG.context_parallel_group()
|
|
if cp_group.size() > 1:
|
|
dist.broadcast(input_, src=dist.get_global_rank(cp_group, src), group=cp_group)
|
|
return input_
|
|
|
|
|
|
def cp_broadcast_multi_tensor(inputs: list[torch.Tensor], src: int = 0) -> None:
|
|
"""Broadcast multiple tensors in place across the CP group."""
|
|
cp_group = PG.context_parallel_group()
|
|
if cp_group.size() == 1:
|
|
return
|
|
# TODO: Maybe use deepep multi_tensor_broadcast later.
|
|
for tensor in inputs:
|
|
dist.broadcast(tensor, src=dist.get_global_rank(cp_group, src), group=cp_group)
|
|
|
|
|
|
class TensorMetadata(NamedTuple):
|
|
dtype: torch.dtype
|
|
size: torch.Size
|
|
|
|
|
|
def broadcast_tensor_dict(
|
|
tensor_dict: dict[Any, Any] | None = None,
|
|
src: int = 0,
|
|
) -> dict[Any, Any]:
|
|
"""Broadcast a dict across the CP group: cuda tensors via NCCL, everything else via shm object broadcast."""
|
|
cp_group = PG.context_parallel_group()
|
|
if cp_group.size() == 1:
|
|
return tensor_dict
|
|
assert src == 0, f"Invalid src rank ({src})"
|
|
|
|
tensors = []
|
|
if cp_group.rank() == src:
|
|
assert isinstance(tensor_dict, dict), f"Expecting a dictionary, got {type(tensor_dict)}"
|
|
metadata_list = []
|
|
for key, value in tensor_dict.items():
|
|
if isinstance(value, torch.Tensor):
|
|
assert value.is_cuda, f"Tensor {key}: {value} is not on cuda. Currently we only support broadcasting tensors on cuda."
|
|
metadata_list.append((key, TensorMetadata(value.dtype, value.size())))
|
|
tensors.append(value)
|
|
else:
|
|
metadata_list.append((key, value))
|
|
cp_broadcast_object(metadata_list)
|
|
else:
|
|
tensor_dict = {}
|
|
for key, value in cp_broadcast_object(None):
|
|
if isinstance(value, TensorMetadata):
|
|
tensor_dict[key] = torch.empty(value.size, dtype=value.dtype, device="cuda")
|
|
tensors.append(tensor_dict[key])
|
|
else:
|
|
tensor_dict[key] = value
|
|
cp_broadcast_multi_tensor(tensors, src)
|
|
return tensor_dict
|