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 5, 19:45 UTC
[Bug fix][Quantization] Fix dummy weight loading (#38478)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Asaf Gardin committed
Mar 30, 2026 at 23:38 UTC
1fc69f59bb0838c2ff6efc416dd8875c3e210d04
3 files changed
+19
-13
vllm/model_executor/model_loader/dummy_loader.py
-9
@@ -1,13 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import torch
3
import torch.nn as nn
4
5
from vllm.config import ModelConfig
6
from vllm.config.load import LoadConfig
7
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
9
-from vllm.model_executor.model_loader.reload.meta import materialize_meta_tensor
10
-from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
8
from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights
9
10
@@ -26,12 +23,6 @@ class DummyModelLoader(BaseModelLoader):
23
pass # Nothing to download
24
25
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
29
- # materialize meta tensors as part of online quantization lifecycle
30
- for layer in model.modules():
31
- for name, param in get_layer_tensors(layer).items():
32
- if param.device == torch.device("meta"):
33
- setattr(layer, name, materialize_meta_tensor(param))
34
-
26
# NOTE(woosuk): For accurate performance evaluation, we assign
27
# random values to the weights.
28
initialize_dummy_weights(model, model_config)
vllm/model_executor/model_loader/reload/layerwise.py
+12
-2
@@ -11,7 +11,10 @@ from vllm.config import ModelConfig
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.attention import Attention, MLAAttention
13
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
14
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
14
+from vllm.model_executor.model_loader.weight_utils import (
15
+ default_weight_loader,
16
+ initialize_single_dummy_weight,
17
+)
18
19
from .meta import (
20
capture_layer_to_meta,
@@ -223,7 +226,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
226
elif info.load_numel <= 0:
227
# first load but received no weights. This happens on dummy load
228
if info.kernel_tensors is None:
226
- materialize_layer(layer, info)
229
+ _layerwise_process(layer, info)
230
+ continue
231
232
# reloading: place kernel tensors back as a fallback
233
else:
@@ -258,6 +262,12 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
262
# Materialize layer tensors onto device
263
materialize_layer(layer, info)
264
265
+ # If no weights were loaded (e.g. dummy loading), initialize with
266
+ # small random values to avoid NaN from zero/garbage data
267
+ if len(info.loaded_weights) <= 0:
268
+ for tensor in get_layer_tensors(layer).values():
269
+ initialize_single_dummy_weight(tensor)
270
+
271
# Reset online quantization flag so process_weights_after_loading
272
# will run again during reload
273
if hasattr(layer, "_already_called_process_weights_after_loading"):
vllm/model_executor/model_loader/weight_utils.py
+7
-2
@@ -239,10 +239,12 @@ def convert_bin_to_safetensor_file(
239
sf_size = os.stat(sf_filename).st_size
240
pt_size = os.stat(pt_filename).st_size
241
if (sf_size - pt_size) / pt_size > 0.01:
242
- raise RuntimeError(f"""The file size different is more than 1%:
242
+ raise RuntimeError(
243
+ f"""The file size different is more than 1%:
244
- {sf_filename}: {sf_size}
245
- {pt_filename}: {pt_size}
245
- """)
246
+ """
247
+ )
248
249
# check if the tensors are the same
250
reloaded = load_file(sf_filename)
@@ -1377,6 +1379,9 @@ def initialize_single_dummy_weight(
1379
high: float = 1e-3,
1380
seed: int = 1234,
1381
) -> None:
1382
+ if param.device.type == "meta":
1383
+ return # deferred to finalize_layerwise_processing (e.g. online quant)
1384
+
1385
if not torch.is_floating_point(param):
1386
if current_platform.is_rocm():
1387
# On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left