From 51aac28fe6591a53d9a8be906ef46a6b4bb03b6a Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Mon, 22 Jun 2026 21:07:59 +0000 Subject: [PATCH] model_loader: add generic checkpoint-completeness verification (dormant) Separate "load the weights" from "verify the checkpoint was complete". A new `_load_and_verify_weights` runs after a full load and fails fast on any required parameter that was left uninitialized -- computed as pure set algebra (`unloaded_required_params`), with legitimately-absent params (PP-distributed, quant/kv scales, role-specific) excluded by an optional per-model `is_optional_weight(name)` predicate. Gated by a per-model `verify_weights_on_load` flag (default off): no model opts in here, so this is dormant -- zero behavior change. It is the mechanism for moving the "weights are not initialized from checkpoints" diagnostic out of individual models' `load_weights` (where it false-positives when the online/RL weight-update path reuses `load_weights` with partial tensor slices) and into the loader, which runs it only on full checkpoint loads. See RFC #28585. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/model_loader/loader.py | 42 ++++++++++++++----- .../srt/model_loader/weight_completeness.py | 14 +++++++ .../test_load_and_verify_weights.py | 34 +++++++++++++++ 3 files changed, 79 insertions(+), 11 deletions(-) create mode 100644 python/sglang/srt/model_loader/weight_completeness.py create mode 100644 test/registered/unit/model_loader/test_load_and_verify_weights.py diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 0ddc5d9ce604..4672d29752e0 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -44,6 +44,7 @@ get_remote_instance_transfer_engine_info_per_rank, register_memory_region, ) +from sglang.srt.model_loader.weight_completeness import unloaded_required_params from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import get_available_gpu_memory @@ -329,6 +330,25 @@ def _post_load_weights(model: nn.Module) -> None: model.post_load_weights() +def _not_optional(name: str) -> bool: + return False + + +def _load_and_verify_weights(model: nn.Module, weights) -> Optional[set]: + loaded_params = model.load_weights(weights) + if loaded_params is None or not getattr(model, "verify_weights_on_load", False): + return loaded_params + is_optional = getattr(model, "is_optional_weight", _not_optional) + missing = unloaded_required_params( + (name for name, _ in model.named_parameters()), loaded_params, is_optional + ) + if missing: + raise RuntimeError( + f"Some weights are not initialized from checkpoints: {sorted(missing)}" + ) + return loaded_params + + class BaseModelLoader(ABC): """Base class for model loaders.""" @@ -794,12 +814,12 @@ def load_weights_and_postprocess(model, weights, target_device): TRTLLM_DISABLE_FP4_QUANT_FAST_MATH="1", FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH="1", ): - model.load_weights(weights) + _load_and_verify_weights(model, weights) if target_device.type == "cuda": torch.cuda.synchronize() torch.cuda.empty_cache() else: - model.load_weights(weights) + _load_and_verify_weights(model, weights) # Used in tests to verify memory savings when using online quantization. if is_cuda_alike(): @@ -988,15 +1008,14 @@ def load_weights_and_postprocess(self, model, weights, target_device): def load_weights_proxy(weights): if QuantizedRLModelLoader.is_reload_scenario(model): logger.info("[QuantizedRL] Using fast path reload in load_weights") - QuantizedRLModelLoader.rebinding_and_load_weights( + return QuantizedRLModelLoader.rebinding_and_load_weights( model, original_load_weights, weights ) - else: - original_load_weights(weights) + return original_load_weights(weights) model.load_weights = load_weights_proxy - model.load_weights(weights) + _load_and_verify_weights(model, weights) original_weights = dict(model.named_parameters()) # Record pre-quantization state (shape/stride) for torch.as_strided reset @@ -1989,7 +2008,7 @@ def _load_weights(self, model_config: ModelConfig, model: nn.Module) -> None: model_config.model_path, model_config.revision, pre_quant, load_8bit ) - model.load_weights(qweight_iterator) + _load_and_verify_weights(model, qweight_iterator) current_platform.empty_cache() @@ -2181,8 +2200,9 @@ def load_model( with set_default_torch_dtype(model_config.dtype): with target_device: model = _initialize_model(model_config, self.load_config, quant_config) - model.load_weights( - self._get_weights_iterator(local_model_path, gguf_weights_map) + _load_and_verify_weights( + model, + self._get_weights_iterator(local_model_path, gguf_weights_map), ) for _, module in model.named_modules(): @@ -2496,7 +2516,7 @@ def _load_model_from_remote_fs( target_device = torch.device(device_config.device) with set_default_torch_dtype(model_config.dtype): - model.load_weights(self._get_weights_iterator_fs(client)) + _load_and_verify_weights(model, self._get_weights_iterator_fs(client)) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) @@ -2566,7 +2586,7 @@ def load_model_with_cpu_quantization( ) if not isinstance(self, DummyModelLoader): - model.load_weights(self._get_all_weights(model_config, model)) + _load_and_verify_weights(model, self._get_all_weights(model_config, model)) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) diff --git a/python/sglang/srt/model_loader/weight_completeness.py b/python/sglang/srt/model_loader/weight_completeness.py new file mode 100644 index 000000000000..e352cb35c636 --- /dev/null +++ b/python/sglang/srt/model_loader/weight_completeness.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from typing import Callable, Container, Iterable + + +def unloaded_required_params( + param_names: Iterable[str], + loaded: Container[str], + is_optional: Callable[[str], bool], +) -> set[str]: + """Return required parameters that were not loaded.""" + return { + name for name in param_names if name not in loaded and not is_optional(name) + } diff --git a/test/registered/unit/model_loader/test_load_and_verify_weights.py b/test/registered/unit/model_loader/test_load_and_verify_weights.py new file mode 100644 index 000000000000..1f8c59acd823 --- /dev/null +++ b/test/registered/unit/model_loader/test_load_and_verify_weights.py @@ -0,0 +1,34 @@ +import unittest + +from sglang.srt.model_loader.weight_completeness import unloaded_required_params +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _never(name: str) -> bool: + return False + + +class TestUnloadedRequiredParams(unittest.TestCase): + def test_reports_required_params_not_loaded(self): + missing = unloaded_required_params( + ["a.weight", "b.weight"], {"a.weight"}, _never + ) + self.assertEqual(missing, {"b.weight"}) + + def test_optional_params_are_excluded(self): + missing = unloaded_required_params( + ["a.weight", "b.k_scale"], set(), lambda n: n.endswith(".k_scale") + ) + self.assertEqual(missing, {"a.weight"}) + + def test_empty_when_all_required_loaded(self): + missing = unloaded_required_params( + ["a.weight", "b.weight"], {"a.weight", "b.weight"}, _never + ) + self.assertEqual(missing, set()) + + +if __name__ == "__main__": + unittest.main()