A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
[Core] Add `reload_weights` RPC method (#20096)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
Jul 23, 2025 at 14:24 UTC
5c9b807b34d29fa7abfc33d90725663d96431678
5 files changed
+51
-34
tests/v1/worker/test_gpu_model_runner.py
+6
-1
@@ -460,11 +460,16 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
460
{"load_config": {
461
"load_format": original_load_format
462
}})
463
- model_runner_2.load_model() # Load real weights inplace
463
+ model_runner_2.reload_weights() # Load real weights inplace
464
assert str(model_runner.get_model().state_dict()) == str(
465
model_runner_2.get_model().state_dict())
466
467
468
+def test_reload_weights_before_load_model(model_runner):
469
+ with pytest.raises(AssertionError):
470
+ model_runner.reload_weights()
471
+
472
+
473
def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
474
torch.set_default_dtype(torch.float16)
475
layer_0 = "model.layers.0.self_attn.attn"
vllm/v1/worker/gpu_model_runner.py
+10
-11
@@ -1873,17 +1873,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1873
with DeviceMemoryProfiler() as m:
1874
time_before_load = time.perf_counter()
1875
model_loader = get_model_loader(self.load_config)
1876
- if not hasattr(self, "model"):
1877
- logger.info("Loading model from scratch...")
1878
- self.model = model_loader.load_model(
1879
- vllm_config=self.vllm_config,
1880
- model_config=self.model_config)
1881
- else:
1882
- logger.info(
1883
- "Model was already initialized. Loading weights inplace..."
1884
- )
1885
- model_loader.load_weights(self.model,
1886
- model_config=self.model_config)
1876
+ logger.info("Loading model from scratch...")
1877
+ self.model = model_loader.load_model(
1878
+ vllm_config=self.vllm_config, model_config=self.model_config)
1879
if self.lora_config:
1880
self.model = self.load_lora_model(self.model,
1881
self.model_config,
@@ -1916,6 +1908,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1908
rank_mapping,
1909
)
1910
1911
+ def reload_weights(self) -> None:
1912
+ assert getattr(self, "model", None) is not None, \
1913
+ "Cannot reload weights before model is loaded."
1914
+ model_loader = get_model_loader(self.load_config)
1915
+ logger.info("Reloading weights inplace...")
1916
+ model_loader.load_weights(self.model, model_config=self.model_config)
1917
+
1918
def save_tensorized_model(
1919
self,
1920
tensorizer_config: "TensorizerConfig",
vllm/v1/worker/gpu_worker.py
+21
-12
@@ -4,6 +4,7 @@
4
import copy
5
import gc
6
import os
7
+from contextlib import AbstractContextManager, nullcontext
8
from typing import TYPE_CHECKING, Any, Optional
9
10
import torch
@@ -118,6 +119,21 @@ class Worker(WorkerBase):
119
buffer.data.copy_(self._sleep_saved_buffers[name].data)
120
self._sleep_saved_buffers = {}
121
122
+ def _maybe_get_memory_pool_context(self,
123
+ tag: str) -> AbstractContextManager:
124
+ if self.vllm_config.model_config.enable_sleep_mode:
125
+ from vllm.device_allocator.cumem import CuMemAllocator
126
+
127
+ allocator = CuMemAllocator.get_instance()
128
+ if tag == "weights":
129
+ assert allocator.get_current_usage() == 0, (
130
+ "Sleep mode can only be "
131
+ "used for one instance per process.")
132
+ context = allocator.use_memory_pool(tag=tag)
133
+ else:
134
+ context = nullcontext()
135
+ return context
136
+
137
def initialize_cache(self, num_gpu_blocks: int,
138
num_cpu_blocks: int) -> None:
139
self.cache_config.num_gpu_blocks = num_gpu_blocks
@@ -179,24 +195,17 @@ class Worker(WorkerBase):
195
# FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool
196
# to hijack tensor allocation.
197
def load_model(self) -> None:
182
- if self.vllm_config.model_config.enable_sleep_mode:
183
- from vllm.device_allocator.cumem import CuMemAllocator
184
-
185
- allocator = CuMemAllocator.get_instance()
186
- assert allocator.get_current_usage() == 0, (
187
- "Sleep mode can only be "
188
- "used for one instance per process.")
189
- context = allocator.use_memory_pool(tag="weights")
190
- else:
191
- from contextlib import nullcontext
192
- context = nullcontext()
198
eep_scale_up = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"
194
- with context:
199
+ with self._maybe_get_memory_pool_context(tag="weights"):
200
self.model_runner.load_model(eep_scale_up=eep_scale_up)
201
202
def update_config(self, overrides: dict[str, Any]) -> None:
203
self.model_runner.update_config(overrides)
204
205
+ def reload_weights(self) -> None:
206
+ with self._maybe_get_memory_pool_context(tag="weights"):
207
+ self.model_runner.reload_weights()
208
+
209
@torch.inference_mode()
210
def determine_available_memory(self) -> int:
211
"""Profiles the peak memory usage of the model to determine how much
vllm/v1/worker/tpu_model_runner.py
+11
-10
@@ -1174,16 +1174,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1174
mesh=self.mesh)
1175
else:
1176
model_loader = get_model_loader(self.load_config)
1177
- if not hasattr(self, "model"):
1178
- logger.info("Loading model from scratch...")
1179
- model = model_loader.load_model(
1180
- vllm_config=self.vllm_config,
1181
- model_config=self.model_config)
1182
- else:
1183
- logger.info("Model was already initialized. \
1184
- Loading weights inplace...")
1185
- model_loader.load_weights(
1186
- self.model, model_config=self.model_config)
1177
+ logger.info("Loading model from scratch...")
1178
+ model = model_loader.load_model(
1179
+ vllm_config=self.vllm_config,
1180
+ model_config=self.model_config)
1181
except RuntimeError as e:
1182
raise RuntimeError(
1183
f"Unable to load model, a likely reason is the model is "
@@ -1205,6 +1199,13 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1199
self.model = model
1200
self.sampler = TPUSampler()
1201
1202
+ def reload_weights(self) -> None:
1203
+ assert getattr(self, "model", None) is not None, \
1204
+ "Cannot reload weights before model is loaded."
1205
+ model_loader = get_model_loader(self.load_config)
1206
+ logger.info("Reloading weights inplace...")
1207
+ model_loader.load_weights(self.model, model_config=self.model_config)
1208
+
1209
@torch.no_grad()
1210
def _dummy_run(self, num_tokens: int, num_reqs: int,
1211
num_blocks: int) -> None:
vllm/v1/worker/tpu_worker.py
+3
@@ -265,6 +265,9 @@ class TPUWorker:
265
def update_config(self, overrides: dict[str, Any]) -> None:
266
self.model_runner.update_config(overrides)
267
268
+ def reload_weights(self) -> None:
269
+ self.model_runner.reload_weights()
270
+
271
def compile_or_warm_up_model(self) -> None:
272
if not self.model_config.enforce_eager:
273
self.model_runner.capture_model()