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 7, 19:00 UTC
[Misc] Enable weights loading tracking for quantized models (#35074)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Feb 25, 2026 at 14:11 UTC
4572a06afe96d0a6d5d3efacf130c71505dd2bc9
1 file changed
+15
-4
vllm/model_executor/model_loader/default_loader.py
+15
-4
@@ -14,6 +14,7 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
14
from vllm.config import ModelConfig
15
from vllm.config.load import LoadConfig
16
from vllm.logger import init_logger
17
+from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
18
from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
19
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
20
from vllm.model_executor.model_loader.weight_utils import (
@@ -286,7 +287,6 @@ class DefaultModelLoader(BaseModelLoader):
287
):
288
self.load_config.safetensors_load_strategy = "torchao"
289
289
- weights_to_load = {name for name, _ in model.named_parameters()}
290
loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
291
292
self.counter_after_loading_weights = time.perf_counter()
@@ -295,9 +295,20 @@ class DefaultModelLoader(BaseModelLoader):
295
self.counter_after_loading_weights - self.counter_before_loading_weights,
296
scope="local",
297
)
298
- # We only enable strict check for non-quantized models
299
- # that have loaded weights tracking currently.
300
- if model_config.quantization is None and loaded_weights is not None:
298
+ self.track_weights_loading(model, loaded_weights)
299
+
300
+ def track_weights_loading(
301
+ self, model: nn.Module, loaded_weights: set[str] | None
302
+ ) -> None:
303
+ weights_to_load = {name for name, _ in model.named_parameters()}
304
+ if loaded_weights is not None:
305
+ for name, module in model.named_modules():
306
+ quant_method = getattr(module, "quant_method", None)
307
+ # ignore kv_cache scale, which can be missing in checkpoints
308
+ if isinstance(quant_method, BaseKVCacheMethod):
309
+ for param_name, _ in module.named_parameters():
310
+ full_name = f"{name}.{param_name}" if name else param_name
311
+ loaded_weights.add(full_name)
312
weights_not_loaded = weights_to_load - loaded_weights
313
if weights_not_loaded:
314
raise ValueError(