@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 218 lines 6.81 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for GPUWorker weight-transfer pass-through behavior.
4
5 The worker no longer contains transport, layerwise, or sparse logic: it only
6 delegates to the configured weight transfer engine and tracks whether an update
7 session is active. These tests verify that delegation and the session guard.
8 """
9
10 import pytest
11 import torch
12 import torch.nn as nn
13
14 from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config
15 from vllm.lora.layers import BaseLayerWithLoRA
16 from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload
17 from vllm.v1.worker.gpu_worker import Worker
18
19
20 class _RecordingEngine:
21 """Minimal stand-in for a weight transfer engine."""
22
23 def __init__(self, raise_on_update: bool = False):
24 self.raise_on_update = raise_on_update
25 self.started = False
26 self.finished = False
27 self.reset_count = 0
28 self.supports_draft_weight_update = False
29 self.update_calls: list[dict] = []
30 self.seen_configs: list[VllmConfig] = []
31
32 def _record_config(self) -> None:
33 self.seen_configs.append(get_current_vllm_config())
34
35 def start_weight_update(self) -> None:
36 self._record_config()
37 self.started = True
38
39 def update_weights(self, update_info: dict) -> None:
40 self._record_config()
41 self.update_calls.append(update_info)
42 if self.raise_on_update:
43 raise ValueError("boom")
44
45 def finish_weight_update(self) -> None:
46 self._record_config()
47 self.finished = True
48
49 def reset_weight_update_target(self) -> None:
50 self.reset_count += 1
51
52
53 class _RecordingModelRunner:
54 def __init__(self) -> None:
55 self.seen_config: VllmConfig | None = None
56 self.reset_lora_calls = 0
57
58 def reload_weights(self) -> None:
59 self.seen_config = get_current_vllm_config()
60
61 def reset_lora_state(self) -> None:
62 self.reset_lora_calls += 1
63
64
65 def _make_worker(engine: _RecordingEngine | None) -> Worker:
66 worker = object.__new__(Worker)
67 worker.vllm_config = VllmConfig()
68 worker.weight_transfer_engine = engine
69 worker._weight_update_active = False
70 worker._weight_update_is_draft = False
71 worker.model_runner = _RecordingModelRunner()
72 return worker
73
74
75 def test_reload_weights_sets_current_config():
76 worker = _make_worker(None)
77 model_runner = _RecordingModelRunner()
78 worker.model_runner = model_runner
79
80 Worker.reload_weights(worker)
81
82 assert model_runner.seen_config is worker.vllm_config
83
84
85 def test_reload_parameter_lookup_preserves_lora_module_names():
86 base_layer = nn.Module()
87 qweight = nn.Parameter(torch.ones(1))
88 base_layer.register_parameter("qweight", qweight)
89 wrapper = BaseLayerWithLoRA()
90 wrapper.base_layer = base_layer
91 model = nn.Module()
92 model.proj = wrapper
93
94 named_parameters = dict(model.named_parameters())
95 assert set(named_parameters) == {"proj.base_layer.qweight"}
96 assert named_parameters["proj.base_layer.qweight"] is qweight
97 assert model.get_parameter("proj.base_layer.qweight") is qweight
98 assert _get_parameter_for_reload(model, "proj.qweight") is qweight
99
100
101 def test_start_update_finish_delegates_to_engine():
102 engine = _RecordingEngine()
103 worker = _make_worker(engine)
104
105 Worker.start_weight_update(worker)
106 assert engine.started is True
107 assert worker._weight_update_active is True
108
109 Worker.update_weights(worker, {"names": ["w"]})
110 assert engine.update_calls == [{"names": ["w"]}]
111 assert worker._weight_update_active is True
112
113 Worker.finish_weight_update(worker)
114 assert engine.finished is True
115 assert engine.reset_count == 1
116 assert worker._weight_update_active is False
117 assert engine.seen_configs == [worker.vllm_config] * 3
118 assert worker.model_runner.reset_lora_calls == 1
119
120
121 @pytest.mark.parametrize(
122 ("rank", "expected"),
123 [(1, {"names": ["rank-1"]}), (2, {"names": []})],
124 )
125 def test_rank_local_update_selects_worker_payload(rank, expected):
126 engine = _RecordingEngine()
127 worker = _make_worker(engine)
128 worker.rank = rank
129 Worker.start_weight_update(worker)
130
131 Worker.update_weights(
132 worker, [{"names": ["rank-0"]}, {"names": ["rank-1"]}, {"names": []}]
133 )
134
135 assert engine.update_calls == [expected]
136 assert worker._weight_update_active is True
137
138
139 def test_rank_local_update_uses_data_parallel_index_after_reconfigure():
140 engine = _RecordingEngine()
141 worker = _make_worker(engine)
142 worker.rank = 0
143 parallel_config = ParallelConfig(
144 data_parallel_size=4,
145 data_parallel_rank=2,
146 )
147 assert parallel_config.data_parallel_rank == 2
148 assert parallel_config.data_parallel_index == 2
149
150 parallel_config.reconfigure_for_independent_dp_rank()
151 assert parallel_config.data_parallel_rank == 0
152 assert parallel_config.data_parallel_index == 2
153
154 worker.vllm_config.parallel_config = parallel_config
155 Worker.start_weight_update(worker)
156
157 Worker.update_weights(
158 worker,
159 [
160 {"names": ["dp-0"]},
161 {"names": ["dp-1"]},
162 {"names": ["dp-2"]},
163 {"names": ["dp-3"]},
164 ],
165 )
166
167 assert engine.update_calls == [{"names": ["dp-2"]}]
168 assert worker._weight_update_active is True
169
170
171 def test_finish_draft_session_keeps_lora_state():
172 engine = _RecordingEngine()
173 engine.supports_draft_weight_update = True
174 worker = _make_worker(engine)
175 worker._set_draft_weight_update_target = lambda: None
176
177 Worker.start_draft_weight_update(worker)
178 Worker.finish_weight_update(worker)
179
180 assert worker.model_runner.reset_lora_calls == 0
181
182
183 def test_double_start_raises():
184 worker = _make_worker(_RecordingEngine())
185 Worker.start_weight_update(worker)
186 with pytest.raises(RuntimeError, match="already"):
187 Worker.start_weight_update(worker)
188
189
190 def test_update_without_start_raises():
191 worker = _make_worker(_RecordingEngine())
192 with pytest.raises(RuntimeError, match="start_weight_update must be called"):
193 Worker.update_weights(worker, {"names": ["w"]})
194
195
196 def test_finish_without_start_raises():
197 worker = _make_worker(_RecordingEngine())
198 with pytest.raises(RuntimeError, match="without a matching"):
199 Worker.finish_weight_update(worker)
200
201
202 def test_update_resets_active_on_error():
203 engine = _RecordingEngine(raise_on_update=True)
204 worker = _make_worker(engine)
205 Worker.start_weight_update(worker)
206
207 with pytest.raises(ValueError, match="boom"):
208 Worker.update_weights(worker, {"names": ["w"]})
209
210 # A failed update ends the session so the next start is clean.
211 assert engine.reset_count == 1
212 assert worker._weight_update_active is False
213
214
215 def test_missing_engine_raises():
216 worker = _make_worker(None)
217 with pytest.raises(RuntimeError, match="Weight transfer not configured"):
218 Worker.start_weight_update(worker)