@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f857a03f6b
Fork 0

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 9, 07:30 UTC

[QeRL] Layerwise Reloading (#32133)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Kyle Sayers committed Jan 30, 2026 at 10:50 UTC f857a03f6b71101ac3f67b766f9a165c8b617f7b
17 files changed +923 -314
tests/conftest.py
+22 -4
@@ -27,7 +27,7 @@ import threading
27 from collections.abc import Generator
28 from contextlib import nullcontext
29 from enum import Enum
30 -from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING
30 +from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING, Optional
31
32 import numpy as np
33 import pytest
@@ -1024,7 +1024,9 @@ class VllmRunner:
1024 **kwargs,
1025 )
1026
1027 - def generate_prompt_perplexity(self, prompts: list[str]) -> list[float]:
1027 + def generate_prompt_perplexity(
1028 + self, prompts: list[str], mask: Optional[list[str]] = None
1029 + ) -> list[float]:
1030 """
1031 Return the perplexity score associated with generating the prompts
1032
@@ -1035,13 +1037,20 @@ class VllmRunner:
1037 prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
1038 )
1039
1040 + mask_prefix_lens = (
1041 + [len(self.llm.get_tokenizer()(prefix)["input_ids"]) for prefix in mask]
1042 + if mask is not None
1043 + else [0 for _ in range(len(prompts))]
1044 + )
1045 +
1046 perplexities = []
1039 - for output in outputs:
1047 + for output, mask_prefix_len in zip(outputs, mask_prefix_lens):
1048 output = cast(TokensTextLogprobsPromptLogprobs, output)
1049 token_datas = cast(list[dict[int, Logprob] | None], output[3])
1050 assert token_datas[0] is None
1051 +
1052 token_log_probs = []
1044 - for token_data in token_datas[1:]:
1053 + for token_data in token_datas[mask_prefix_len + 1 :]:
1054 assert token_data is not None
1055 assert len(token_data) == 1
1056 token_log_prob = list(token_data.values())[0].logprob
@@ -1122,6 +1131,9 @@ class VllmRunner:
1131 def get_llm(self) -> LLM:
1132 return self.llm
1133
1134 + def collective_rpc(self, *args, **kwargs):
1135 + return self.llm.collective_rpc(*args, **kwargs)
1136 +
1137 def __enter__(self):
1138 return self
1139
@@ -1532,3 +1544,9 @@ def use_fresh_inductor_cache():
1544 """
1545 with fresh_cache():
1546 yield
1547 +
1548 +
1549 +@pytest.fixture(scope="function")
1550 +def enable_pickle(monkeypatch):
1551 + """`LLM.apply_model` requires pickling a function."""
1552 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
tests/model_executor/model_loader/test_reload.py new
+150
@@ -0,0 +1,150 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import gc
4 +import inspect
5 +from weakref import WeakKeyDictionary, ref
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.model_executor.layers.linear import QKVParallelLinear
11 +from vllm.model_executor.model_loader.reload.meta import (
12 + capture_layer_to_meta,
13 + get_numel_loaded,
14 + materialize_layer,
15 + materialize_meta_tensor,
16 + restore_layer_on_meta,
17 + to_meta_tensor,
18 +)
19 +from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo
20 +from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
21 +from vllm.platforms import current_platform
22 +from vllm.utils.torch_utils import cuda_device_count_stateless
23 +
24 +
25 +def test_move_metatensors():
26 + tensor = torch.empty((1, 2, 3))
27 + meta_tensor = to_meta_tensor(tensor)
28 + materialized_tensor = materialize_meta_tensor(meta_tensor)
29 +
30 + assert meta_tensor.device.type == "meta"
31 + assert tensor.device == materialized_tensor.device
32 +
33 + assert tensor.dtype == meta_tensor.dtype == materialized_tensor.dtype
34 + assert tensor.shape == meta_tensor.shape == materialized_tensor.shape
35 + assert tensor.__class__ == meta_tensor.__class__ == materialized_tensor.__class__
36 + assert tensor.__dict__ == meta_tensor.__dict__ == materialized_tensor.__dict__
37 +
38 +
39 +def test_reload_lifecycle():
40 + layer = torch.nn.Linear(2, 3)
41 + info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))
42 +
43 + restore_layer_on_meta(layer, info)
44 + for name, tensor in get_layer_tensors(layer).items():
45 + meta_tensor = getattr(layer, name)
46 + assert tensor.dtype == meta_tensor.dtype
47 + assert tensor.shape == meta_tensor.shape
48 + assert tensor.__class__ == meta_tensor.__class__
49 + assert tensor.__dict__ == meta_tensor.__dict__
50 +
51 + materialize_layer(layer)
52 + for name, tensor in get_layer_tensors(layer).items():
53 + materialized_tensor = getattr(layer, name)
54 + assert tensor.dtype == materialized_tensor.dtype
55 + assert tensor.shape == materialized_tensor.shape
56 + assert tensor.__class__ == materialized_tensor.__class__
57 + assert tensor.__dict__ == materialized_tensor.__dict__
58 +
59 +
60 +def test_model_cleanup(dist_init, default_vllm_config):
61 + layer = QKVParallelLinear(2, 3, 4)
62 + assert layer.weight.weight_loader.__self__ is layer
63 + info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))
64 +
65 + mock_info_dict: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
66 + WeakKeyDictionary()
67 + )
68 + mock_info_dict[layer] = info
69 + layer_ref = ref(layer)
70 +
71 + del layer
72 + gc.collect()
73 +
74 + assert layer_ref() is None
75 + assert len(mock_info_dict) == 0
76 +
77 +
78 +def test_get_numel_loaded():
79 + param = torch.empty(10, device="meta")
80 + loaded_weight = torch.empty(10)
81 +
82 + def complex_weight_loader(param, loaded_weight):
83 + param[:3] = loaded_weight[:3]
84 + param[5:8] = loaded_weight[5:8]
85 + return "value"
86 +
87 + args = inspect.signature(complex_weight_loader).bind(param, loaded_weight)
88 + num_loaded, ret = get_numel_loaded(complex_weight_loader, args)
89 + assert num_loaded == 6
90 + assert ret == "value"
91 +
92 +
93 +@pytest.mark.parametrize("tp_size", [2])
94 +@pytest.mark.parametrize(
95 + "base_model,mul_model,add_model",
96 + [
97 + (
98 + "Qwen/Qwen3-0.6B",
99 + "inference-optimization/Qwen3-0.6B-debug-multiply",
100 + "inference-optimization/Qwen3-0.6B-debug-add",
101 + ),
102 + (
103 + "inference-optimization/Qwen3-0.6B-FP8_BLOCK",
104 + "inference-optimization/Qwen3-0.6B-debug-multiply-FP8_BLOCK",
105 + "inference-optimization/Qwen3-0.6B-debug-add-FP8_BLOCK",
106 + ),
107 + (
108 + "inference-optimization/Qwen3-0.6B-W4A16-G128",
109 + "inference-optimization/Qwen3-0.6B-debug-multiply-W4A16-G128",
110 + "inference-optimization/Qwen3-0.6B-debug-add-W4A16-G128",
111 + ),
112 + (
113 + "inference-optimization/DeepSeek-V3-debug-empty",
114 + "inference-optimization/DeepSeek-V3-debug-multiply",
115 + "inference-optimization/DeepSeek-V3-debug-add",
116 + ),
117 + (
118 + "inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC",
119 + "inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC",
120 + "inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC",
121 + ),
122 + (
123 + "inference-optimization/DeepSeek-V3-debug-empty-NVFP4A16",
124 + "inference-optimization/DeepSeek-V3-debug-multiply-NVFP4A16",
125 + "inference-optimization/DeepSeek-V3-debug-add-NVFP4A16",
126 + ),
127 + ],
128 +)
129 +def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
130 + if cuda_device_count_stateless() < tp_size:
131 + pytest.skip(reason="Not enough CUDA devices")
132 +
133 + if "FP8" in base_model and not current_platform.supports_fp8():
134 + pytest.skip(reason="Requires FP8 support")
135 +
136 + with vllm_runner(
137 + model_name=base_model,
138 + tensor_parallel_size=tp_size,
139 + enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
140 + enable_prefix_caching=False,
141 + ) as llm:
142 + llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
143 + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
144 + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
145 + assert mul_perp < add_perp
146 +
147 + llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})
148 + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
149 + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
150 + assert add_perp < mul_perp
tests/quantization/test_torchao.py
+15 -4
@@ -1,11 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import importlib.metadata
3 import importlib.util
4
5 import pytest
6 import torch
7
8 +from vllm.model_executor.model_loader import get_model_loader
9 from vllm.platforms import current_platform
10
11 DTYPE = ["bfloat16"]
@@ -105,8 +105,8 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):
105
106
107 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
108 -def test_online_quant_config_dict_json(vllm_runner):
109 - """Testing on the fly quantization, load_weights integration point,
108 +def test_online_quant_config_dict_json(vllm_runner, enable_pickle):
109 + """Testing online quantization, load_weights integration point,
110 with config dict serialized to json string
111 """
112 torch._dynamo.reset()
@@ -135,7 +135,18 @@ def test_online_quant_config_dict_json(vllm_runner):
135 ) as llm:
136 output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
137
138 - assert output
138 + load_config = llm.llm.llm_engine.vllm_config.load_config
139 + model_config = llm.llm.llm_engine.vllm_config.model_config
140 +
141 + def load_weights(model):
142 + model_loader = get_model_loader(load_config)
143 + weights_iterator = model_loader.get_all_weights(model_config, model)
144 + model.load_weights(weights_iterator)
145 +
146 + llm.apply_model(load_weights)
147 +
148 + reload_output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
149 + assert output[0][0] == reload_output[0][0]
150
151
152 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
tests/v1/worker/test_gpu_model_runner.py
+1 -1
@@ -543,7 +543,7 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
543
544
545 def test_reload_weights_before_load_model(model_runner):
546 - with pytest.raises(AssertionError):
546 + with pytest.raises(ValueError):
547 model_runner.reload_weights()
548
549
vllm/model_executor/model_loader/online_quantization.py deleted
-275
@@ -1,275 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import types
5 -from collections.abc import Iterable
6 -
7 -import torch
8 -from torch import nn
9 -
10 -from vllm.config import ModelConfig
11 -from vllm.logger import init_logger
12 -from vllm.model_executor.model_loader.utils import process_weights_after_loading
13 -
14 -logger = init_logger(__name__)
15 -
16 -# Notes for Online Quantization
17 -# In terms of state of checkpoints, quantization config and their
18 -# correspondance to online quantization:
19 -# | Use Case | Checkpoints | model_config.quantization |
20 -# | no quant | high precision | None |
21 -# | offline quant | quantized | fp8, torchao etc. |
22 -# | online quant | high precision | torchao etc. |
23 -#
24 -# The process for loading non-quantized checkpoint
25 -# 1. load non-quantized weights (load_weights)
26 -# 2. do any additional post processing (process_weights_after_loading)
27 -#
28 -# The process for loading offline quantized checkpoint
29 -# 1. load offline-quantized weights (load_weights)
30 -# 2. do any additional post processing (process_weights_after_loading)
31 -
32 -# The process for unquantized model reloading
33 -# (repeated run in RL training loop)
34 -# first run
35 -# UI1. load_weights: load bfloat16 weights
36 -# UI2. process_weights_after_loading: any additional post processing
37 -# subsequent run
38 -# UC1: load_weights: load bfloat16 weights
39 -# (shouldn't be any issues since we didn't change any attributes
40 -# of the weights)
41 -# UC2: process_weights_after_loading: any additional post processing
42 -
43 -# The process for weight reloading with online quantization
44 -# (repeated run in RL training loop)
45 -# first run
46 -# I1. load_weights: load bfloat16 weights
47 -# I2. process_weights_after_loading:
48 -# record weight metadata and attributes for R1 and R2
49 -# quantize weights to fp8
50 -# subsequent run
51 -# (beginning model weight is in fp8)
52 -# load_weights:
53 -# R1. restore bfloat16 model weight metadata
54 -# R2. restore the model weight attributes
55 -# R3. reload bfloat16 weights
56 -# R4. quantize weights (by calling process_weights_after_loading),
57 -# also set `process_weights_after_loading_already_called` to
58 -# True to stop it from running again
59 -# R5. (workaround for cudagraph), we restore the weight params to original quantized
60 -# weights params, and use original_weight_param.copy_(updated_weight_param) so that
61 -# the weight update work well with cudagraph
62 -# process_weights_after_loading (if called):
63 -# this will be skipped since it's already ran in
64 -# load_weights
65 -
66 -
67 -def maybe_save_metadata_and_attributes_for_weight_reloading(
68 - model: nn.Module, model_config: ModelConfig
69 -):
70 - # following is to support on the fly quantization, currently only supported
71 - # for torchao
72 - if model_config.quantization != "torchao":
73 - return
74 -
75 - from vllm.model_executor.model_loader.weight_utils import get_quant_config
76 -
77 - quant_config = get_quant_config(model_config, None)
78 -
79 - # If checkpoint is already torchao serialized, this means it's
80 - # pre-quantized quantization case, we'll skip saving the metadata
81 - # Otherwise, this is Step I2 of initialization steps of
82 - # online quantization
83 - # This step record the weights metadata and weight attributes so we can
84 - # restore the bfloat16 model weights during the relad step (R1 and R2)
85 - # see Notes in online_quantization.py for more details
86 - if not (
87 - hasattr(quant_config, "is_checkpoint_torchao_serialized")
88 - and not quant_config.is_checkpoint_torchao_serialized
89 - ):
90 - return
91 -
92 - # This is the I2 step of online quantiztion that saves
93 - # metadata and attributes of weights so they can be used in R1 and
94 - # R2 step, note that we only save these during initialization
95 -
96 - # Includes two things
97 - # 1. save floating point metadata (shape, dtype, device) for init
98 - # 2. save weight attributes, e.g. `output_dim`, `weight_loader` for init
99 -
100 - if getattr(model, "weight_metadata_and_attr_saved", False):
101 - return
102 -
103 - # save the dtype, shape and device for model parameter, used for
104 - # restoring the model high precision parameters before
105 - # reloading the weights
106 - assert not hasattr(model, "original_weights_rebuild_keys")
107 - model.original_weights_rebuild_keys = {}
108 - for name, p in model.named_parameters():
109 - model.original_weights_rebuild_keys[name] = {
110 - "shape": p.shape,
111 - "dtype": p.dtype,
112 - "device": p.device,
113 - }
114 -
115 - # record the weight attributes (loader functions etc.)
116 - # so these can be recovered later when we reload the weights
117 - # structure: {"weight_name": {"weight_attr_key": attr}}
118 - assert not hasattr(model, "recorded_weight_attr")
119 - model.recorded_weight_attr = {}
120 - for name, param in model.named_parameters():
121 - model.recorded_weight_attr[name] = {}
122 - for key in param.__dict__:
123 - if hasattr(param, key):
124 - attr = getattr(param, key)
125 - if not callable(attr):
126 - model.recorded_weight_attr[name][key] = attr
127 - elif hasattr(attr, "__self__") and param is attr.__self__:
128 - # if attr is a bonded method for an instance, and
129 - # attr.__self__ points to the instance (param)
130 - # we'll record the underlying function object
131 - model.recorded_weight_attr[name][key] = attr.__func__
132 - else:
133 - model.recorded_weight_attr[name][key] = attr
134 - # mark the metadata and attributes saved so we don't run it again
135 - model._model_config = model_config
136 - model.weight_metadata_and_attr_saved = True
137 -
138 -
139 -def _bond_method_to_cls(func, obj):
140 - if hasattr(func, "__self__") or not callable(func):
141 - # If the function is already bound to an instance, return it as is
142 - return func
143 - else:
144 - return types.MethodType(func, obj)
145 -
146 -
147 -def support_quantized_model_reload_from_hp_weights(original_load_weights):
148 - """Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
149 - reloading high precision (bfloat16/float16/float32) weight for an already quantized
150 - model, this involves restoring the weights to a high precision weights and
151 - then online quantize the weights
152 - """
153 - # online quantization, right now only enabled for
154 - # torchao
155 - # R1, R2, R3, R4, R5 in the Notes
156 -
157 - def patched_model_load_weights(
158 - auto_weight_loader, weights: Iterable[tuple[str, torch.Tensor]], *, mapper=None
159 - ) -> set[str]:
160 - model = auto_weight_loader.module
161 - offline_quantization_or_first_run_of_online_quantization = not getattr(
162 - model, "weight_metadata_and_attr_saved", False
163 - )
164 -
165 - # if we don't have `model.weight_metadata_and_attr_saved` defined and
166 - # set to True, it means that this is either offline quantization case
167 - # or the first run of online quantization
168 - # see Notes in this file for more details
169 - if offline_quantization_or_first_run_of_online_quantization:
170 - # case 1: offline quantized checkpoint
171 - # case 2: Step I1 first run of weight loading with
172 - # online quantization
173 - return original_load_weights(auto_weight_loader, weights, mapper=mapper)
174 -
175 - model_config = model._model_config
176 -
177 - # TODO: Add fp8 support
178 - assert model_config.quantization == "torchao", (
179 - "online quantization is only enabled for torchao currently"
180 - )
181 - # TODO: use create_weights to restore the weights to original state
182 -
183 - # Step R1: First restore the quantized weights to original bfloat16
184 - # weights, with original metadata (shape, dtype, device)
185 - # and attributes, so that bfloat16 weights can be loaded properly
186 - # TODO: maybe set remove_duplicate to True?
187 - original_quantized_weight_dict = dict(
188 - model.named_parameters(remove_duplicate=False)
189 - )
190 - named_modules = dict(model.named_modules(remove_duplicate=False))
191 - model_device = None
192 -
193 - for name, d in model.original_weights_rebuild_keys.items():
194 - _shape = d["shape"]
195 - _dtype = d["dtype"]
196 - _device = d["device"]
197 - if model_device is not None:
198 - assert model_device == _device, (
199 - "Expecting all weights "
200 - "to be in the same device for now, got both: "
201 - f"{model_device} and {_device}"
202 - )
203 - else:
204 - model_device = _device
205 -
206 - if name in original_quantized_weight_dict:
207 - module_name, weight_name = name.rsplit(".", 1)
208 - module = named_modules[module_name]
209 - setattr(
210 - module,
211 - weight_name,
212 - torch.nn.Parameter(
213 - torch.empty(_shape, dtype=_dtype, device=_device),
214 - requires_grad=False,
215 - ),
216 - )
217 -
218 - # Step R2: recover the weight attributes to the state before first loading
219 - # recorded_weight_attr is
220 - # {"weight_name": {"weight_attr_key": attr}}
221 - # e.g.
222 - # {
223 - # {
224 - # "layer.0.weight": {
225 - # "weight_loader": weight_loader_function_object,
226 - # "input_dim": 0, ...
227 - # },
228 - # "layer.1.weight": ...,
229 - # }
230 - # }
231 - for full_weight_name, weight_attr_dict in model.recorded_weight_attr.items():
232 - for attr_name, attr in weight_attr_dict.items():
233 - module_name, weight_name = full_weight_name.rsplit(".", 1)
234 - module = named_modules[module_name]
235 - weight = getattr(module, weight_name)
236 - if not hasattr(weight, attr_name):
237 - setattr(weight, attr_name, _bond_method_to_cls(attr, weight))
238 -
239 - # Step R3: reload bfloat16 / high precision weights
240 - updated_params = original_load_weights(
241 - auto_weight_loader, weights, mapper=mapper
242 - )
243 -
244 - # Step R4: online quantize the weights
245 - # manually process weights after loading
246 - model.process_weights_after_loading_already_called = False
247 - if model_device is not None:
248 - process_weights_after_loading(model, model_config, model_device)
249 - else:
250 - logger.warning_once(
251 - "model_device is None, skip calling process_weights_after_loading"
252 - )
253 -
254 - # Step R5 (workaround for cudagraph): restore the original quantized weights
255 - # and do a copy_ of the currents weights to the original weights
256 - updated_quantized_weights = dict(model.named_parameters(remove_duplicate=False))
257 - for name in model.original_weights_rebuild_keys:
258 - if name in original_quantized_weight_dict:
259 - original_quantized_weight = original_quantized_weight_dict[name]
260 - updated_quantized_weight = updated_quantized_weights[name]
261 -
262 - module_name, weight_name = name.rsplit(".", 1)
263 - module = named_modules[module_name]
264 - setattr(module, weight_name, original_quantized_weight)
265 - with torch.no_grad():
266 - original_quantized_weight.copy_(updated_quantized_weight)
267 -
268 - del original_quantized_weight_dict
269 - del named_modules
270 - del updated_quantized_weight
271 -
272 - model.process_weights_after_loading_already_called = True
273 - return updated_params
274 -
275 - return patched_model_load_weights
vllm/model_executor/model_loader/reload/__init__.py new
+37
@@ -0,0 +1,37 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Layerwise weight reloading utilities for vLLM.
5 +
6 +This module provides functionality to reload model weights layer-by-layer,
7 +which is useful for weight updates without full model reconstruction.
8 +
9 +Limitations:
10 +1. Composition with CPU offloading has not been implemented
11 +2. Reloading Attention/MLA weights (q_scale, k_scale, v_scale) has not been implemented
12 +3. Tied parameters will only reflect processing from one of the parent layers (for
13 + example, only processing from embed_tokens will have an effect)
14 +4. This design assumes that the number of weights loaded from disk is the same as the
15 + number of weights created at model init time. This is not true for quant methods
16 + which (1) pad weights or (2) load qkv weights into the same parameter. Both of these
17 + cases are non-issues for today's quant methods, but future quantizations may cause
18 + reloading to fail
19 +"""
20 +
21 +__all__ = [
22 + "record_metadata_for_reloading",
23 + "initialize_layerwise_reload",
24 + "finalize_layerwise_reload",
25 + "set_torchao_reload_attrs",
26 + "support_quantized_model_reload_from_hp_weights",
27 +]
28 +
29 +from .layerwise import (
30 + finalize_layerwise_reload,
31 + initialize_layerwise_reload,
32 + record_metadata_for_reloading,
33 +)
34 +from .torchao_decorator import (
35 + set_torchao_reload_attrs,
36 + support_quantized_model_reload_from_hp_weights,
37 +)
vllm/model_executor/model_loader/reload/layerwise.py new
+270
@@ -0,0 +1,270 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import inspect
4 +from collections.abc import Callable
5 +from functools import wraps
6 +from weakref import WeakKeyDictionary
7 +
8 +import torch
9 +
10 +from vllm.attention.layer import Attention, MLAAttention
11 +from vllm.config import ModelConfig
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
14 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15 +
16 +from .meta import (
17 + capture_layer_to_meta,
18 + get_numel_loaded,
19 + materialize_layer,
20 + restore_layer_on_meta,
21 +)
22 +from .types import LayerReloadingInfo
23 +from .utils import get_layer_params_buffers, get_layer_size, get_layer_tensors
24 +
25 +logger = init_logger(__name__)
26 +
27 +__all__ = [
28 + "get_layerwise_info",
29 + "record_metadata_for_reloading",
30 + "initialize_layerwise_reload",
31 + "finalize_layerwise_reload",
32 +]
33 +
34 +
35 +# Global dict storing information used for layerwise restoring, loading, and processing.
36 +# For more information regarding what info is stored when, see `LayerReloadingInfo`
37 +#
38 +# Use a weak ref dictionary so that modules can be freed when the model is freed.
39 +# Values are sanitized from references to the layer key in order to avoid circular refs
40 +LAYERWISE_INFO: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
41 + WeakKeyDictionary()
42 +)
43 +
44 +
45 +def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:
46 + """
47 + Get information related to restoring and layerwise processing. If no previous
48 + information existed, a new entry is constructed
49 + """
50 + if layer not in LAYERWISE_INFO:
51 + LAYERWISE_INFO[layer] = LayerReloadingInfo()
52 +
53 + return LAYERWISE_INFO[layer]
54 +
55 +
56 +def record_metadata_for_reloading(model: torch.nn.Module):
57 + """
58 + Record layer metadata needed for later reloading.
59 +
60 + Stores parameter and buffer metadata as meta tensors for restoration.
61 + Must be called before `initialize_layerwise_reload`.
62 + """
63 + for layer in model.modules():
64 + info = get_layerwise_info(layer)
65 + info.restore_metadata = capture_layer_to_meta(layer)
66 +
67 +
68 +@torch.no_grad()
69 +def initialize_layerwise_reload(model: torch.nn.Module):
70 + """
71 + Set up layerwise weight loading with deferred processing.
72 +
73 + Must be called after `record_metadata_for_reloading`. This function:
74 + 1. Saves current kernel tensors for later copying
75 + 2. Restores layer parameters/buffers from metadata (on meta device)
76 + 3. Wraps weight loaders to defer processing until all weights are loaded
77 +
78 + When all weights for a layer are loaded, the wrapped loaders will:
79 + 1. Materialize the layer onto the target device
80 + 2. Load all cached weights
81 + 3. Run quantization processing if applicable
82 + 4. Copy processed values back to original tensor storage
83 + """
84 + # disable torchao reloading to avoid infinite recursion
85 + model._original_do_torchao_reload = getattr(model, "_do_torchao_reload", False)
86 + model._do_torchao_reload = False
87 +
88 + for layer in model.modules():
89 + info = get_layerwise_info(layer)
90 +
91 + # Skip if the layer has already been initialized
92 + if info.can_process():
93 + continue
94 +
95 + # Save current tensors for later copying
96 + info.kernel_tensors = get_layer_params_buffers(layer)
97 +
98 + # Restore layer parameters/buffers onto meta device
99 + restore_layer_on_meta(layer, info)
100 +
101 + # Track loading progress to determine when to process/copy
102 + info.load_numel = 0
103 + info.load_numel_total = get_layer_size(layer)
104 +
105 + # Wrap each parameter's weight loader
106 + # Note that nested wrapping will occur for shared tensors
107 + for name, tensor in get_layer_tensors(layer).items():
108 + if _get_weight_loader(tensor).__name__ != "online_process_loader":
109 + tensor.weight_loader = make_online_process_loader(layer, name)
110 +
111 +
112 +def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Callable:
113 + """Create a wrapped weight loader that defers processing."""
114 + info = get_layerwise_info(layer)
115 + param = getattr(layer, param_name)
116 + original_loader = _get_original_loader(param)
117 + loader_signature = inspect.signature(original_loader)
118 +
119 + @wraps(original_loader, assigned=("__doc__", "__annotations__"))
120 + def online_process_loader(*args, **kwargs):
121 + if not info.can_process():
122 + # Unfortunately, some qconfigs are set up to load the same weight
123 + # multiple times. For example, CT_WNA16 loads `weight_shape` for
124 + # each of the qkv partitions. This results in layers loading extra
125 + # weights (beyond load_numel_total) after it's already processed.
126 + #
127 + # Best solution is to ensure that `load_numel_total` reflects the
128 + # actual number of weights loaded, either by modifying qconfigs to
129 + # create as many weights as loaded (see padding issue as well)
130 + # or maybe capturing how many weights are loaded on first pass
131 + #
132 + # For now, `load_numel_total` is still safe to use as long as
133 + # there's no way to reach `load_numel_total` without loading all
134 + # necessary weights. `weight_shape` is very small, so this is safe.
135 + # see Limitations(4)
136 + logger.debug("%s: Excessive loading", layer.__class__.__name__)
137 + return
138 +
139 + # Bind and normalize arguments
140 + bound_args = loader_signature.bind(*args, **kwargs)
141 + bound_args.apply_defaults()
142 +
143 + # Cache loaded weights, track loading progress
144 + info.loaded_weights.append((param_name, bound_args))
145 + num_loaded, ret = get_numel_loaded(original_loader, bound_args)
146 + info.load_numel += num_loaded
147 +
148 + logger.debug(
149 + "%s: %d / %d",
150 + layer.__class__.__name__,
151 + info.load_numel,
152 + info.load_numel_total,
153 + )
154 +
155 + # Process and copy when all weights are loaded
156 + if info.load_numel >= info.load_numel_total and not isinstance( # type: ignore[operator]
157 + layer, (Attention, MLAAttention)
158 + ):
159 + _layerwise_process(layer, info)
160 +
161 + return ret
162 +
163 + return online_process_loader
164 +
165 +
166 +def finalize_layerwise_reload(model: torch.nn.Module, model_config: ModelConfig):
167 + """
168 + Remove the outermost layer of weight loading wrappers.
169 +
170 + This function should be applied after `initialize_layerwise_reload` is applied
171 + unwrap the layerwise weight loaders.
172 +
173 + Also processes Attention/MLA layers, which must be processed after all other layers
174 + """
175 + model._do_torchao_reload = model._original_do_torchao_reload
176 +
177 + for layer in model.modules():
178 + info = get_layerwise_info(layer)
179 +
180 + # Attention/MLA layers are processed after all other layers
181 + if isinstance(layer, (Attention, MLAAttention)):
182 + if info.load_numel > 0:
183 + raise NotImplementedError(
184 + "Layerwise reloading of Q/K/V scale weights is not implemented yet"
185 + )
186 +
187 + else:
188 + _place_kernel_tensors(layer, info)
189 + layer.process_weights_after_loading(model_config.dtype)
190 +
191 + # No weights were loaded, place kernel tensors back
192 + elif info.can_process() and info.load_numel <= 0:
193 + _place_kernel_tensors(layer, info)
194 +
195 + # Process non-attention layers which did not load all elements. This can happen
196 + # if the created weight has extra padding elements which are not loaded
197 + # Having too many of these delayed layers can lead to execess memory usage
198 + # see Limitations(4)
199 + elif info.load_numel > 0 and info.load_numel < info.load_numel_total: # type: ignore[operator]
200 + logger.debug("%s: Delayed processing", layer.__class__.__name__)
201 + _layerwise_process(layer, info)
202 +
203 + info.reset()
204 +
205 +
206 +def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
207 + """
208 + Finalize layer loading after all weights have been cached.
209 +
210 + This function:
211 + 1. Materializes the layer onto the target device
212 + 2. Loads all cached weights
213 + 3. Runs quantization processing if applicable
214 + 4. Copies processed values back to original tensor storage
215 + """
216 + # Materialize layer tensors onto device
217 + materialize_layer(layer)
218 +
219 + # Unwrap layerwise loading wrappers
220 + for param in get_layer_tensors(layer).values():
221 + param.weight_loader = _get_original_loader(param)
222 +
223 + # Load all cached weights into materialized layer (using original loaders)
224 + for name, args in info.loaded_weights:
225 + param = getattr(layer, name)
226 + args.arguments["param"] = param
227 + param.weight_loader(*args.args, **args.kwargs)
228 +
229 + # Process weights (quantization, repacking, etc.)
230 + # Attention/MLA are processed in `finalize_layerwise_reload`
231 + quant_method = getattr(layer, "quant_method", None)
232 + if isinstance(quant_method, QuantizeMethodBase):
233 + quant_method.process_weights_after_loading(layer)
234 +
235 + # Copy processed values into original tensor storage (preserves cudagraph refs)
236 + # this code is a no-op if not reloading (because kernel tensors is empty)
237 + parameters, buffers = info.kernel_tensors
238 + for name, param in parameters.items():
239 + param.data.copy_(getattr(layer, name))
240 + for name, buffer in buffers.items():
241 + buffer.data.copy_(getattr(layer, name))
242 +
243 + _place_kernel_tensors(layer, info)
244 +
245 + info.reset()
246 + logger.debug("%s: Processed", layer.__class__.__name__)
247 +
248 +
249 +def _get_original_loader(tensor: torch.Tensor) -> Callable:
250 + """Return the weight loader with any layerwise wrappers removed"""
251 + loader = _get_weight_loader(tensor)
252 + while loader.__name__ == "online_process_loader":
253 + loader = loader.__wrapped__ # type: ignore[union-attr]
254 +
255 + return loader
256 +
257 +
258 +def _get_weight_loader(tensor: torch.Tensor):
259 + return getattr(tensor, "weight_loader", default_weight_loader)
260 +
261 +
262 +def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
263 + for name in get_layer_tensors(layer):
264 + delattr(layer, name)
265 +
266 + parameters, buffers = info.kernel_tensors
267 + for name, param in parameters.items():
268 + layer.register_parameter(name, param)
269 + for name, buffer in buffers.items():
270 + layer.register_buffer(name, buffer)
vllm/model_executor/model_loader/reload/meta.py new
+146
@@ -0,0 +1,146 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import inspect
4 +from collections.abc import Callable
5 +
6 +import torch
7 +from torch.utils._python_dispatch import TorchDispatchMode
8 +
9 +from .sanitize import restore_layer_refs, sanitize_layer_refs
10 +from .types import LayerReloadingInfo, LayerTensors
11 +from .utils import get_layer_params_buffers, get_layer_tensors
12 +
13 +__all__ = [
14 + "to_meta_tensor",
15 + "materialize_meta_tensor",
16 + "capture_layer_to_meta",
17 + "restore_layer_on_meta",
18 + "materialize_layer",
19 + "get_numel_loaded",
20 +]
21 +
22 +SKIP_MODULES: set[str] = {"HadamardTransform"}
23 +
24 +SKIP_TENSORS: set[str] = {
25 + "_expert_map",
26 + "expert_mask",
27 + "expert_global_to_physical",
28 + "expert_physical_to_global",
29 + "expert_local_to_global",
30 +}
31 +
32 +
33 +def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor:
34 + """Convert a tensor to a meta tensor while preserving class and attributes."""
35 + meta_tensor = tensor.data.to("meta")
36 + meta_tensor.__class__ = tensor.__class__
37 + meta_tensor.__dict__ = tensor.__dict__.copy()
38 + return meta_tensor
39 +
40 +
41 +def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor:
42 + """
43 + Materialize a meta tensor into an actual tensor on the current device.
44 + Should be called within the torch device context for the given rank.
45 + """
46 + tensor = torch.empty_strided(
47 + size=tuple(meta_tensor.size()),
48 + stride=tuple(meta_tensor.stride()),
49 + dtype=meta_tensor.dtype,
50 + requires_grad=False,
51 + )
52 + tensor.__class__ = meta_tensor.__class__
53 + tensor.__dict__ = meta_tensor.__dict__.copy()
54 + return tensor
55 +
56 +
57 +def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors:
58 + if layer.__class__.__name__ in SKIP_MODULES:
59 + return ({}, {})
60 +
61 + params, buffers = get_layer_params_buffers(layer)
62 + return (
63 + {
64 + name: sanitize_layer_refs(to_meta_tensor(param), layer)
65 + for name, param in params.items()
66 + if name not in SKIP_TENSORS
67 + },
68 + {
69 + name: sanitize_layer_refs(to_meta_tensor(buffer), layer)
70 + for name, buffer in buffers.items()
71 + if name not in SKIP_TENSORS
72 + },
73 + )
74 +
75 +
76 +def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):
77 + """Restore a layer to model format with tensors on the meta device"""
78 + if layer.__class__.__name__ in SKIP_MODULES:
79 + return
80 +
81 + for name in get_layer_tensors(layer):
82 + if name not in SKIP_TENSORS:
83 + delattr(layer, name)
84 +
85 + restore_params, restore_buffers = info.restore_metadata
86 + for name, param in restore_params.items():
87 + if name not in SKIP_TENSORS:
88 + param = restore_layer_refs(param, layer)
89 + layer.register_parameter(name, param)
90 +
91 + for name, buffer in restore_buffers.items():
92 + if name not in SKIP_TENSORS:
93 + buffer = restore_layer_refs(buffer, layer)
94 + layer.register_buffer(name, buffer)
95 +
96 +
97 +def materialize_layer(layer: torch.nn.Module) -> None:
98 + """Materialize all meta tensors in a layer to actual tensors."""
99 + if layer.__class__.__name__ in SKIP_MODULES:
100 + return
101 +
102 + for name, tensor in get_layer_tensors(layer).items():
103 + if name not in SKIP_TENSORS:
104 + setattr(layer, name, materialize_meta_tensor(tensor))
105 +
106 +
107 +class MetaCopyCounter(TorchDispatchMode):
108 + """
109 + Tracks total number of elements modified with `copy_`.
110 +
111 + Useful for keeping track of weight loading where underlying weights can be
112 + arbitrarily transformed (such as with `narrow`) before calling copy.
113 +
114 + Note: Assumes that copy kwargs are not used.
115 + """
116 +
117 + def __init__(self):
118 + super().__init__()
119 + self.copied_numel = 0
120 +
121 + def __torch_dispatch__(self, func, types, args=(), kwargs=None):
122 + if kwargs is None:
123 + kwargs = {}
124 +
125 + if func is torch.ops.aten.copy_.default and args[0].device.type == "meta":
126 + assert args[0].numel() == args[1].numel()
127 + self.copied_numel += args[0].numel()
128 +
129 + return func(*args, **kwargs)
130 +
131 +
132 +def get_numel_loaded(
133 + weight_loader: Callable, args: inspect.BoundArguments
134 +) -> tuple[int, object]:
135 + """
136 + Determine how many elements would be loaded by a weight loader call.
137 +
138 + :param weight loader: used to load weights
139 + :param args: bound arguments to weight loader
140 + :return: number of elements loaded by the weight loader, the return value of the
141 + weight loader
142 + """
143 + assert args.arguments["param"].device.type == "meta"
144 + with MetaCopyCounter() as counter:
145 + return_value = weight_loader(*args.args, **args.kwargs)
146 + return counter.copied_numel, return_value
vllm/model_executor/model_loader/reload/sanitize.py new
+50
@@ -0,0 +1,50 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from types import MethodType
4 +
5 +import torch
6 +
7 +__all__ = ["sanitize_layer_refs", "restore_layer_refs"]
8 +
9 +
10 +layer_ref_sentinel = object()
11 +
12 +
13 +def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor:
14 + """
15 + Removes references to layer held by tensor attributes. Specifically, removes the
16 + `__self__` attribute of weight loader methods attached to the tensor.
17 +
18 + Used by `capture_layer_to_meta` to avoid circular references to layers in
19 + `LAYERWISE_INFO`, leading to modules never being cleaned up. Without sanitation,
20 + tensors will reference layers, and the WeakKeyDictionary will never evict entries,
21 + even when the model is deleted.
22 +
23 + :param tensor: tensor to be sanitized
24 + :param layer: layer whose references should be removed
25 + :return: sanitized tensor
26 + """
27 + for key, value in tensor.__dict__.items():
28 + if isinstance(value, MethodType) and value.__self__ is layer:
29 + tensor.__dict__[key] = value.__func__.__get__(layer_ref_sentinel)
30 +
31 + return tensor
32 +
33 +
34 +def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor:
35 + """
36 + Restores references to layer held by tensor attributes.
37 +
38 + Used by `restore_layer_on_meta` to add back layer references, allowing for proper
39 + weight loading.
40 +
41 + :param tensor: tensor to be sanitized
42 + :param layer: layer whose references should be removed
43 + :return: sanitized tensor
44 +
45 + """
46 + for key, value in tensor.__dict__.items():
47 + if isinstance(value, MethodType) and value.__self__ is layer_ref_sentinel:
48 + tensor.__dict__[key] = value.__func__.__get__(layer)
49 +
50 + return tensor
vllm/model_executor/model_loader/reload/torchao_decorator.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable
5 +from functools import wraps
6 +from types import FunctionType
7 +from typing import TYPE_CHECKING
8 +
9 +import torch
10 +
11 +from vllm.config import ModelConfig
12 +
13 +from .layerwise import (
14 + finalize_layerwise_reload,
15 + initialize_layerwise_reload,
16 +)
17 +
18 +if TYPE_CHECKING:
19 + from vllm.model_executor.models.utils import AutoWeightsLoader
20 +
21 +__all__ = ["set_torchao_reload_attrs", "support_quantized_model_reload_from_hp_weights"]
22 +
23 +
24 +def set_torchao_reload_attrs(model: torch.nn.Module, model_config: ModelConfig):
25 + model._do_torchao_reload = True
26 + model._model_config = model_config
27 +
28 +
29 +def support_quantized_model_reload_from_hp_weights(original_load_weights: FunctionType):
30 + """
31 + Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
32 + reloading high precision (bfloat16/float16/float32) weight for an already quantized
33 + model, this involves restoring the weights to a high precision weights and
34 + then online quantize the weights.
35 +
36 + Only applies to torchao quantized models. Assumes that all model weights are
37 + loaded within a single weights iterator (cannot perform batched updates)
38 + """
39 +
40 + @wraps(original_load_weights)
41 + def patched_model_load_weights(
42 + self: "AutoWeightsLoader",
43 + weights: Iterable[tuple[str, torch.Tensor]],
44 + *args,
45 + **kwargs,
46 + ):
47 + model = self.module
48 +
49 + if not getattr(model, "_do_torchao_reload", False):
50 + return original_load_weights(self, weights, *args, **kwargs)
51 +
52 + initialize_layerwise_reload(model)
53 + loaded_weights = original_load_weights(self, weights, *args, **kwargs)
54 + finalize_layerwise_reload(model, model._model_config)
55 +
56 + return loaded_weights
57 +
58 + return patched_model_load_weights
vllm/model_executor/model_loader/reload/types.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass, field
4 +from inspect import BoundArguments
5 +
6 +import torch
7 +
8 +__all__ = ["LayerTensors", "LayerReloadingInfo"]
9 +
10 +# encodes both parameters and buffers separately
11 +LayerTensors = tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]
12 +
13 +
14 +@dataclass
15 +class LayerReloadingInfo:
16 + # model format (meta), populated by `record_metadata_for_reloading`
17 + restore_metadata: LayerTensors = field(default_factory=lambda: ({}, {}))
18 +
19 + # kernel format (device)
20 + kernel_tensors: LayerTensors = field(default_factory=lambda: ({}, {}))
21 +
22 + # track how many restored elements are ready for loading
23 + load_numel: int = 0
24 + load_numel_total: int | None = None
25 +
26 + # stores arguments and tensors ready for loading
27 + loaded_weights: list[tuple[str, BoundArguments]] = field(default_factory=list)
28 +
29 + def reset(self):
30 + self.__init__(restore_metadata=self.restore_metadata) # type: ignore[misc]
31 +
32 + def can_process(self) -> bool:
33 + return self.load_numel_total is not None
vllm/model_executor/model_loader/reload/utils.py new
+31
@@ -0,0 +1,31 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from .types import LayerTensors
7 +
8 +__all__ = [
9 + "get_layer_tensors",
10 + "get_layer_params_buffers",
11 + "get_layer_size",
12 +]
13 +
14 +
15 +def get_layer_tensors(layer: torch.nn.Module) -> dict[str, torch.Tensor]:
16 + """Get all parameters and buffers from a module as a dict."""
17 + params, buffers = get_layer_params_buffers(layer)
18 + return params | buffers
19 +
20 +
21 +def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
22 + """Get all parameters and buffers of a module as a tuple of dicts."""
23 + return (
24 + {name: param for name, param in layer._parameters.items() if param is not None},
25 + {name: buffer for name, buffer in layer._buffers.items() if buffer is not None},
26 + )
27 +
28 +
29 +def get_layer_size(layer: torch.nn.Module) -> int:
30 + """Calculate total number of elements across all tensors in a layer."""
31 + return sum(tensor.numel() for tensor in get_layer_tensors(layer).values())
vllm/model_executor/model_loader/utils.py
+16 -17
@@ -18,6 +18,10 @@ from vllm.model_executor.layers.quantization.base_config import (
18 QuantizationConfig,
19 QuantizeMethodBase,
20 )
21 +from vllm.model_executor.model_loader.reload import (
22 + record_metadata_for_reloading,
23 + set_torchao_reload_attrs,
24 +)
25 from vllm.model_executor.models.interfaces import SupportsQuant
26 from vllm.utils.platform_utils import is_pin_memory_available
27
@@ -45,7 +49,9 @@ def initialize_model(
49 if "vllm_config" in all_params and "prefix" in all_params:
50 # new-style model class
51 with set_current_vllm_config(vllm_config, check_compile=True, prefix=prefix):
48 - return model_class(vllm_config=vllm_config, prefix=prefix)
52 + model = model_class(vllm_config=vllm_config, prefix=prefix)
53 + record_metadata_for_reloading(model)
54 + return model
55
56 msg = (
57 "vLLM model class should accept `vllm_config` and `prefix` as "
@@ -75,27 +81,15 @@ def initialize_model(
81 if "scheduler_config" in all_params:
82 kwargs["scheduler_config"] = vllm_config.scheduler_config
83 with set_current_vllm_config(vllm_config, check_compile=True, prefix=prefix):
78 - return model_class(**kwargs)
84 + model = model_class(**kwargs)
85 + record_metadata_for_reloading(model)
86 +
87 + return model
88
89
90 def process_weights_after_loading(
91 model: nn.Module, model_config: ModelConfig, target_device: torch.device
92 ) -> None:
84 - if getattr(model, "process_weights_after_loading_already_called", False):
85 - # In case `process_weights_after_loading` is called multiple times
86 - # we'll skip it at later times
87 - logger.debug_once(
88 - "process_weights_after_loading already called for model %s", model
89 - )
90 - return
91 -
92 - # to avoid circular dependency
93 - from vllm.model_executor.model_loader.online_quantization import (
94 - maybe_save_metadata_and_attributes_for_weight_reloading,
95 - )
96 -
97 - maybe_save_metadata_and_attributes_for_weight_reloading(model, model_config)
98 -
93 for _, module in model.named_modules():
94 quant_method = getattr(module, "quant_method", None)
95 if isinstance(quant_method, QuantizeMethodBase):
@@ -117,6 +111,11 @@ def process_weights_after_loading(
111 # of process_weights_after_loading
112 module.process_weights_after_loading(model_config.dtype)
113
114 + # Needed for torchao model reloading via model.reload_weights
115 + # @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights`
116 + if model_config.quantization == "torchao":
117 + set_torchao_reload_attrs(model, model_config)
118 +
119
120 @contextmanager
121 def device_loading_context(module: torch.nn.Module, target_device: torch.device):
vllm/model_executor/models/utils.py
+1 -1
@@ -22,7 +22,7 @@ from vllm.logger import init_logger
22 from vllm.model_executor.layers.quantization.base_config import (
23 QuantizationConfig,
24 )
25 -from vllm.model_executor.model_loader.online_quantization import (
25 +from vllm.model_executor.model_loader.reload import (
26 support_quantized_model_reload_from_hp_weights,
27 )
28 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
vllm/model_executor/parameter.py
+1 -2
@@ -522,8 +522,7 @@ class SharedWeightParameter(BasevLLMParameter):
522 @property
523 def data(self):
524 raise ValueError(
525 - "Accessing `data` of a "
526 - "`PartitionedModelWeightParameter` is not allowed. "
525 + "Accessing `data` of a `SharedWeightParameter` is not allowed. "
526 "Instead, use `get_partition` to get the weight of "
527 "the particular partition you want to access"
528 )
vllm/v1/worker/gpu_model_runner.py
+90 -8
@@ -7,7 +7,7 @@ import itertools
7 import threading
8 import time
9 from collections import defaultdict
10 -from collections.abc import Iterator, Sequence
10 +from collections.abc import Iterable, Iterator, Sequence
11 from contextlib import contextmanager
12 from copy import copy, deepcopy
13 from dataclasses import dataclass
@@ -59,6 +59,10 @@ from vllm.model_executor.layers.rotary_embedding import (
59 XDRotaryEmbedding,
60 )
61 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
62 +from vllm.model_executor.model_loader.reload import (
63 + finalize_layerwise_reload,
64 + initialize_layerwise_reload,
65 +)
66 from vllm.model_executor.models.interfaces import (
67 MultiModalEmbeddings,
68 SupportsMRoPE,
@@ -2524,8 +2528,10 @@ class GPUModelRunner(
2528 return mm_embeds, is_mm_embed
2529
2530 def get_model(self) -> nn.Module:
2527 - # get raw model out of the cudagraph wrapper.
2531 + if not hasattr(self, "model"):
2532 + raise ValueError("Cannot get model before model has been initialized")
2533 if isinstance(self.model, (CUDAGraphWrapper, UBatchWrapper)):
2534 + # get raw model out of the cudagraph wrapper.
2535 return self.model.unwrap()
2536 return self.model
2537
@@ -4270,13 +4276,89 @@ class GPUModelRunner(
4276
4277 return None
4278
4273 - def reload_weights(self) -> None:
4274 - assert getattr(self, "model", None) is not None, (
4275 - "Cannot reload weights before model is loaded."
4279 + def reload_weights(
4280 + self,
4281 + weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,
4282 + weights_path: str | None = None,
4283 + is_checkpoint_format: bool = True,
4284 + ) -> None:
4285 + """
4286 + Reload weights from a weights iterator or from disk
4287 +
4288 + :param weights_iterator: weights to load into model
4289 + :param weights_path: path to load weights from if weights_iterator is not
4290 + provided. Use path of original model if neither is provided.
4291 + :param is_checkpoint_format: set to False if weights have already been processed
4292 + into kernel format (repacking, renaming, ect.)
4293 + """
4294 + # TODO(@kylesayrs): generalize to all runners and loaders
4295 + # argument validation
4296 + if weights_iterator is None and not is_checkpoint_format:
4297 + logger.warning(
4298 + "Reloading from disk means that weights will be in checkpoint format. "
4299 + "Please use `is_checkpoint_format=True` "
4300 + "to avoid weight reloading errors"
4301 + )
4302 +
4303 + model = self.get_model()
4304 + weights_to_load = {name for name, _ in model.named_parameters()}
4305 + counter_before_reloading = time.perf_counter()
4306 +
4307 + # load weights from disk if none are provided
4308 + if weights_iterator is None:
4309 + model_loader = get_model_loader(self.load_config)
4310 + if not hasattr(model_loader, "get_all_weights"):
4311 + raise NotImplementedError(
4312 + f"Model reloading with `{self.load_config.load_format}` format"
4313 + )
4314 +
4315 + if weights_path is not None:
4316 + self.model_config.model = weights_path
4317 + weights_iterator = model_loader.get_all_weights(self.model_config, model)
4318 + weights_iterator = cast(
4319 + Iterable[tuple[str, torch.Tensor]], weights_iterator
4320 + )
4321 +
4322 + # begin loading weights
4323 + logger.info_once("Reloading weights inplace...", scope="local")
4324 + load_device = (
4325 + self.vllm_config.load_config.device or self.vllm_config.device_config.device
4326 )
4277 - model_loader = get_model_loader(self.load_config)
4278 - logger.info("Reloading weights inplace...")
4279 - model_loader.load_weights(self.get_model(), model_config=self.model_config)
4327 + with torch.device(load_device):
4328 + if is_checkpoint_format:
4329 + # load weights from checkpoint/ original model format
4330 + initialize_layerwise_reload(model)
4331 + loaded_weights = model.load_weights(weights_iterator)
4332 + finalize_layerwise_reload(model, self.model_config)
4333 +
4334 + else:
4335 + # load weights from kernel format
4336 + logger.warning_once(
4337 + "Reloading with `is_checkpoint_format=True` requires that "
4338 + "weights be in kernel format and already sharded",
4339 + scope="local",
4340 + )
4341 + loaded_weights = set()
4342 + for name, loaded_weight in weights_iterator:
4343 + param = model.get_parameter(name) # TODO: buffers?
4344 + param.copy_(loaded_weight)
4345 + loaded_weights.add(name)
4346 +
4347 + # logging and validation
4348 + counter_after_reloading = time.perf_counter()
4349 + diff_seconds = counter_after_reloading - counter_before_reloading
4350 + logger.info_once(
4351 + "Reloading and processing weights took %.2f seconds",
4352 + diff_seconds,
4353 + scope="local",
4354 + )
4355 + if self.model_config.quantization is None and loaded_weights is not None:
4356 + weights_not_loaded = weights_to_load - loaded_weights
4357 + if weights_not_loaded:
4358 + logger.warning(
4359 + "Following weights were not loaded from checkpoint: %s",
4360 + weights_not_loaded,
4361 + )
4362
4363 def save_tensorized_model(
4364 self,
vllm/v1/worker/gpu_worker.py
+2 -2
@@ -280,8 +280,8 @@ class Worker(WorkerBase):
280 def update_config(self, overrides: dict[str, Any]) -> None:
281 self.model_runner.update_config(overrides)
282
283 - def reload_weights(self) -> None:
284 - self.model_runner.reload_weights()
283 + def reload_weights(self, *args, **kwargs) -> None:
284 + self.model_runner.reload_weights(*args, **kwargs)
285
286 @torch.inference_mode()
287 def determine_available_memory(self) -> int: