| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for weight transfer APIs via LLM class. |
| 4 | |
| 5 | These tests use a mock weight transfer engine to verify that the API |
| 6 | calls the correct methods with the right arguments, without requiring |
| 7 | actual NCCL communication. |
| 8 | """ |
| 9 | |
| 10 | import os |
| 11 | import weakref |
| 12 | from dataclasses import dataclass |
| 13 | from unittest.mock import patch |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | from vllm.config import WeightTransferConfig |
| 19 | from vllm.distributed.weight_transfer.base import ( |
| 20 | WeightTransferEngine, |
| 21 | WeightTransferInitInfo, |
| 22 | WeightTransferInitRequest, |
| 23 | WeightTransferUpdateInfo, |
| 24 | WeightTransferUpdateRequest, |
| 25 | ) |
| 26 | |
| 27 | from ...utils import create_new_process_for_each_test |
| 28 | |
| 29 | # Use a tiny model for fast testing |
| 30 | MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" |
| 31 | |
| 32 | |
| 33 | # --- Mock Weight Transfer Engine --- |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class MockInitInfo(WeightTransferInitInfo): |
| 38 | """Mock initialization info.""" |
| 39 | |
| 40 | test_param: str = "test" |
| 41 | |
| 42 | |
| 43 | @dataclass |
| 44 | class MockUpdateInfo(WeightTransferUpdateInfo): |
| 45 | """Mock update info.""" |
| 46 | |
| 47 | names: list[str] | None = None |
| 48 | dtype_names: list[str] | None = None |
| 49 | shapes: list[list[int]] | None = None |
| 50 | |
| 51 | |
| 52 | class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]): |
| 53 | """Mock weight transfer engine that tracks method calls.""" |
| 54 | |
| 55 | init_info_cls = MockInitInfo |
| 56 | update_info_cls = MockUpdateInfo |
| 57 | |
| 58 | # Class-level tracking for verification across processes |
| 59 | init_transfer_engine_called: bool = False |
| 60 | start_called: bool = False |
| 61 | receive_weights_called: bool = False |
| 62 | finish_called: bool = False |
| 63 | shutdown_called: bool = False |
| 64 | last_init_info: MockInitInfo | None = None |
| 65 | last_update_info: MockUpdateInfo | None = None |
| 66 | |
| 67 | def __init__(self, config, vllm_config, device, model): |
| 68 | super().__init__(config, vllm_config, device, model) |
| 69 | # Reset tracking on init |
| 70 | MockWeightTransferEngine.init_transfer_engine_called = False |
| 71 | MockWeightTransferEngine.start_called = False |
| 72 | MockWeightTransferEngine.receive_weights_called = False |
| 73 | MockWeightTransferEngine.finish_called = False |
| 74 | MockWeightTransferEngine.shutdown_called = False |
| 75 | MockWeightTransferEngine.last_init_info = None |
| 76 | MockWeightTransferEngine.last_update_info = None |
| 77 | |
| 78 | def init_transfer_engine(self, init_info: MockInitInfo) -> None: |
| 79 | MockWeightTransferEngine.init_transfer_engine_called = True |
| 80 | MockWeightTransferEngine.last_init_info = init_info |
| 81 | |
| 82 | def start_weight_update(self) -> None: |
| 83 | MockWeightTransferEngine.start_called = True |
| 84 | |
| 85 | def finish_weight_update(self) -> None: |
| 86 | MockWeightTransferEngine.finish_called = True |
| 87 | |
| 88 | def receive_weights(self, update_info: MockUpdateInfo) -> None: |
| 89 | MockWeightTransferEngine.receive_weights_called = True |
| 90 | MockWeightTransferEngine.last_update_info = update_info |
| 91 | |
| 92 | def shutdown(self) -> None: |
| 93 | MockWeightTransferEngine.shutdown_called = True |
| 94 | |
| 95 | |
| 96 | def mock_create_engine(config, vllm_config, device, model): |
| 97 | """Mock factory function that returns our mock engine.""" |
| 98 | return MockWeightTransferEngine(config, vllm_config, device, model) |
| 99 | |
| 100 | |
| 101 | # --- Tests --- |
| 102 | |
| 103 | |
| 104 | @create_new_process_for_each_test() |
| 105 | def test_get_world_size_tp1(vllm_runner): |
| 106 | """Test world_size is correctly configured for TP=1.""" |
| 107 | if torch.accelerator.device_count() < 1: |
| 108 | pytest.skip("Need at least 1 GPU for this test") |
| 109 | |
| 110 | with vllm_runner( |
| 111 | MODEL_NAME, |
| 112 | enforce_eager=True, |
| 113 | load_format="dummy", |
| 114 | tensor_parallel_size=1, |
| 115 | weight_transfer_config=WeightTransferConfig(backend="nccl"), |
| 116 | ) as runner: |
| 117 | world_size = runner.llm.llm_engine.vllm_config.parallel_config.world_size |
| 118 | assert world_size == 1 |
| 119 | |
| 120 | |
| 121 | @create_new_process_for_each_test() |
| 122 | def test_init_weight_transfer_engine_calls_engine(vllm_runner): |
| 123 | """Test that init_weight_transfer_engine calls the engine's |
| 124 | init_transfer_engine method.""" |
| 125 | if torch.accelerator.device_count() < 1: |
| 126 | pytest.skip("Need at least 1 GPU for this test") |
| 127 | |
| 128 | # Run in-process so mock.patch works (spawn won't inherit the mock) |
| 129 | os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| 130 | # Enable insecure serialization to allow pickling functions for collective_rpc |
| 131 | os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" |
| 132 | |
| 133 | with ( |
| 134 | patch( |
| 135 | "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", |
| 136 | mock_create_engine, |
| 137 | ), |
| 138 | vllm_runner( |
| 139 | MODEL_NAME, |
| 140 | enforce_eager=True, |
| 141 | load_format="dummy", |
| 142 | tensor_parallel_size=1, |
| 143 | weight_transfer_config=WeightTransferConfig(backend="nccl"), |
| 144 | ) as runner, |
| 145 | ): |
| 146 | llm = weakref.proxy(runner.llm) |
| 147 | |
| 148 | # Verify engine was created |
| 149 | def check_engine_exists(self): |
| 150 | return self.weight_transfer_engine is not None |
| 151 | |
| 152 | results = llm.collective_rpc(check_engine_exists) |
| 153 | assert all(results), "Weight transfer engine should be initialized" |
| 154 | |
| 155 | # Call init_weight_transfer_engine |
| 156 | llm.init_weight_transfer_engine( |
| 157 | WeightTransferInitRequest(init_info={"test_param": "hello"}) |
| 158 | ) |
| 159 | |
| 160 | # Verify init_transfer_engine was called on the engine |
| 161 | def check_init_called(self): |
| 162 | engine = self.weight_transfer_engine |
| 163 | return ( |
| 164 | engine.init_transfer_engine_called, |
| 165 | engine.last_init_info.test_param if engine.last_init_info else None, |
| 166 | ) |
| 167 | |
| 168 | results = llm.collective_rpc(check_init_called) |
| 169 | for called, param in results: |
| 170 | assert called, "init_transfer_engine should have been called" |
| 171 | assert param == "hello", f"Expected 'hello', got {param}" |
| 172 | |
| 173 | |
| 174 | @create_new_process_for_each_test() |
| 175 | def test_update_weights_calls_engine(vllm_runner): |
| 176 | """Test that update_weights calls the engine's receive_weights method.""" |
| 177 | if torch.accelerator.device_count() < 1: |
| 178 | pytest.skip("Need at least 1 GPU for this test") |
| 179 | |
| 180 | # Run in-process so mock.patch works (spawn won't inherit the mock) |
| 181 | os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| 182 | # Enable insecure serialization to allow pickling functions for collective_rpc |
| 183 | os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" |
| 184 | |
| 185 | with ( |
| 186 | patch( |
| 187 | "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", |
| 188 | mock_create_engine, |
| 189 | ), |
| 190 | vllm_runner( |
| 191 | MODEL_NAME, |
| 192 | enforce_eager=True, |
| 193 | load_format="dummy", |
| 194 | tensor_parallel_size=1, |
| 195 | weight_transfer_config=WeightTransferConfig(backend="nccl"), |
| 196 | ) as runner, |
| 197 | ): |
| 198 | llm = weakref.proxy(runner.llm) |
| 199 | |
| 200 | # First init the weight transfer |
| 201 | llm.init_weight_transfer_engine( |
| 202 | WeightTransferInitRequest(init_info={"test_param": "init"}) |
| 203 | ) |
| 204 | llm.start_weight_update() |
| 205 | |
| 206 | # Call update_weights |
| 207 | test_names = ["layer.weight", "layer.bias"] |
| 208 | test_dtypes = ["float32", "float32"] |
| 209 | test_shapes = [[10, 10], [10]] |
| 210 | |
| 211 | llm.update_weights( |
| 212 | WeightTransferUpdateRequest( |
| 213 | update_info={ |
| 214 | "names": test_names, |
| 215 | "dtype_names": test_dtypes, |
| 216 | "shapes": test_shapes, |
| 217 | } |
| 218 | ) |
| 219 | ) |
| 220 | |
| 221 | # Verify receive_weights was called with correct info |
| 222 | def check_update_called(self): |
| 223 | engine = self.weight_transfer_engine |
| 224 | if not engine.receive_weights_called: |
| 225 | return False, None, None, None |
| 226 | info = engine.last_update_info |
| 227 | return (True, info.names, info.dtype_names, info.shapes) |
| 228 | |
| 229 | results = llm.collective_rpc(check_update_called) |
| 230 | for called, names, dtypes, shapes in results: |
| 231 | assert called, "receive_weights should have been called" |
| 232 | assert names == test_names |
| 233 | assert dtypes == test_dtypes |
| 234 | assert shapes == test_shapes |
| 235 | |
| 236 | llm.finish_weight_update() |
| 237 | assert llm.get_weight_version() == "default" |
| 238 | |
| 239 | |
| 240 | @create_new_process_for_each_test() |
| 241 | def test_full_weight_transfer_flow(vllm_runner): |
| 242 | """Test the complete weight transfer flow: init -> start -> update -> finish.""" |
| 243 | if torch.accelerator.device_count() < 1: |
| 244 | pytest.skip("Need at least 1 GPU for this test") |
| 245 | |
| 246 | # Run in-process so mock.patch works (spawn won't inherit the mock) |
| 247 | os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| 248 | # Enable insecure serialization to allow pickling functions for collective_rpc |
| 249 | os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" |
| 250 | |
| 251 | with ( |
| 252 | patch( |
| 253 | "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", |
| 254 | mock_create_engine, |
| 255 | ), |
| 256 | vllm_runner( |
| 257 | MODEL_NAME, |
| 258 | enforce_eager=True, |
| 259 | load_format="dummy", |
| 260 | tensor_parallel_size=1, |
| 261 | weight_transfer_config=WeightTransferConfig(backend="nccl"), |
| 262 | ) as runner, |
| 263 | ): |
| 264 | llm = weakref.proxy(runner.llm) |
| 265 | |
| 266 | assert llm.get_weight_version() == "default" |
| 267 | |
| 268 | # Step 1: Initialize weight transfer engine |
| 269 | llm.init_weight_transfer_engine( |
| 270 | WeightTransferInitRequest(init_info={"test_param": "flow_test"}) |
| 271 | ) |
| 272 | |
| 273 | # Step 2: Start weight update |
| 274 | llm.start_weight_update() |
| 275 | |
| 276 | # Step 3: Update weights |
| 277 | llm.update_weights( |
| 278 | WeightTransferUpdateRequest( |
| 279 | update_info={ |
| 280 | "names": ["test.weight"], |
| 281 | "dtype_names": ["bfloat16"], |
| 282 | "shapes": [[100, 100]], |
| 283 | } |
| 284 | ) |
| 285 | ) |
| 286 | |
| 287 | assert llm.get_weight_version() == "default" |
| 288 | |
| 289 | # Step 4: Finish weight update |
| 290 | llm.finish_weight_update("step-42") |
| 291 | |
| 292 | assert llm.get_weight_version() == "step-42" |
| 293 | |
| 294 | llm.update_weight_version("manual-version") |
| 295 | assert llm.get_weight_version() == "manual-version" |
| 296 | |
| 297 | # Verify the full flow completed |
| 298 | def check_flow(self): |
| 299 | engine = self.weight_transfer_engine |
| 300 | return { |
| 301 | "init_called": engine.init_transfer_engine_called, |
| 302 | "update_called": engine.receive_weights_called, |
| 303 | "init_param": ( |
| 304 | engine.last_init_info.test_param if engine.last_init_info else None |
| 305 | ), |
| 306 | "update_names": ( |
| 307 | engine.last_update_info.names if engine.last_update_info else None |
| 308 | ), |
| 309 | } |
| 310 | |
| 311 | results = llm.collective_rpc(check_flow) |
| 312 | for result in results: |
| 313 | assert result["init_called"], "init_transfer_engine should be called" |
| 314 | assert result["update_called"], "receive_weights should be called" |
| 315 | assert result["init_param"] == "flow_test" |
| 316 | assert result["update_names"] == ["test.weight"] |
| 317 | |
| 318 | |
| 319 | @create_new_process_for_each_test() |
| 320 | def test_weight_transfer_config_backend(vllm_runner): |
| 321 | """Test that WeightTransferConfig backend is properly configured.""" |
| 322 | if torch.accelerator.device_count() < 1: |
| 323 | pytest.skip("Need at least 1 GPU for this test") |
| 324 | |
| 325 | # Test with nccl backend |
| 326 | with vllm_runner( |
| 327 | MODEL_NAME, |
| 328 | enforce_eager=True, |
| 329 | load_format="dummy", |
| 330 | tensor_parallel_size=1, |
| 331 | weight_transfer_config=WeightTransferConfig(backend="nccl"), |
| 332 | ) as runner: |
| 333 | config = runner.llm.llm_engine.vllm_config.weight_transfer_config |
| 334 | assert config is not None |
| 335 | assert config.backend == "nccl" |