yangzhuxinyzx
58beb3433b
This reverts commit f4c3451a2b.
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.distributed.parallel_state import GroupCoordinator
|
|
|
|
|
|
class _DummyWork:
|
|
def wait(self) -> None:
|
|
pass
|
|
|
|
|
|
def _group(rank_in_group: int, world_size: int = 2) -> GroupCoordinator:
|
|
group = GroupCoordinator.__new__(GroupCoordinator)
|
|
group.world_size = world_size
|
|
group.rank_in_group = rank_in_group
|
|
group.ranks = list(range(world_size))
|
|
group.use_cpu_custom_send_recv = False
|
|
group.device_group = None
|
|
group.cpu_group = None
|
|
return group
|
|
|
|
|
|
def test_static_tensor_dict_transfer_uses_caller_owned_buffers(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
calls: list[tuple[str, int, torch.Tensor]] = []
|
|
|
|
def fake_isend(tensor: torch.Tensor, *, dst: int, group: Any) -> _DummyWork:
|
|
del group
|
|
calls.append(("send", dst, tensor))
|
|
return _DummyWork()
|
|
|
|
def fake_irecv(tensor: torch.Tensor, *, src: int, group: Any) -> _DummyWork:
|
|
del group
|
|
tensor.fill_(7)
|
|
calls.append(("recv", src, tensor))
|
|
return _DummyWork()
|
|
|
|
monkeypatch.setattr(torch.distributed, "isend", fake_isend)
|
|
monkeypatch.setattr(torch.distributed, "irecv", fake_irecv)
|
|
|
|
sent = torch.arange(8, dtype=torch.float32)
|
|
send_handles = _group(0).isend_tensor_dict_static({"hidden_states": sent})
|
|
|
|
received = torch.empty_like(sent)
|
|
recv_handles = _group(1).irecv_tensor_dict_static({"hidden_states": received})
|
|
|
|
assert len(send_handles) == 1
|
|
assert len(recv_handles) == 1
|
|
assert calls == [("send", 1, sent), ("recv", 0, received)]
|
|
torch.testing.assert_close(received, torch.full_like(received, 7))
|
|
|
|
|
|
def test_static_tensor_dict_transfer_rejects_non_tensors() -> None:
|
|
with pytest.raises(TypeError, match="only accepts tensors"):
|
|
_group(0).isend_tensor_dict_static( # type: ignore[arg-type]
|
|
{"hidden_states": object()}
|
|
)
|