[Refactor] Add opt-in loader-side weight completeness verification

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)

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)
+    }

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()