wehub-resource-sync
1217 lines
40 KiB
Python
1217 lines
40 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for weight transfer engine backends.
|
|
|
|
Unit tests for engine classes (parsing, validation, registry).
|
|
Integration tests for NCCL and IPC weight transfer between processes using Ray.
|
|
"""
|
|
|
|
import pickle
|
|
from unittest.mock import MagicMock
|
|
|
|
import pybase64 as base64
|
|
import pytest
|
|
import ray
|
|
import torch
|
|
from torch.multiprocessing.reductions import reduce_tensor
|
|
|
|
from vllm.config.parallel import ParallelConfig
|
|
from vllm.config.weight_transfer import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
|
|
from vllm.distributed.weight_transfer.ipc_engine import (
|
|
IPCWeightTransferEngine,
|
|
IPCWeightTransferInitInfo,
|
|
IPCWeightTransferUpdateInfo,
|
|
)
|
|
from vllm.distributed.weight_transfer.nccl_engine import (
|
|
NCCLWeightTransferEngine,
|
|
NCCLWeightTransferInitInfo,
|
|
NCCLWeightTransferUpdateInfo,
|
|
)
|
|
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
|
|
SparseNCCLWeightTransferEngine,
|
|
SparseNCCLWeightTransferUpdateInfo,
|
|
SparseWeightPatch,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.network_utils import get_open_port
|
|
|
|
|
|
def _init_ray_for_weight_transfer() -> None:
|
|
if ray.is_initialized():
|
|
return
|
|
ray.init(
|
|
ignore_reinit_error=True,
|
|
runtime_env={
|
|
"env_vars": {
|
|
"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
|
|
"RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1",
|
|
"RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1",
|
|
}
|
|
},
|
|
)
|
|
|
|
|
|
def _get_ray_assigned_device() -> torch.device:
|
|
gpu_ids = ray.get_gpu_ids()
|
|
if not gpu_ids:
|
|
return torch.device("cuda:0")
|
|
return torch.device(f"cuda:{int(gpu_ids[0])}")
|
|
|
|
|
|
def _set_ray_assigned_device() -> torch.device:
|
|
device = _get_ray_assigned_device()
|
|
current_platform.set_device(device)
|
|
return device
|
|
|
|
|
|
def create_mock_parallel_config(
|
|
rank: int = 0,
|
|
world_size: int = 1,
|
|
dp_rank: int = 0,
|
|
) -> ParallelConfig:
|
|
"""Create a mock ParallelConfig for testing."""
|
|
config = MagicMock(spec=ParallelConfig)
|
|
config.rank = rank
|
|
config.world_size = world_size
|
|
config.data_parallel_rank = dp_rank
|
|
config.data_parallel_index = dp_rank
|
|
return config
|
|
|
|
|
|
def create_mock_vllm_config(
|
|
rank: int = 0,
|
|
world_size: int = 1,
|
|
dp_rank: int = 0,
|
|
) -> MagicMock:
|
|
"""Create a mock VllmConfig exposing parallel_config and model_config."""
|
|
vllm_config = MagicMock()
|
|
vllm_config.parallel_config = create_mock_parallel_config(rank, world_size, dp_rank)
|
|
vllm_config.model_config = MagicMock()
|
|
return vllm_config
|
|
|
|
|
|
# --- Unit Tests: NCCLWeightTransferUpdateInfo Validation ---
|
|
|
|
|
|
class TestNCCLWeightTransferUpdateInfoValidation:
|
|
"""Test NCCLWeightTransferUpdateInfo dataclass validation."""
|
|
|
|
def test_valid_update_info(self):
|
|
info = NCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32", "float32"],
|
|
shapes=[[10, 10], [10]],
|
|
)
|
|
assert info.names == ["layer.weight", "layer.bias"]
|
|
assert info.dtype_names == ["float32", "float32"]
|
|
assert info.shapes == [[10, 10], [10]]
|
|
|
|
def test_mismatched_dtype_names_raises(self):
|
|
with pytest.raises(ValueError, match="dtype_names"):
|
|
NCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32"], # Only one dtype
|
|
shapes=[[10, 10], [10]],
|
|
)
|
|
|
|
def test_mismatched_shapes_raises(self):
|
|
with pytest.raises(ValueError, match="shapes"):
|
|
NCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32", "float32"],
|
|
shapes=[[10, 10]], # Only one shape
|
|
)
|
|
|
|
def test_empty_lists_valid(self):
|
|
info = NCCLWeightTransferUpdateInfo(names=[], dtype_names=[], shapes=[])
|
|
assert len(info.names) == 0
|
|
|
|
|
|
# --- Unit Tests: SparseNCCLWeightTransferUpdateInfo Validation ---
|
|
|
|
|
|
class TestSparseNCCLWeightTransferUpdateInfoValidation:
|
|
"""Test SparseNCCLWeightTransferUpdateInfo dataclass validation."""
|
|
|
|
def test_valid_sparse_update_info(self):
|
|
info = SparseNCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32", "bfloat16"],
|
|
shapes=[[10, 10], [10]],
|
|
num_updates_list=[4, 2],
|
|
)
|
|
assert info.num_updates_list == [4, 2]
|
|
|
|
def test_mismatched_dtype_names_raises(self):
|
|
with pytest.raises(ValueError, match="dtype_names"):
|
|
SparseNCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10], [10]],
|
|
num_updates_list=[4, 2],
|
|
)
|
|
|
|
def test_rejects_empty_num_updates_list(self):
|
|
with pytest.raises(ValueError, match="cannot be empty"):
|
|
SparseNCCLWeightTransferUpdateInfo(
|
|
names=[],
|
|
dtype_names=[],
|
|
shapes=[],
|
|
num_updates_list=[],
|
|
)
|
|
|
|
def test_rejects_mismatched_num_updates(self):
|
|
with pytest.raises(ValueError, match="`num_updates_list`"):
|
|
SparseNCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32", "float32"],
|
|
shapes=[[10, 10], [10]],
|
|
num_updates_list=[3],
|
|
)
|
|
|
|
def test_rejects_negative_num_updates(self):
|
|
with pytest.raises(ValueError, match="non-negative"):
|
|
SparseNCCLWeightTransferUpdateInfo(
|
|
names=["layer.weight"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10]],
|
|
num_updates_list=[-1],
|
|
)
|
|
|
|
|
|
# --- Unit Tests: Engine Parsing ---
|
|
|
|
|
|
class TestNCCLEngineParsing:
|
|
"""Test NCCLWeightTransferEngine parsing methods."""
|
|
|
|
def _make_engine(self):
|
|
config = WeightTransferConfig(backend="nccl")
|
|
return NCCLWeightTransferEngine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
|
|
def test_parse_init_info_valid(self):
|
|
engine = self._make_engine()
|
|
init_info = engine.parse_init_info(
|
|
{
|
|
"master_address": "127.0.0.1",
|
|
"master_port": 12345,
|
|
"rank_offset": 1,
|
|
"world_size": 3,
|
|
}
|
|
)
|
|
assert isinstance(init_info, NCCLWeightTransferInitInfo)
|
|
assert init_info.master_address == "127.0.0.1"
|
|
assert init_info.master_port == 12345
|
|
assert init_info.rank_offset == 1
|
|
assert init_info.world_size == 3
|
|
|
|
def test_parse_init_info_missing_field_raises(self):
|
|
engine = self._make_engine()
|
|
with pytest.raises(ValueError, match="Invalid init_info"):
|
|
engine.parse_init_info({"master_address": "127.0.0.1"})
|
|
|
|
def test_parse_update_info_valid(self):
|
|
engine = self._make_engine()
|
|
update_info = engine.parse_update_info(
|
|
{
|
|
"names": ["w1", "w2"],
|
|
"dtype_names": ["float32", "bfloat16"],
|
|
"shapes": [[100, 100], [50]],
|
|
}
|
|
)
|
|
assert isinstance(update_info, NCCLWeightTransferUpdateInfo)
|
|
assert update_info.names == ["w1", "w2"]
|
|
assert update_info.dtype_names == ["float32", "bfloat16"]
|
|
assert update_info.shapes == [[100, 100], [50]]
|
|
|
|
|
|
# --- Unit Tests: Engine Registry ---
|
|
|
|
|
|
class TestEngineRegistry:
|
|
"""Test weight transfer engine registry."""
|
|
|
|
def test_create_engine_nccl(self):
|
|
config = WeightTransferConfig(backend="nccl")
|
|
engine = WeightTransferEngineFactory.create_engine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
assert isinstance(engine, NCCLWeightTransferEngine)
|
|
|
|
def test_create_engine_ipc(self):
|
|
config = WeightTransferConfig(backend="ipc")
|
|
engine = WeightTransferEngineFactory.create_engine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
assert isinstance(engine, IPCWeightTransferEngine)
|
|
|
|
def test_create_engine_sparse_nccl(self):
|
|
config = WeightTransferConfig(backend="sparse_nccl")
|
|
engine = WeightTransferEngineFactory.create_engine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
assert isinstance(engine, SparseNCCLWeightTransferEngine)
|
|
|
|
def test_create_engine_invalid_backend(self):
|
|
config = WeightTransferConfig(backend="invalid")
|
|
with pytest.raises(ValueError, match="Invalid weight transfer backend"):
|
|
WeightTransferEngineFactory.create_engine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
|
|
def test_register_duplicate_raises(self):
|
|
with pytest.raises(ValueError, match="already registered"):
|
|
WeightTransferEngineFactory.register_engine(
|
|
"nccl", NCCLWeightTransferEngine
|
|
)
|
|
|
|
|
|
# --- Unit Tests: Sparse patch application (CPU) ---
|
|
|
|
|
|
class TestSparseNCCLPatchApplication:
|
|
"""Test SparseNCCLWeightTransferEngine._apply_patch on a real param."""
|
|
|
|
def _make_engine(self, model):
|
|
config = WeightTransferConfig(backend="sparse_nccl")
|
|
return SparseNCCLWeightTransferEngine(
|
|
config, create_mock_vllm_config(), torch.device("cpu"), model
|
|
)
|
|
|
|
def _make_model(self, numel: int = 8):
|
|
model = torch.nn.Module()
|
|
model.register_parameter(
|
|
"w", torch.nn.Parameter(torch.zeros(numel), requires_grad=False)
|
|
)
|
|
|
|
def get_parameter(name):
|
|
assert name == "w"
|
|
return model.w
|
|
|
|
model.get_parameter = get_parameter
|
|
return model
|
|
|
|
def test_apply_patch_updates_only_selected_entries(self):
|
|
model = self._make_model(8)
|
|
engine = self._make_engine(model)
|
|
engine._apply_patch(
|
|
SparseWeightPatch(
|
|
name="w",
|
|
indices=torch.tensor([1, 3], dtype=torch.int32),
|
|
values=torch.tensor([5.0, 7.0], dtype=torch.float32),
|
|
)
|
|
)
|
|
expected = torch.zeros(8)
|
|
expected[1] = 5.0
|
|
expected[3] = 7.0
|
|
assert torch.equal(model.w.data, expected)
|
|
|
|
def test_apply_patch_rejects_mismatched_lengths(self):
|
|
model = self._make_model(8)
|
|
engine = self._make_engine(model)
|
|
with pytest.raises(ValueError, match="matching lengths"):
|
|
engine._apply_patch(
|
|
SparseWeightPatch(
|
|
name="w",
|
|
indices=torch.tensor([1, 3], dtype=torch.int32),
|
|
values=torch.tensor([5.0], dtype=torch.float32),
|
|
)
|
|
)
|
|
|
|
def test_apply_patch_rejects_non_int32_indices(self):
|
|
model = self._make_model(8)
|
|
engine = self._make_engine(model)
|
|
with pytest.raises(ValueError, match="int32 indices"):
|
|
engine._apply_patch(
|
|
SparseWeightPatch(
|
|
name="w",
|
|
indices=torch.tensor([1], dtype=torch.int64),
|
|
values=torch.tensor([5.0], dtype=torch.float32),
|
|
)
|
|
)
|
|
|
|
def test_apply_patch_rejects_dtype_mismatch(self):
|
|
model = self._make_model(8)
|
|
engine = self._make_engine(model)
|
|
with pytest.raises(ValueError, match="does not match"):
|
|
engine._apply_patch(
|
|
SparseWeightPatch(
|
|
name="w",
|
|
indices=torch.tensor([1], dtype=torch.int32),
|
|
values=torch.tensor([5.0], dtype=torch.bfloat16),
|
|
)
|
|
)
|
|
|
|
def test_apply_patch_rejects_non_contiguous_param(self):
|
|
model = torch.nn.Module()
|
|
model.register_parameter(
|
|
"w",
|
|
torch.nn.Parameter(
|
|
torch.arange(12, dtype=torch.float32).view(3, 4).t(),
|
|
requires_grad=False,
|
|
),
|
|
)
|
|
model.get_parameter = lambda name: model.w
|
|
engine = self._make_engine(model)
|
|
with pytest.raises(NotImplementedError, match="contiguous params"):
|
|
engine._apply_patch(
|
|
SparseWeightPatch(
|
|
name="w",
|
|
indices=torch.tensor([1], dtype=torch.int32),
|
|
values=torch.tensor([1.0], dtype=torch.float32),
|
|
)
|
|
)
|
|
|
|
|
|
# --- Test receive_weights without init raises ---
|
|
|
|
|
|
def test_nccl_receive_weights_without_init_raises():
|
|
"""Test that receive_weights raises if init_transfer_engine wasn't called."""
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
config = WeightTransferConfig(backend="nccl")
|
|
engine = NCCLWeightTransferEngine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
|
|
update_info = NCCLWeightTransferUpdateInfo(
|
|
names=["w"], dtype_names=["float32"], shapes=[[10]]
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized"):
|
|
engine.receive_weights(update_info)
|
|
|
|
|
|
def test_sparse_nccl_receive_weights_without_init_raises():
|
|
"""Test that sparse receive raises if init_transfer_engine wasn't called."""
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
config = WeightTransferConfig(backend="sparse_nccl")
|
|
engine = SparseNCCLWeightTransferEngine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
|
|
update_info = SparseNCCLWeightTransferUpdateInfo(
|
|
names=["w"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10]],
|
|
num_updates_list=[2],
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized"):
|
|
engine.receive_weights(update_info)
|
|
|
|
|
|
# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
def trainer_broadcast_tensor(
|
|
master_address: str,
|
|
master_port: int,
|
|
world_size: int,
|
|
tensor_shape: list[int],
|
|
tensor_dtype: str,
|
|
) -> bool:
|
|
"""Trainer task that broadcasts a tensor via NCCL."""
|
|
import torch
|
|
|
|
device = _set_ray_assigned_device()
|
|
|
|
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
|
|
from vllm.distributed.utils import StatelessProcessGroup
|
|
|
|
# Create process group as rank 0 (trainer)
|
|
pg = StatelessProcessGroup.create(
|
|
host=master_address,
|
|
port=master_port,
|
|
rank=0,
|
|
world_size=world_size,
|
|
)
|
|
comm = PyNcclCommunicator(pg, device=device.index)
|
|
|
|
# Create and broadcast the tensor
|
|
dtype = getattr(torch, tensor_dtype)
|
|
tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device=device)
|
|
comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream())
|
|
torch.accelerator.synchronize()
|
|
|
|
return True
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
def inference_receive_tensor(
|
|
master_address: str,
|
|
master_port: int,
|
|
world_size: int,
|
|
tensor_shape: list[int],
|
|
tensor_dtype: str,
|
|
) -> dict:
|
|
"""Inference task that receives tensor via NCCLWeightTransferEngine."""
|
|
import contextlib
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
_set_ray_assigned_device()
|
|
|
|
from vllm.config.parallel import ParallelConfig
|
|
from vllm.config.weight_transfer import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer.nccl_engine import (
|
|
NCCLWeightTransferEngine,
|
|
NCCLWeightTransferInitInfo,
|
|
NCCLWeightTransferUpdateInfo,
|
|
)
|
|
|
|
class Recorder(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.received = []
|
|
|
|
def load_weights(self, weights):
|
|
for name, tensor in weights:
|
|
self.received.append((name, tensor.clone()))
|
|
|
|
config = WeightTransferConfig(backend="nccl")
|
|
vllm_config = MagicMock()
|
|
parallel_config = MagicMock(spec=ParallelConfig)
|
|
parallel_config.rank = 0
|
|
parallel_config.world_size = 1
|
|
parallel_config.data_parallel_rank = 0
|
|
parallel_config.data_parallel_index = 0
|
|
vllm_config.parallel_config = parallel_config
|
|
vllm_config.model_config = MagicMock()
|
|
|
|
recorder = Recorder()
|
|
engine = NCCLWeightTransferEngine(
|
|
config, vllm_config, torch.device("cuda"), recorder
|
|
)
|
|
# Transport-only test: bypass the set_current_vllm_config context that
|
|
# receive_weights enters, since vllm_config here is a mock.
|
|
import vllm.config as _vllm_config_mod
|
|
|
|
_vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
|
|
|
|
# Initialize the engine (joins as rank 1)
|
|
init_info = NCCLWeightTransferInitInfo(
|
|
master_address=master_address,
|
|
master_port=master_port,
|
|
rank_offset=1, # Trainer is rank 0, we become rank 1
|
|
world_size=world_size,
|
|
)
|
|
engine.init_transfer_engine(init_info)
|
|
|
|
update_info = NCCLWeightTransferUpdateInfo(
|
|
names=["test.weight"],
|
|
dtype_names=[tensor_dtype],
|
|
shapes=[tensor_shape],
|
|
)
|
|
engine.receive_weights(update_info)
|
|
torch.accelerator.synchronize()
|
|
|
|
# Verify we received the tensor
|
|
success = False
|
|
received_shape = None
|
|
received_sum = None
|
|
|
|
if len(recorder.received) == 1:
|
|
name, tensor = recorder.received[0]
|
|
received_shape = list(tensor.shape)
|
|
received_sum = tensor.sum().item()
|
|
if received_shape == tensor_shape:
|
|
expected_sum = 1.0 * torch.tensor(tensor_shape).prod().item()
|
|
if abs(received_sum - expected_sum) < 0.01:
|
|
success = True
|
|
|
|
engine.shutdown()
|
|
|
|
return {
|
|
"success": success,
|
|
"received_shape": received_shape,
|
|
"received_sum": received_sum,
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
torch.accelerator.device_count() < 2,
|
|
reason="Need at least 2 GPUs to run NCCL weight transfer test.",
|
|
)
|
|
def test_nccl_weight_transfer_between_processes():
|
|
"""Test NCCL weight transfer from trainer to inference process using Ray.
|
|
|
|
This test verifies that the NCCLWeightTransferEngine can receive
|
|
tensors broadcast by a trainer process via NCCL.
|
|
"""
|
|
_init_ray_for_weight_transfer()
|
|
|
|
master_address = "127.0.0.1"
|
|
master_port = get_open_port()
|
|
world_size = 2 # 1 trainer + 1 inference worker
|
|
|
|
tensor_shape = [100, 100]
|
|
tensor_dtype = "float32"
|
|
|
|
inference_future = inference_receive_tensor.remote(
|
|
master_address, master_port, world_size, tensor_shape, tensor_dtype
|
|
)
|
|
trainer_future = trainer_broadcast_tensor.remote(
|
|
master_address, master_port, world_size, tensor_shape, tensor_dtype
|
|
)
|
|
|
|
trainer_result, result = ray.get([trainer_future, inference_future])
|
|
|
|
assert trainer_result, "Trainer should complete successfully"
|
|
assert result["success"], (
|
|
f"Weight transfer failed. "
|
|
f"Received shape: {result['received_shape']}, "
|
|
f"Received sum: {result['received_sum']}"
|
|
)
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
def trainer_broadcast_sparse_tensor(
|
|
master_address: str,
|
|
master_port: int,
|
|
world_size: int,
|
|
) -> bool:
|
|
"""Trainer task that broadcasts sparse patches via NCCL."""
|
|
import torch
|
|
|
|
device = _set_ray_assigned_device()
|
|
|
|
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
|
|
from vllm.distributed.utils import StatelessProcessGroup
|
|
from vllm.distributed.weight_transfer.nccl_engine import (
|
|
NCCLTrainerSendWeightsArgs,
|
|
)
|
|
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
|
|
SparseNCCLWeightTransferEngine,
|
|
SparseWeightPatch,
|
|
)
|
|
|
|
pg = StatelessProcessGroup.create(
|
|
host=master_address,
|
|
port=master_port,
|
|
rank=0,
|
|
world_size=world_size,
|
|
)
|
|
comm = PyNcclCommunicator(pg, device=device.index)
|
|
|
|
patch = SparseWeightPatch(
|
|
name="test.weight",
|
|
indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device),
|
|
values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device),
|
|
)
|
|
SparseNCCLWeightTransferEngine.trainer_send_weights(
|
|
iter([patch]),
|
|
NCCLTrainerSendWeightsArgs(group=comm),
|
|
)
|
|
torch.accelerator.synchronize()
|
|
return True
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
def inference_receive_sparse_tensor(
|
|
master_address: str,
|
|
master_port: int,
|
|
world_size: int,
|
|
) -> dict:
|
|
"""Inference task that receives sparse patches via the sparse engine."""
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
device = _set_ray_assigned_device()
|
|
|
|
from vllm.config.parallel import ParallelConfig
|
|
from vllm.config.weight_transfer import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
|
|
SparseNCCLWeightTransferEngine,
|
|
SparseNCCLWeightTransferUpdateInfo,
|
|
)
|
|
|
|
config = WeightTransferConfig(backend="sparse_nccl")
|
|
vllm_config = MagicMock()
|
|
parallel_config = MagicMock(spec=ParallelConfig)
|
|
parallel_config.rank = 0
|
|
parallel_config.world_size = 1
|
|
parallel_config.data_parallel_rank = 0
|
|
parallel_config.data_parallel_index = 0
|
|
vllm_config.parallel_config = parallel_config
|
|
vllm_config.model_config = MagicMock()
|
|
|
|
# Real module holding the target parameter the patch will modify.
|
|
model = torch.nn.Module()
|
|
model.register_parameter(
|
|
"w", torch.nn.Parameter(torch.zeros(30, device="cuda"), requires_grad=False)
|
|
)
|
|
model.get_parameter = lambda name: model.w
|
|
|
|
update_info = SparseNCCLWeightTransferUpdateInfo(
|
|
names=["w"],
|
|
dtype_names=["float32"],
|
|
shapes=[[30]],
|
|
num_updates_list=[3],
|
|
)
|
|
|
|
engine = SparseNCCLWeightTransferEngine(
|
|
config, vllm_config, torch.device("cuda"), model
|
|
)
|
|
from vllm.distributed.weight_transfer.nccl_common import (
|
|
NCCLWeightTransferInitInfo,
|
|
)
|
|
|
|
engine.init_transfer_engine(
|
|
NCCLWeightTransferInitInfo(
|
|
master_address=master_address,
|
|
master_port=master_port,
|
|
rank_offset=1,
|
|
world_size=world_size,
|
|
)
|
|
)
|
|
engine.receive_weights(update_info)
|
|
torch.accelerator.synchronize()
|
|
|
|
expected = torch.zeros(30, dtype=torch.float32, device=device)
|
|
expected[[1, 7, 25]] = torch.tensor(
|
|
[10.0, 20.0, 30.0], dtype=torch.float32, device=device
|
|
)
|
|
success = torch.equal(model.w.data, expected)
|
|
engine.shutdown()
|
|
return {
|
|
"success": success,
|
|
"selected_values": model.w.data[[1, 7, 25]].cpu().tolist(),
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
torch.accelerator.device_count() < 2,
|
|
reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
|
|
)
|
|
def test_nccl_sparse_weight_transfer_between_processes():
|
|
"""Test NCCL sparse weight transfer from trainer to inference process."""
|
|
_init_ray_for_weight_transfer()
|
|
|
|
master_address = "127.0.0.1"
|
|
master_port = get_open_port()
|
|
world_size = 2
|
|
|
|
inference_future = inference_receive_sparse_tensor.remote(
|
|
master_address, master_port, world_size
|
|
)
|
|
trainer_future = trainer_broadcast_sparse_tensor.remote(
|
|
master_address, master_port, world_size
|
|
)
|
|
|
|
trainer_result, result = ray.get([trainer_future, inference_future])
|
|
|
|
assert trainer_result, "Trainer should complete successfully"
|
|
assert result["success"], (
|
|
"Sparse weight transfer failed. "
|
|
f"Received selected values: {result['selected_values']}"
|
|
)
|
|
|
|
|
|
# --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
|
|
|
|
|
|
class TestIPCWeightTransferUpdateInfoValidation:
|
|
"""Test IPCWeightTransferUpdateInfo dataclass validation."""
|
|
|
|
def test_valid_update_info(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
_, ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}]
|
|
|
|
info = IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10]],
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
assert info.names == ["layer.weight"]
|
|
assert info.dtype_names == ["float32"]
|
|
assert info.shapes == [[10, 10]]
|
|
assert len(info.ipc_handles) == 1
|
|
|
|
def test_mismatched_dtype_names_raises(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
_, ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
|
|
|
|
with pytest.raises(ValueError, match="dtype_names"):
|
|
IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32"], # Only one dtype
|
|
shapes=[[10, 10], [10]],
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
|
|
def test_mismatched_shapes_raises(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
_, ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
|
|
|
|
with pytest.raises(ValueError, match="shapes"):
|
|
IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32", "float32"],
|
|
shapes=[[10, 10]], # Only one shape
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
|
|
def test_mismatched_ipc_handles_raises(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
_, ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}] # Only one handle
|
|
|
|
with pytest.raises(ValueError, match="ipc_handles"):
|
|
IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight", "layer.bias"],
|
|
dtype_names=["float32", "float32"],
|
|
shapes=[[10, 10], [10]],
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
|
|
def test_valid_update_info_from_pickled(self, monkeypatch):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}]
|
|
|
|
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
|
|
|
info = IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10]],
|
|
ipc_handles_pickled=pickled,
|
|
)
|
|
assert info.ipc_handles == ipc_handles
|
|
assert info.ipc_handles_pickled is None
|
|
|
|
def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
|
|
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
|
|
|
|
with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
|
|
IPCWeightTransferUpdateInfo(
|
|
names=[],
|
|
dtype_names=[],
|
|
shapes=[],
|
|
ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
|
|
)
|
|
|
|
def test_both_handles_and_pickled_raises(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}]
|
|
|
|
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
|
|
|
with pytest.raises(ValueError, match="Cannot specify both"):
|
|
IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10]],
|
|
ipc_handles=ipc_handles,
|
|
ipc_handles_pickled=pickled,
|
|
)
|
|
|
|
def test_neither_handles_nor_pickled_raises(self):
|
|
with pytest.raises(ValueError, match="must be provided"):
|
|
IPCWeightTransferUpdateInfo(
|
|
names=["layer.weight"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10]],
|
|
)
|
|
|
|
def test_empty_lists_valid(self):
|
|
info = IPCWeightTransferUpdateInfo(
|
|
names=[],
|
|
dtype_names=[],
|
|
shapes=[],
|
|
ipc_handles=[],
|
|
)
|
|
assert len(info.names) == 0
|
|
|
|
|
|
# --- Unit Tests: IPC Engine Parsing ---
|
|
|
|
|
|
class TestIPCEngineParsing:
|
|
"""Test IPCWeightTransferEngine parsing methods."""
|
|
|
|
def _make_engine(self):
|
|
config = WeightTransferConfig(backend="ipc")
|
|
return IPCWeightTransferEngine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
|
|
def test_parse_update_info_valid(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
engine = self._make_engine()
|
|
|
|
dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
|
|
dummy_tensor2 = torch.ones(50, device="cuda:0")
|
|
_, ipc_args1 = reduce_tensor(dummy_tensor1)
|
|
_, ipc_args2 = reduce_tensor(dummy_tensor2)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]
|
|
|
|
update_info = engine.parse_update_info(
|
|
{
|
|
"names": ["w1", "w2"],
|
|
"dtype_names": ["float32", "bfloat16"],
|
|
"shapes": [[100, 100], [50]],
|
|
"ipc_handles": ipc_handles,
|
|
}
|
|
)
|
|
|
|
assert isinstance(update_info, IPCWeightTransferUpdateInfo)
|
|
assert update_info.names == ["w1", "w2"]
|
|
assert update_info.dtype_names == ["float32", "bfloat16"]
|
|
assert update_info.shapes == [[100, 100], [50]]
|
|
assert len(update_info.ipc_handles) == 2
|
|
|
|
def test_parse_update_info_pickled(self, monkeypatch):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
|
|
|
engine = self._make_engine()
|
|
|
|
dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
|
|
dummy_tensor2 = torch.ones(50, device="cuda:0")
|
|
_, ipc_args1 = reduce_tensor(dummy_tensor1)
|
|
_, ipc_args2 = reduce_tensor(dummy_tensor2)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]
|
|
|
|
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
|
|
|
update_info = engine.parse_update_info(
|
|
{
|
|
"names": ["w1", "w2"],
|
|
"dtype_names": ["float32", "bfloat16"],
|
|
"shapes": [[100, 100], [50]],
|
|
"ipc_handles_pickled": pickled,
|
|
}
|
|
)
|
|
|
|
assert isinstance(update_info, IPCWeightTransferUpdateInfo)
|
|
assert update_info.names == ["w1", "w2"]
|
|
assert len(update_info.ipc_handles) == 2
|
|
assert gpu_uuid in update_info.ipc_handles[0]
|
|
assert gpu_uuid in update_info.ipc_handles[1]
|
|
|
|
def test_parse_update_info_ignores_none_pickled_handles(self):
|
|
engine = self._make_engine()
|
|
ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
|
|
|
|
update_info = engine.parse_update_info(
|
|
{
|
|
"names": ["w1"],
|
|
"dtype_names": ["float32"],
|
|
"shapes": [[1]],
|
|
"ipc_handles": ipc_handles,
|
|
"ipc_handles_pickled": None,
|
|
}
|
|
)
|
|
|
|
assert isinstance(update_info, IPCWeightTransferUpdateInfo)
|
|
assert update_info.ipc_handles == ipc_handles
|
|
|
|
def test_parse_update_info_both_handles_and_pickled_raises(self):
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
engine = self._make_engine()
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
_, ipc_handle = reduce_tensor(dummy_tensor)
|
|
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
|
ipc_handles = [{gpu_uuid: ipc_handle}]
|
|
|
|
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
|
|
|
with pytest.raises(ValueError, match="Cannot specify both"):
|
|
engine.parse_update_info(
|
|
{
|
|
"names": ["layer.weight"],
|
|
"dtype_names": ["float32"],
|
|
"shapes": [[10, 10]],
|
|
"ipc_handles": ipc_handles,
|
|
"ipc_handles_pickled": pickled,
|
|
}
|
|
)
|
|
|
|
|
|
# --- Integration Test: IPC Weight Transfer Between Ray Tasks ---
|
|
|
|
|
|
def get_physical_gpu_id(device_index: int = 0) -> str:
|
|
"""Get physical GPU UUID for a device."""
|
|
props = torch.cuda.get_device_properties(device_index)
|
|
return str(props.uuid)
|
|
|
|
|
|
@ray.remote(num_gpus=0.5)
|
|
class TrainerActor:
|
|
"""Trainer actor that creates and holds CUDA IPC handles."""
|
|
|
|
def __init__(self, tensor_shape: list[int], tensor_dtype: str):
|
|
device = _set_ray_assigned_device()
|
|
|
|
# Create tensor on GPU and keep it alive
|
|
dtype = getattr(torch, tensor_dtype)
|
|
self.tensor = torch.ones(tensor_shape, dtype=dtype, device=device)
|
|
self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer
|
|
|
|
_, ipc_args = reduce_tensor(self.tensor)
|
|
gpu_uuid = get_physical_gpu_id(device.index)
|
|
|
|
torch.accelerator.synchronize()
|
|
|
|
self.ipc_handle_dict = {
|
|
"ipc_handle": ipc_args,
|
|
"gpu_uuid": gpu_uuid,
|
|
"shape": tensor_shape,
|
|
"dtype": tensor_dtype,
|
|
}
|
|
|
|
def get_ipc_handle_dict(self) -> dict:
|
|
"""Return IPC handle dict. Tensor stays alive in this actor."""
|
|
return self.ipc_handle_dict
|
|
|
|
|
|
@ray.remote(num_gpus=0.5)
|
|
def inference_receive_ipc_tensor(
|
|
ipc_handle_dict: dict,
|
|
mode: str = "ray",
|
|
) -> dict:
|
|
"""Inference task that receives tensor via IPCWeightTransferEngine."""
|
|
import contextlib
|
|
import os
|
|
|
|
# Worker-side: ipc_handles_pickled is deserialized via pickle.
|
|
if mode == "http":
|
|
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
_set_ray_assigned_device()
|
|
|
|
from vllm.config.parallel import ParallelConfig
|
|
from vllm.config.weight_transfer import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer.ipc_engine import (
|
|
IPCWeightTransferEngine,
|
|
)
|
|
|
|
class Recorder(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.received = []
|
|
|
|
def load_weights(self, weights):
|
|
for name, tensor in weights:
|
|
self.received.append((name, tensor.clone()))
|
|
|
|
config = WeightTransferConfig(backend="ipc")
|
|
vllm_config = MagicMock()
|
|
parallel_config = MagicMock(spec=ParallelConfig)
|
|
parallel_config.rank = 0
|
|
parallel_config.world_size = 1
|
|
parallel_config.data_parallel_rank = 0
|
|
parallel_config.data_parallel_index = 0
|
|
vllm_config.parallel_config = parallel_config
|
|
vllm_config.model_config = MagicMock()
|
|
|
|
recorder = Recorder()
|
|
engine = IPCWeightTransferEngine(
|
|
config, vllm_config, _get_ray_assigned_device(), recorder
|
|
)
|
|
# Transport-only test: bypass the set_current_vllm_config context that
|
|
# receive_weights enters, since vllm_config here is a mock.
|
|
import vllm.config as _vllm_config_mod
|
|
|
|
_vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
|
|
|
|
init_info = IPCWeightTransferInitInfo()
|
|
engine.init_transfer_engine(init_info)
|
|
|
|
ipc_handles = [{ipc_handle_dict["gpu_uuid"]: ipc_handle_dict["ipc_handle"]}]
|
|
|
|
if mode == "ray":
|
|
update_dict: dict = {
|
|
"names": ["test.weight"],
|
|
"dtype_names": [ipc_handle_dict["dtype"]],
|
|
"shapes": [ipc_handle_dict["shape"]],
|
|
"ipc_handles": ipc_handles,
|
|
}
|
|
elif mode == "http":
|
|
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
|
|
update_dict = {
|
|
"names": ["test.weight"],
|
|
"dtype_names": [ipc_handle_dict["dtype"]],
|
|
"shapes": [ipc_handle_dict["shape"]],
|
|
"ipc_handles_pickled": pickled,
|
|
}
|
|
else:
|
|
raise ValueError(f"Unknown mode: {mode}")
|
|
|
|
update_info = engine.parse_update_info(update_dict)
|
|
engine.receive_weights(update_info)
|
|
torch.accelerator.synchronize()
|
|
|
|
success = False
|
|
received_shape = None
|
|
received_sum = None
|
|
|
|
if len(recorder.received) == 1:
|
|
name, tensor = recorder.received[0]
|
|
received_shape = list(tensor.shape)
|
|
received_sum = tensor.sum().item()
|
|
if received_shape == ipc_handle_dict["shape"]:
|
|
expected_sum = 42.0 * torch.tensor(ipc_handle_dict["shape"]).prod().item()
|
|
if abs(received_sum - expected_sum) < 0.01:
|
|
success = True
|
|
|
|
engine.shutdown()
|
|
|
|
return {
|
|
"success": success,
|
|
"received_shape": received_shape,
|
|
"received_sum": received_sum,
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
torch.accelerator.device_count() < 1,
|
|
reason="Need at least 1 GPU to run IPC weight transfer test.",
|
|
)
|
|
@pytest.mark.parametrize("mode", ["ray", "http"])
|
|
def test_ipc_weight_transfer_between_processes(mode: str):
|
|
"""Test IPC weight transfer from trainer to inference process using Ray."""
|
|
from ray.util.placement_group import placement_group
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
|
|
_init_ray_for_weight_transfer()
|
|
|
|
pg = placement_group([{"GPU": 1, "CPU": 2}])
|
|
ray.get(pg.ready())
|
|
|
|
scheduling_strategy = PlacementGroupSchedulingStrategy(
|
|
placement_group=pg,
|
|
placement_group_capture_child_tasks=True,
|
|
)
|
|
|
|
tensor_shape = [100, 100]
|
|
tensor_dtype = "float32"
|
|
|
|
trainer_actor = TrainerActor.options( # type: ignore[attr-defined]
|
|
scheduling_strategy=scheduling_strategy
|
|
).remote(tensor_shape, tensor_dtype)
|
|
|
|
ipc_handle_dict = ray.get(trainer_actor.get_ipc_handle_dict.remote())
|
|
|
|
inference_result = ray.get(
|
|
inference_receive_ipc_tensor.options(
|
|
scheduling_strategy=scheduling_strategy
|
|
).remote(ipc_handle_dict, mode=mode)
|
|
)
|
|
|
|
assert inference_result["success"], (
|
|
f"IPC weight transfer failed (mode={mode}). "
|
|
f"Received shape: {inference_result['received_shape']}, "
|
|
f"Received sum: {inference_result['received_sum']}"
|
|
)
|
|
|
|
|
|
def test_ipc_receive_weights_missing_gpu_uuid_raises():
|
|
"""Test that receive_weights raises if GPU UUID not found in IPC handles."""
|
|
if torch.accelerator.device_count() < 1:
|
|
pytest.skip("Need at least 1 GPU for this test")
|
|
|
|
config = WeightTransferConfig(backend="ipc")
|
|
engine = IPCWeightTransferEngine(
|
|
config,
|
|
create_mock_vllm_config(),
|
|
torch.device("cuda:0"),
|
|
MagicMock(spec=torch.nn.Module),
|
|
)
|
|
|
|
dummy_tensor = torch.ones(10, 10, device="cuda:0")
|
|
_, ipc_handle = reduce_tensor(dummy_tensor)
|
|
wrong_uuid = "wrong-uuid-12345"
|
|
ipc_handles = [{wrong_uuid: ipc_handle}]
|
|
|
|
update_info = IPCWeightTransferUpdateInfo(
|
|
names=["w"],
|
|
dtype_names=["float32"],
|
|
shapes=[[10, 10]],
|
|
ipc_handles=ipc_handles,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="IPC handle not found"):
|
|
engine.receive_weights(update_info)
|