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 4, 12:30 UTC
[UX] Allow enable/disable model weights loading tracking by config (#41086)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Copilot <copilot@github.com>
Isotr0py committed
Apr 29, 2026 at 12:04 UTC
99255f3cb5cec7466bf9fa5310fd310baf87d711
1 file changed
+39
-4
vllm/model_executor/model_loader/default_loader.py
+39
-4
@@ -76,7 +76,11 @@ class DefaultModelLoader(BaseModelLoader):
76
self.local_expert_ids: set[int] | None = None
77
78
extra_config = load_config.model_loader_extra_config
79
- allowed_keys = {"enable_multithread_load", "num_threads"}
79
+ allowed_keys = {
80
+ "enable_multithread_load",
81
+ "num_threads",
82
+ "enable_weights_track",
83
+ }
84
unexpected_keys = set(extra_config.keys()) - allowed_keys
85
86
if unexpected_keys:
@@ -86,6 +90,10 @@ class DefaultModelLoader(BaseModelLoader):
90
f"{unexpected_keys}"
91
)
92
93
+ self.enable_weights_track: bool | None = extra_config.get(
94
+ "enable_weights_track", None
95
+ )
96
+
97
def _prepare_weights(
98
self,
99
model_name_or_path: str,
@@ -377,7 +385,6 @@ class DefaultModelLoader(BaseModelLoader):
385
386
self._init_ep_weight_filter(model_config)
387
380
- weights_to_load = {name for name, _ in model.named_parameters()}
388
loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
389
390
self.counter_after_loading_weights = time.perf_counter()
@@ -386,8 +393,36 @@ class DefaultModelLoader(BaseModelLoader):
393
self.counter_after_loading_weights - self.counter_before_loading_weights,
394
)
395
# We only enable strict check for non-quantized models
389
- # that have loaded weights tracking currently.
390
- if model_config.quantization is None and loaded_weights is not None:
396
+ # that have loaded weights tracking by default.
397
+ default_enable_weights_track = (
398
+ model_config.quantization is None and loaded_weights is not None
399
+ )
400
+ enable_weights_track = (
401
+ self.enable_weights_track
402
+ if self.enable_weights_track is not None
403
+ else default_enable_weights_track
404
+ )
405
+ if enable_weights_track:
406
+ self.track_weights_loading(model, loaded_weights)
407
+
408
+ def track_weights_loading(
409
+ self, model: nn.Module, loaded_weights: set[str] | None
410
+ ) -> None:
411
+ weights_to_load = {name for name, _ in model.named_parameters()}
412
+ if loaded_weights is not None:
413
+ # ignore online quantization scales
414
+ for name, module in model.named_modules():
415
+ quant_method = getattr(module, "quant_method", None)
416
+ has_online_quant = getattr(quant_method, "uses_meta_device", False)
417
+ has_postprocess_quant = getattr(
418
+ quant_method, "process_weights_after_loading", None
419
+ )
420
+ # ignore kv_cache scale and online quant scale,
421
+ # which can be missing in checkpoints
422
+ if has_online_quant or has_postprocess_quant:
423
+ for param_name, _ in module.named_parameters():
424
+ full_name = f"{name}.{param_name}" if name else param_name
425
+ loaded_weights.add(full_name)
426
weights_not_loaded = weights_to_load - loaded_weights
427
if weights_not_loaded:
428
raise ValueError(