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, 17:15 UTC
[Quantization] Consolidate dummy format logic into DummyModelLoader (#38637)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Asaf Gardin committed
Apr 1, 2026 at 01:20 UTC
3dc01ef352726e20b8a2e7c8430ae4544b7bb284
2 files changed
+42
-15
vllm/model_executor/model_loader/dummy_loader.py
+40
-4
@@ -4,8 +4,19 @@ import torch.nn as nn
4
5
from vllm.config import ModelConfig
6
from vllm.config.load import LoadConfig
7
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
8
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
8
-from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights
9
+from vllm.model_executor.model_loader.reload.layerwise import (
10
+ _get_original_loader,
11
+ get_layerwise_info,
12
+)
13
+from vllm.model_executor.model_loader.reload.meta import materialize_layer
14
+from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo
15
+from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
16
+from vllm.model_executor.model_loader.weight_utils import (
17
+ initialize_dummy_weights,
18
+ initialize_single_dummy_weight,
19
+)
20
21
22
class DummyModelLoader(BaseModelLoader):
@@ -23,6 +34,31 @@ class DummyModelLoader(BaseModelLoader):
34
pass # Nothing to download
35
36
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
26
- # NOTE(woosuk): For accurate performance evaluation, we assign
27
- # random values to the weights.
28
- initialize_dummy_weights(model, model_config)
37
+ for layer in model.modules():
38
+ info = get_layerwise_info(layer)
39
+ if info.can_load():
40
+ self._process_online_quant_layer(layer, info)
41
+ else:
42
+ # NOTE(woosuk): For accurate performance evaluation, we assign
43
+ # random values to the weights.
44
+ initialize_dummy_weights(layer, model_config)
45
+
46
+ def _process_online_quant_layer(
47
+ self,
48
+ layer: nn.Module,
49
+ info: LayerReloadingInfo,
50
+ ) -> None:
51
+ """Materialize, apply dummy weights, and run quantization processing."""
52
+ materialize_layer(layer, info)
53
+
54
+ for tensor in get_layer_tensors(layer).values():
55
+ initialize_single_dummy_weight(tensor)
56
+
57
+ for param in get_layer_tensors(layer).values():
58
+ param.weight_loader = _get_original_loader(param)
59
+
60
+ quant_method = getattr(layer, "quant_method", None)
61
+ if isinstance(quant_method, QuantizeMethodBase):
62
+ quant_method.process_weights_after_loading(layer)
63
+
64
+ info.reset()
vllm/model_executor/model_loader/reload/layerwise.py
+2
-11
@@ -11,10 +11,7 @@ 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 (
15
- default_weight_loader,
16
- initialize_single_dummy_weight,
17
-)
14
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
16
from .meta import (
17
capture_layer_to_meta,
@@ -224,7 +221,7 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
221
222
# No weights were loaded
223
elif info.load_numel <= 0:
227
- # first load but received no weights. This happens on dummy load
224
+ # first load: checkpoint did not contain weights for this layer
225
if info.kernel_tensors is None:
226
_layerwise_process(layer, info)
227
continue
@@ -262,12 +259,6 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
259
# Materialize layer tensors onto device
260
materialize_layer(layer, info)
261
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
-
262
# Reset online quantization flag so process_weights_after_loading
263
# will run again during reload
264
if hasattr(layer, "_already_called_process_weights_after_loading"):